Spaces:
Running on Zero
Running on Zero
| # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved | |
| """ | |
| Model Architecture modified from DETR for OmniShotCut Model | |
| """ | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn | |
| # Import files from the local fodler | |
| # from util import box_ops | |
| from util.misc import NestedTensor | |
| from .backbone import build_backbone | |
| from .transformer import build_transformer | |
| from .loss import Criterion | |
| from datasets.utils import nested_tensor_from_tensor_list | |
| class MLP(nn.Module): | |
| """ Very simple multi-layer perceptron (also called FFN)""" | |
| def __init__(self, input_dim, hidden_dim, output_dim, num_layers): | |
| super().__init__() | |
| self.num_layers = num_layers | |
| h = [hidden_dim] * (num_layers - 1) | |
| self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])) | |
| def forward(self, x): | |
| for i, layer in enumerate(self.layers): | |
| x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x) | |
| return x | |
| class OmniShotCut(nn.Module): | |
| """ This is the OmniShotCut module that performs object detection """ | |
| def __init__(self, backbone, transformer, num_intra_relation_classes, num_inter_relation_classes, num_frames, 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 | |
| OmniShotCut 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.num_frames = num_frames | |
| self.aux_loss = aux_loss | |
| # Trainable Parameters | |
| self.backbone = backbone | |
| self.transformer = transformer | |
| hidden_dim = transformer.d_model | |
| self.intra_relation_class_embed = nn.Linear(hidden_dim, num_intra_relation_classes) | |
| self.inter_relation_class_embed = nn.Linear(hidden_dim, num_inter_relation_classes) | |
| self.range_class_embed = nn.Linear(hidden_dim, num_frames + 2) # TODO: s+2 is to add some padding | |
| self.query_embed = nn.Embedding(num_queries, hidden_dim) | |
| self.input_proj = nn.Conv2d(backbone.num_channels, hidden_dim, kernel_size = 1) | |
| 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: | |
| - "intra_clip_logits": the classification logits (including no-object) for all queries. | |
| - "inter_clip_logits": the classification logits (including no-object) for all queries. | |
| - "pred_shot_logits": The normalized ranges 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) | |
| # Call the Backbone (ResNet18) | |
| features, pos = self.backbone(samples) # Joiner outputs ResNet features and the position embedding | |
| pos = pos[-1] # output shape is (B, F, 384(C), H, W) | |
| src, mask = features[-1].decompose() # src shape is [B * F, C, H, W] | |
| assert mask is not None | |
| # Reshape Vision Inputs | |
| src_proj = self.input_proj(src) | |
| n, c, h, w = src_proj.shape | |
| # Reshape src_proj to (B, C, F, H*W); mask to (B, F, H*W) | |
| src_proj = src_proj.reshape(n // self.num_frames, self.num_frames, c, h, w).permute(0, 2, 1, 3, 4).flatten(-2) | |
| mask = mask.reshape(n // self.num_frames, self.num_frames, h*w) | |
| pos = pos.permute(0, 2, 1, 3, 4).flatten(-2) | |
| # Call Transformer; output shape is (6, B, #Queries, F*H*W) | |
| hs = self.transformer(src_proj, mask, self.query_embed.weight, pos)[0] # src_proj shape is (B, C, F, H*W); mask shape is (B, F, H*W); pos shape is (B, C, F, H*W) | |
| # Output Class will be modified for Cut Anything | |
| outputs_intra_class = self.intra_relation_class_embed(hs) | |
| outputs_inter_class = self.inter_relation_class_embed(hs) | |
| outputs_shot_class = self.range_class_embed(hs) | |
| out = { | |
| 'intra_clip_logits': outputs_intra_class[-1], | |
| 'inter_clip_logits': outputs_inter_class[-1], | |
| 'pred_shot_logits': outputs_shot_class[-1] | |
| } # Last Layer | |
| if self.aux_loss: | |
| out['aux_outputs'] = self._set_aux_loss(outputs_intra_class, outputs_inter_class, outputs_shot_class) # Fetch Central Layer (Except the last layer) | |
| return out | |
| def _set_aux_loss(self, outputs_intra_class, outputs_inter_class, outputs_shot_class): | |
| # 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 [{'intra_clip_logits': a1, 'inter_clip_logits': a2, 'pred_shot_logits': b} | |
| for a1, a2, b in zip(outputs_intra_class[:-1], outputs_inter_class[:-1], outputs_shot_class[:-1])] | |
| def build_model(args): | |
| # num_intra_relation_classes is preset in argument | |
| num_intra_relation_classes = args.num_intra_relation_classes | |
| num_inter_relation_classes = args.num_inter_relation_classes | |
| # Init Model | |
| backbone = build_backbone(args) | |
| transformer = build_transformer(args) | |
| model = OmniShotCut( | |
| backbone, | |
| transformer, | |
| num_intra_relation_classes = num_intra_relation_classes, | |
| num_inter_relation_classes = num_inter_relation_classes, | |
| num_frames = args.max_process_window_length, | |
| num_queries = args.num_queries, | |
| aux_loss = args.aux_loss, | |
| ) | |
| # Return | |
| return model | |