Spaces:
Runtime error
Runtime error
| # ------------------------------------------------------------------------ | |
| # HOTR official code : hotr/models/detr.py | |
| # Copyright (c) Kakao Brain, Inc. and its affiliates. All Rights Reserved | |
| # ------------------------------------------------------------------------ | |
| # Modified from DETR (https://github.com/facebookresearch/detr) | |
| # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved | |
| # ------------------------------------------------------------------------ | |
| """ | |
| DETR & HOTR model and criterion classes. | |
| """ | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn | |
| from hotr.util.misc import (NestedTensor, nested_tensor_from_tensor_list) | |
| from .backbone import build_backbone | |
| from .detr_matcher import build_matcher | |
| from .hotr_matcher import build_hoi_matcher | |
| from .transformer import build_transformer, build_hoi_transformer | |
| from .criterion import SetCriterion | |
| from .post_process import PostProcess | |
| from .feed_forward import MLP | |
| from .hotr import HOTR | |
| from .hotr_v1 import HOTR_V1 | |
| class DETR(nn.Module): | |
| """ This is the DETR module that performs object detection """ | |
| def __init__(self, backbone, transformer, num_classes, num_queries, aux_loss=False): | |
| """ Initializes the model. | |
| Parameters: | |
| backbone: torch module of the backbone to be used. See backbone.py | |
| transformer: torch module of the transformer architecture. See transformer.py | |
| num_classes: number of object classes | |
| num_queries: number of object queries, ie detection slot. This is the maximal number of objects | |
| DETR can detect in a single image. For COCO, we recommend 100 queries. | |
| aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used. | |
| """ | |
| super().__init__() | |
| self.num_queries = num_queries | |
| self.transformer = transformer | |
| hidden_dim = transformer.d_model | |
| self.class_embed = nn.Linear(hidden_dim, num_classes + 1) | |
| self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3) | |
| self.query_embed = nn.Embedding(num_queries, hidden_dim) | |
| self.input_proj = nn.Conv2d(backbone.num_channels, hidden_dim, kernel_size=1) | |
| self.backbone = backbone | |
| self.aux_loss = aux_loss | |
| def forward(self, samples: NestedTensor): | |
| """ The forward expects a NestedTensor, which consists of: | |
| - samples.tensor: batched images, of shape [batch_size x 3 x H x W] | |
| - samples.mask: a binary mask of shape [batch_size x H x W], containing 1 on padded pixels | |
| It returns a dict with the following elements: | |
| - "pred_logits": the classification logits (including no-object) for all queries. | |
| Shape= [batch_size x num_queries x (num_classes + 1)] | |
| - "pred_boxes": The normalized boxes coordinates for all queries, represented as | |
| (center_x, center_y, height, width). These values are normalized in [0, 1], | |
| relative to the size of each individual image (disregarding possible padding). | |
| See PostProcess for information on how to retrieve the unnormalized bounding box. | |
| - "aux_outputs": Optional, only returned when auxilary losses are activated. It is a list of | |
| dictionnaries containing the two above keys for each decoder layer. | |
| """ | |
| if isinstance(samples, (list, torch.Tensor)): | |
| samples = nested_tensor_from_tensor_list(samples) | |
| features, pos = self.backbone(samples) | |
| src, mask = features[-1].decompose() | |
| assert mask is not None | |
| hs = self.transformer(self.input_proj(src), mask, self.query_embed.weight, pos[-1])[0] | |
| outputs_class = self.class_embed(hs) | |
| outputs_coord = self.bbox_embed(hs).sigmoid() | |
| out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]} | |
| if self.aux_loss: | |
| out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord) | |
| return out | |
| def _set_aux_loss(self, outputs_class, outputs_coord): | |
| # this is a workaround to make torchscript happy, as torchscript | |
| # doesn't support dictionary with non-homogeneous values, such | |
| # as a dict having both a Tensor and a list. | |
| return [{'pred_logits': a, 'pred_boxes': b} | |
| for a, b in zip(outputs_class[:-1], outputs_coord[:-1])] | |
| def build(args): | |
| device = torch.device(args.device) | |
| backbone = build_backbone(args) | |
| transformer = build_transformer(args) | |
| model = DETR( | |
| backbone, | |
| transformer, | |
| num_classes=args.num_classes, | |
| num_queries=args.num_queries, | |
| aux_loss=args.aux_loss, | |
| ) | |
| matcher = build_matcher(args) | |
| weight_dict = {'loss_ce': 1, 'loss_bbox': args.bbox_loss_coef} | |
| weight_dict['loss_giou'] = args.giou_loss_coef | |
| # TODO this is a hack | |
| if args.aux_loss: | |
| aux_weight_dict = {} | |
| for i in range(args.dec_layers - 1): | |
| aux_weight_dict.update({k + f'_{i}': v for k, v in weight_dict.items()}) | |
| weight_dict.update(aux_weight_dict) | |
| losses = ['labels', 'boxes', 'cardinality'] if args.frozen_weights is None else [] | |
| if args.HOIDet: | |
| hoi_matcher = build_hoi_matcher(args) | |
| hoi_losses = [] | |
| hoi_losses.append('pair_labels') | |
| hoi_losses.append('pair_actions') | |
| if args.dataset_file == 'hico-det': hoi_losses.append('pair_targets') | |
| hoi_weight_dict={} | |
| hoi_weight_dict['loss_hidx'] = args.hoi_idx_loss_coef | |
| hoi_weight_dict['loss_oidx'] = args.hoi_idx_loss_coef | |
| hoi_weight_dict['loss_h_consistency'] = args.hoi_idx_consistency_loss_coef | |
| hoi_weight_dict['loss_o_consistency'] = args.hoi_idx_consistency_loss_coef | |
| hoi_weight_dict['loss_act'] = args.hoi_act_loss_coef | |
| hoi_weight_dict['loss_act_consistency'] = args.hoi_act_consistency_loss_coef | |
| if args.dataset_file == 'hico-det': | |
| hoi_weight_dict['loss_tgt'] = args.hoi_tgt_loss_coef | |
| hoi_weight_dict['loss_tgt_consistency'] = args.hoi_tgt_consistency_loss_coef | |
| if args.hoi_aux_loss: | |
| hoi_aux_weight_dict = {} | |
| for i in range(args.hoi_dec_layers): | |
| hoi_aux_weight_dict.update({k + f'_{i}': v for k, v in hoi_weight_dict.items()}) | |
| hoi_weight_dict.update(hoi_aux_weight_dict) | |
| criterion = SetCriterion(args.num_classes, matcher=matcher, weight_dict=hoi_weight_dict, | |
| eos_coef=args.eos_coef, losses=losses, num_actions=args.num_actions, | |
| HOI_losses=hoi_losses, HOI_matcher=hoi_matcher, args=args) | |
| interaction_transformer = build_hoi_transformer(args) # if (args.share_enc and args.pretrained_dec) else None | |
| kwargs = {} | |
| if args.dataset_file == 'hico-det': kwargs['return_obj_class'] = args.valid_obj_ids | |
| if args.sep_enc_forward: | |
| model = HOTR_V1( | |
| detr=model, | |
| num_hoi_queries=args.num_hoi_queries, | |
| num_actions=args.num_actions, | |
| interaction_transformer=interaction_transformer, | |
| augpath_name = args.augpath_name, | |
| share_dec_param = args.share_dec_param, | |
| stop_grad_stage = args.stop_grad_stage, | |
| freeze_detr=(args.frozen_weights is not None), | |
| share_enc=args.share_enc, | |
| pretrained_dec=args.pretrained_dec, | |
| temperature=args.temperature, | |
| hoi_aux_loss=args.hoi_aux_loss, | |
| **kwargs # only return verb class for HICO-DET dataset | |
| ) | |
| else: | |
| model = HOTR( | |
| detr=model, | |
| num_hoi_queries=args.num_hoi_queries, | |
| num_actions=args.num_actions, | |
| interaction_transformer=interaction_transformer, | |
| augpath_name = args.augpath_name, | |
| share_dec_param = args.share_dec_param, | |
| stop_grad_stage = args.stop_grad_stage, | |
| freeze_detr=(args.frozen_weights is not None), | |
| share_enc=args.share_enc, | |
| pretrained_dec=args.pretrained_dec, | |
| temperature=args.temperature, | |
| hoi_aux_loss=args.hoi_aux_loss, | |
| **kwargs # only return verb class for HICO-DET dataset | |
| ) | |
| postprocessors = {'hoi': PostProcess(args.HOIDet)} | |
| else: | |
| criterion = SetCriterion(args.num_classes, matcher=matcher, weight_dict=weight_dict, | |
| eos_coef=args.eos_coef, losses=losses) | |
| postprocessors = {'bbox': PostProcess(args.HOIDet)} | |
| criterion.to(device) | |
| return model, criterion, postprocessors |