content: # ------------------------------------------------------------------------
  # UP-DETR
  # Copyright (c) Tencent, 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
  # ------------------------------------------------------------------------
  import torch
  
  from models.backbone import build_backbone
  from models.detr import DETR, SetCriterion, PostProcess
  from models.matcher import build_matcher
  from models.segmentation import DETRsegm, PostProcessSegm, PostProcessPanoptic
  from models.transformer import build_transformer
  from models.updetr import UPDETR
  
  
  def build_model(args):
      num_classes = 20 if args.dataset_file != 'coco' else 91
      if args.dataset_file == "coco_panoptic":
          num_classes = 250
      if args.dataset_file=="ImageNet":
          num_classes = 2 # feel free to this num_classes, positive integer larger than 1 is OK.
      device = torch.device(args.device)
  
      backbone = build_backbone(args)
  
      transformer = build_transformer(args)
  
      if args.dataset_file=="ImageNet":
          model = UPDETR(
              backbone,
              transformer,
              num_classes=num_classes,
              num_queries=args.num_queries,
              aux_loss=args.aux_loss,
              num_patches=args.num_patches,
              feature_recon=args.feature_recon,
              query_shuffle=args.query_shuffle)
      else:
          model = DETR(
              backbone,
              transformer,
              num_classes=num_classes,
              num_queries=args.num_queries,
              aux_loss=args.aux_loss,
          )
          if args.masks:
              model = DETRsegm(model, freeze_detr=(args.frozen_weights is not None))
      matcher = build_matcher(args)
      weight_dict = {'loss_ce': 1, 'loss_bbox': args.bbox_loss_coef, 'loss_giou': args.giou_loss_coef}
      if args.dataset_file == 'ImageNet' and args.feature_recon:
          weight_dict['loss_feature'] = 1
      if args.masks:
          weight_dict["loss_mask"] = args.mask_loss_coef
          weight_dict["loss_dice"] = args.dice_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.dataset_file == 'ImageNet' and args.feature_recon:
          losses += ['feature']
      if args.masks:
          losses += ["masks"]
      criterion = SetCriterion(num_classes, matcher=matcher, weight_dict=weight_dict,
                               eos_coef=args.eos_coef, losses=losses)
      criterion.to(device)
      postprocessors = {'bbox': PostProcess()}
      if args.masks:
          postprocessors['segm'] = PostProcessSegm()
          if args.dataset_file == "coco_panoptic":
              is_thing_map = {i: i <= 90 for i in range(201)}
              postprocessors["panoptic"] = PostProcessPanoptic(is_thing_map, threshold=0.85)
  
      return model, criterion, postprocessors