make vidt support ymir

.dockerignore

@@ -0,0 +1,5 @@
+*.pth
+__pycache__
+*.egg-info
+ops/build
+ops/dist

.flake8

@@ -0,0 +1,26 @@
+[flake8]
+ignore =
+    # 'toml' imported but unused
+    # F401,
+    # closing bracket does not match visual indentation
+    # E124,
+    # continuation line over-indented for hanging indent
+    # E126,
+    # visually indented line with same indent as next logical line,
+    # E129,
+    # line break before binary operator
+    # W503,
+    # line break after binary operator
+    W504,
+    # camelcase 'PostProcess' imported as lowercase 'postprocess_vidt'
+    N813,
+    # lowercase 'datasets.transforms' imported as non lowercase 'T'
+    N812,
+    # N806 variable 'N' in function should be lowercase
+    N806
+    N803
+    N813
+
+disable-noqa
+max-line-length = 121
+exclude = methods/**, datasets/*, util/*, ops/*, arguments.py, engine.py, fps_calculator.py, main.py

.gitignore

@@ -0,0 +1,5 @@
+*.pth
+__pycache__
+*.egg-info
+ops/build
+ops/dist

.style.yapf

@@ -0,0 +1,3 @@
+[style]
+based_on_style = pep8
+column_limit = 120

arguments.py

@@ -7,6 +7,7 @@
 
 import argparse
 
+
 def str2bool(v, bool):
 
     if isinstance(v, bool):
@@ -77,15 +78,18 @@ def get_args_parser():
     parser.add_argument('--remove_difficult', action='store_true')
 
     # * Device and Log
-    parser.add_argument('--output_dir', default='',
+    parser.add_argument('--output_dir', default='/out/models',
                         help='path where to save, empty for no saving')
+    parser.add_argument('--tensorboard_dir', default='/out/tensorboard', help='path to save tensorboard logs')
     parser.add_argument('--device', default='cuda',
                         help='device to use for training / testing')
     parser.add_argument('--seed', default=42, type=int)
     parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
                         help='start epoch')
     parser.add_argument('--eval', default=False, type=lambda x: (str(x).lower() == 'true'), help='eval mode')
     parser.add_argument('--resume', default='', help='resume from checkpoint')
+    parser.add_argument('--load_from', default='', help='load from checkpoint, not work with --resume')
+    parser.add_argument('--save_interval', default=100, type=int, help='the interval to save the checkpoint')
 
     # * Training setup
     parser.add_argument('--dist-url', default='tcp://127.0.0.1:3457', type=str,
@@ -137,13 +141,38 @@ def get_args_parser():
     parser.add_argument('--distil_model_path', default=None, type=str, help="Distillation model path to load")
     #######
 
-    # cross-scale fusion
-    parser.add_argument('--cross_scale_fusion', default=False, type=lambda x: (str(x).lower() == 'true'), help='use of scale fusion')
+    ####### For ViDT+
+
+    ## EPFF
+    parser.add_argument('--epff', default=False, type=lambda x: (str(x).lower() == 'true'),
+                        help='use of EPFF module for pyramid feature fusion')
+
+    ## UQR
+    parser.add_argument('--masks', default=False, action='store_true', help="Train segmentation head if the flag is provided")
+    parser.add_argument('--with_vector', default=False, type=bool)
+    parser.add_argument('--n_keep', default=256, type=int,
+                        help="Number of coeffs to be remained")
+    parser.add_argument('--gt_mask_len', default=128, type=int,
+                        help="Size of target mask")
+    parser.add_argument('--vector_loss_coef', default=3.0, type=float)
+    parser.add_argument('--vector_hidden_dim', default=256, type=int,
+                        help="Size of the vector embeddings (dimension of the transformer)")
+    parser.add_argument('--no_vector_loss_norm', default=False, action='store_true')
+    parser.add_argument('--activation', default='relu', type=str, help="Activation function to use")
+    parser.add_argument('--checkpoint', default=False, action='store_true')
+    parser.add_argument('--vector_start_stage', default=0, type=int)
+    parser.add_argument('--num_machines', default=1, type=int)
+    parser.add_argument('--loss_type', default='l1', type=str)
+    parser.add_argument('--dcn', default=False, action='store_true')
+
+    ## New losses
     # iou-aware
-    parser.add_argument('--iou_aware', default=False, type=lambda x: (str(x).lower() == 'true'), help='use of iou-aware loss')
+    parser.add_argument('--iou_aware', default=False, type=lambda x: (str(x).lower() == 'true'),
+                        help='use of iou-aware loss')
     parser.add_argument('--iouaware_loss_coef', default=2, type=float)
     # token label
-    parser.add_argument('--token_label', default=False, type=lambda x: (str(x).lower() == 'true'), help='use of token label loss')
+    parser.add_argument('--token_label', default=False, type=lambda x: (str(x).lower() == 'true'),
+                        help='use of token label loss')
     parser.add_argument('--token_loss_coef', default=2, type=float)
     #######
 

datasets/transforms.py

@@ -4,14 +4,15 @@
 """
 import random
 
+import numpy as np
 import PIL
 import torch
 import torchvision.transforms as T
 import torchvision.transforms.functional as F
 
 from util.box_ops import box_xyxy_to_cxcywh
 from util.misc import interpolate
-import numpy as np
+
 
 def crop(image, target, region):
     cropped_image = F.crop(image, *region)
@@ -74,6 +75,9 @@ def hflip(image, target):
 
 
 def resize(image, target, size, max_size=None):
+    """
+    PIL image input
+    """
     # size can be min_size (scalar) or (w, h) tuple
     # import pdb;pdb.set_trace()
     maxs = size
@@ -142,7 +146,7 @@ def get_size(image_size, size, max_size=None):
     if "masks" in target:
         target['masks'] = interpolate(
             target['masks'][:, None].float(), size, mode="nearest")[:, 0] > 0.5
-    
+
     # max_size = max(rescaled_image.size)
     # maxs = max(size)
     # padding_size = [maxs-rescaled_image.size[0], maxs-rescaled_image.size[1]]
@@ -270,6 +274,7 @@ def __call__(self, image, target=None):
         target = target.copy()
         h, w = image.shape[-2:]
         if "boxes" in target:
+            target['xyxy_boxes'] = target["boxes"]
             boxes = target["boxes"]
             boxes = box_xyxy_to_cxcywh(boxes)
             boxes = boxes / torch.tensor([w, h, w, h], dtype=torch.float32)

datasets/voc.py

@@ -25,11 +25,12 @@
 #         c1, c2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
 #         cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
 #     cv2.imwrite(file_path, img)
-    
+
 
 class CocoDetection(torchvision.datasets.CocoDetection):
-    def __init__(self, img_folder, ann_file, transforms, return_masks):
+    def __init__(self, image_set, img_folder, ann_file, transforms, return_masks):
         super(CocoDetection, self).__init__(img_folder, ann_file)
+        self.image_set = image_set
         self._transforms = transforms
         self.prepare = ConvertCocoPolysToMask(return_masks)
 
@@ -125,7 +126,7 @@ def __call__(self, image, target):
         return image, target
 
 
-def make_coco_transforms(image_set):
+def make_coco_transforms(image_set, args):
 
     normalize = T.Compose([
         T.ToTensor(),
@@ -170,7 +171,7 @@ def make_coco_transforms(image_set):
 
     if image_set == 'val':
         return T.Compose([
-            T.RandomResize([512], max_size=800),
+            T.RandomResize([args.eval_size], max_size=max(args.eval_size, 800)),
             normalize,
         ])
 

datasets/ymir.py

@@ -0,0 +1,21 @@
+from pathlib import Path
+
+from ymir_exc.util import convert_ymir_to_coco
+
+from .voc import CocoDetection, make_coco_transforms
+
+
+def build(image_set, args):
+    root = Path(args.coco_path)
+    assert root.exists(), f'provided ymir path {root} does not exist'
+
+    data_info = convert_ymir_to_coco(cat_id_from_zero=True)
+
+    datasets = []
+    for split in image_set:
+        img_folder = data_info[split]['img_dir']
+        ann_file = data_info[split]['ann_file']
+        dataset = CocoDetection(split, img_folder, ann_file,
+            transforms=make_coco_transforms(split, args), return_masks=False)
+        datasets.append(dataset)
+    return datasets