| 366 | |
| 367 | |
| 368 | class MOTR(nn.Module): |
| 369 | def __init__(self, backbone, transformer, num_classes, num_queries, num_feature_levels, criterion, track_embed, |
| 370 | aux_loss=True, with_box_refine=False, two_stage=False, memory_bank=None): |
| 371 | """ Initializes the model. |
| 372 | Parameters: |
| 373 | backbone: torch module of the backbone to be used. See backbone.py |
| 374 | transformer: torch module of the transformer architecture. See transformer.py |
| 375 | num_classes: number of object classes |
| 376 | num_queries: number of object queries, ie detection slot. This is the maximal number of objects |
| 377 | DETR can detect in a single image. For COCO, we recommend 100 queries. |
| 378 | aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used. |
| 379 | with_box_refine: iterative bounding box refinement |
| 380 | two_stage: two-stage Deformable DETR |
| 381 | """ |
| 382 | super().__init__() |
| 383 | self.num_queries = num_queries |
| 384 | self.track_embed = track_embed |
| 385 | self.transformer = transformer |
| 386 | hidden_dim = transformer.d_model |
| 387 | self.num_classes = num_classes |
| 388 | self.class_embed = nn.Linear(hidden_dim, num_classes) |
| 389 | self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3) |
| 390 | self.num_feature_levels = num_feature_levels |
| 391 | if not two_stage: |
| 392 | self.query_embed = nn.Embedding(num_queries, hidden_dim * 2) |
| 393 | if num_feature_levels > 1: |
| 394 | num_backbone_outs = len(backbone.strides) |
| 395 | input_proj_list = [] |
| 396 | for _ in range(num_backbone_outs): |
| 397 | in_channels = backbone.num_channels[_] |
| 398 | input_proj_list.append(nn.Sequential( |
| 399 | nn.Conv2d(in_channels, hidden_dim, kernel_size=1), |
| 400 | nn.GroupNorm(32, hidden_dim), |
| 401 | )) |
| 402 | for _ in range(num_feature_levels - num_backbone_outs): |
| 403 | input_proj_list.append(nn.Sequential( |
| 404 | nn.Conv2d(in_channels, hidden_dim, kernel_size=3, stride=2, padding=1), |
| 405 | nn.GroupNorm(32, hidden_dim), |
| 406 | )) |
| 407 | in_channels = hidden_dim |
| 408 | self.input_proj = nn.ModuleList(input_proj_list) |
| 409 | else: |
| 410 | self.input_proj = nn.ModuleList([ |
| 411 | nn.Sequential( |
| 412 | nn.Conv2d(backbone.num_channels[0], hidden_dim, kernel_size=1), |
| 413 | nn.GroupNorm(32, hidden_dim), |
| 414 | )]) |
| 415 | self.backbone = backbone |
| 416 | self.aux_loss = aux_loss |
| 417 | self.with_box_refine = with_box_refine |
| 418 | self.two_stage = two_stage |
| 419 | |
| 420 | prior_prob = 0.01 |
| 421 | bias_value = -math.log((1 - prior_prob) / prior_prob) |
| 422 | self.class_embed.bias.data = torch.ones(num_classes) * bias_value |
| 423 | nn.init.constant_(self.bbox_embed.layers[-1].weight.data, 0) |
| 424 | nn.init.constant_(self.bbox_embed.layers[-1].bias.data, 0) |
| 425 | for proj in self.input_proj: |