feat: initial commit

.gitignore

@@ -0,0 +1,13 @@
+#storage folders
+storage/**
+
+error_causing_batch.pth
+
+configs/__pycache__
+src/__pycache__
+src/detection/__pycache__
+tools/__pycache__
+src/detection/vision/__pycache__
+
+custom_out.png
+baseline_out.png

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 sir3mat
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,13 +1,75 @@
----
-title: SchoolInAiProjectWork
-emoji: 🏢
-colorFrom: yellow
-colorTo: purple
-sdk: gradio
-sdk_version: 3.9
-app_file: app.py
-pinned: false
-license: mit
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# School in AI Project Work
+
+This repository contains the code to train and evaluate a pedestrian detector for 
+the "School in Ai 2° edition"@[@UNIMORE](https://www.unimore.it/)
+
+## Installation
+
+N.B.: Installation only avaiable in win64 environments
+
+Create and activate an environment with all required packages:
+
+```
+conda create --name ped_detector --file deps/wins/conda_environment.txt
+# or conda env create -f deps/win/conda_environment.yml
+conda activate cvcspw
+pip install -r deps/win/pip_requirements.txt
+```
+
+## Dataset download and preparation:
+### Solution 1 - From Google Drice
+Download the storage folder directly from Google Drive [here](link google drive)
+and place it in the root dir of the project
+After runnning this step, your storage directory should look like this:
+```text
+storage
+    ├── MOTChallenge
+        ├── MOT17
+        ├── motcha_coco_annotations
+    ├── MOTSynth
+        ├── annotations
+        ├── comb_annotations
+        ├── frames
+    ├── motsynth_output
+```
+### Solution 2 - From scratch
+#### Prepare MOTSynth dataset
+1. Download MOTSynth_1.
+```
+wget -P ./storage/MOTSynth https://motchallenge.net/data/MOTSynth_1.zip
+unzip ./storage/MOTSynth/MOTSynth_1.zip
+rm ./storage/MOTSynth/MOTSynth_1.zip
+```
+2. Delete video from 123 to 256
+2. Extract frames from the videos
+```
+python tools/anns/to_frames.py --motsynth-root ./storage/MOTSynth
+
+# now you can delete other videos
+rm -r ./storage/MOTSynth/MOTSynth_1
+```
+3. Download and extract annotations
+```
+wget -P ./storage/MOTSynth https://motchallenge.net/data/MOTSynth_coco_annotations.zip
+unzip ./storage/MOTSynth/MOTSynth_coco_annotations.zip
+rm ./storage/MOTSynth/MOTSynth_coco_annotations.zip
+```
+4. Prepare combined annotations for MOTSynth from the original coco annotations
+```
+python tools/anns/combine_anns.py --motsynth-path ./storage/MOTSynth
+```
+#### Prepare MOT17 dataset
+
+
+## Colab Usage
+
+You can also use [Google Colab](https://colab.research.google.com) if you need remote resources like GPUs.
+In the notebook folder you can find some useful .ipynb files and remember to load the storage folder in your GDrive before usage.
+
+## Object Detection
+
+An adaption of torchvision's detection reference code is done to train Faster R-CNN on a portion of the MOTSynth dataset. To train the model you can run:
+```
+./scripts/train_detector
+```
+

app.py

@@ -0,0 +1,63 @@
+import os.path as osp
+from tkinter.ttk import Style
+import gradio as gr
+import torch
+import logging
+import torchvision
+from torchvision.models.detection.faster_rcnn import fasterrcnn_resnet50_fpn_v2
+from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
+from configs.path_cfg import MOTCHA_ROOT, OUTPUT_DIR
+from src.detection.graph_utils import add_bbox
+from src.detection.vision import presets
+logging.getLogger('PIL').setLevel(logging.CRITICAL)
+
+
+def load_model(baseline: bool = False):
+    if baseline:
+        model = fasterrcnn_resnet50_fpn_v2(
+            weights="DEFAULT")
+    else:
+        model = fasterrcnn_resnet50_fpn_v2()
+        in_features = model.roi_heads.box_predictor.cls_score.in_features
+        model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 2)
+        checkpoint = torch.load(osp.join(OUTPUT_DIR, "detection_logs",
+                                "fasterrcnn_training", "checkpoint.pth"), map_location="cpu")
+        model.load_state_dict(checkpoint["model"])
+    model.eval()
+    return model
+
+
+def detect_with_resnet50Model_finetuning_motsynth(image):
+    model = load_model()
+    transformEval = presets.DetectionPresetEval()
+    image_tensor = transformEval(image, None)[0]
+    prediction = model([image_tensor])[0]
+    image_w_bbox = add_bbox(image_tensor, prediction, 0.85)
+    torchvision.io.write_png(image_w_bbox, "custom_out.png")
+    return "custom_out.png"
+
+
+def detect_with_resnet50Model_baseline(image):
+    model = load_model(baseline=True)
+    transformEval = presets.DetectionPresetEval()
+    image_tensor = transformEval(image, None)[0]
+    prediction = model([image_tensor])[0]
+    image_w_bbox = add_bbox(image_tensor, prediction, 0.85)
+    torchvision.io.write_png(image_w_bbox, "baseline_out.png")
+    return "baseline_out.png"
+
+
+title = "Performance comparision of Faster R-CNN for people detection with syntetic data"
+description = "<p style='text-align: center'>Performance comparision of Faster R-CNN models for people detecion using MOTSynth and MOT17"
+examples = [[osp.join(MOTCHA_ROOT, "MOT17", "train",
+                      "MOT17-09-DPM", "img1", "000001.jpg")]]
+
+
+io_baseline = gr.Interface(detect_with_resnet50Model_baseline, gr.Image(type="pil"), gr.Image(
+    type="file", shape=(1920, 1080), label="FasterR-CNN_Resnet50_COCO"))
+
+io_custom = gr.Interface(detect_with_resnet50Model_finetuning_motsynth, gr.Image(type="pil"), gr.Image(
+    type="file", shape=(1920, 1080), label="FasterR-CNN_Resnet50_FinteTuning_MOTSynth"))
+
+gr.Parallel(io_baseline, io_custom, title=title,
+            description=description, examples=examples).launch(enable_queue=True)

configs/path_cfg.py

@@ -0,0 +1,19 @@
+import os
+import sys
+import os
+
+IN_COLAB = False
+if 'COLAB_GPU' in os.environ:
+   IN_COLAB=True
+
+cwd = os.getcwd()
+
+if(IN_COLAB):
+    MOTSYNTH_ROOT = '/content/gdrive/MyDrive/CVCS/storage/MOTSynth'
+    MOTCHA_ROOT = '/content/gdrive/MyDrive/CVCS/storage/MOTChallenge'
+    OUTPUT_DIR = '/content/gdrive/MyDrive/CVCS/storage/motsynth_output'
+else:
+    # windows config
+    MOTSYNTH_ROOT = cwd + '\storage\MOTSynth'
+    MOTCHA_ROOT = cwd + '\storage\MOTChallenge'
+    OUTPUT_DIR = cwd + '\storage\motsynth_output'

deps/win/conda_environment.yml

@@ -0,0 +1,42 @@
+name: cvcspw
+channels:
+  - defaults
+dependencies:
+  - ca-certificates=2022.07.19=haa95532_0
+  - certifi=2022.6.15=py38haa95532_0
+  - openssl=1.1.1q=h2bbff1b_0
+  - pip=22.1.2=py38haa95532_0
+  - python=3.8.13=h6244533_0
+  - setuptools=63.4.1=py38haa95532_0
+  - sqlite=3.39.2=h2bbff1b_0
+  - vc=14.2=h21ff451_1
+  - vs2015_runtime=14.27.29016=h5e58377_2
+  - wheel=0.37.1=pyhd3eb1b0_0
+  - wincertstore=0.2=py38haa95532_2
+  - pip:
+    - charset-normalizer==2.1.1
+    - coloredlogs==15.0.1
+    - cycler==0.11.0
+    - fonttools==4.37.1
+    - humanfriendly==10.0
+    - idna==3.3
+    - kiwisolver==1.4.4
+    - matplotlib==3.5.3
+    - numpy==1.23.2
+    - packaging==21.3
+    - pandas==1.4.4
+    - pillow==9.2.0
+    - pycocotools==2.0.4
+    - pyparsing==3.0.9
+    - pyreadline3==3.4.1
+    - python-dateutil==2.8.2
+    - pytz==2022.2.1
+    - requests==2.28.1
+    - seaborn==0.12.0
+    - six==1.16.0
+    - torch==1.12.1+cu116
+    - torchaudio==0.12.1+cu116
+    - torchvision==0.13.1+cu116
+    - typing-extensions==4.3.0
+    - urllib3==1.26.12
+prefix: C:\Users\matte\anaconda3\envs\cvcspw

deps/win/conda_requirements.txt

@@ -0,0 +1,39 @@
+# This file may be used to create an environment using:
+# $ conda create --name <env> --file <this file>
+# platform: win-64
+ca-certificates=2022.07.19=haa95532_0
+certifi=2022.6.15=py38haa95532_0
+charset-normalizer=2.1.1=pypi_0
+coloredlogs=15.0.1=pypi_0
+cycler=0.11.0=pypi_0
+fonttools=4.37.1=pypi_0
+humanfriendly=10.0=pypi_0
+idna=3.3=pypi_0
+kiwisolver=1.4.4=pypi_0
+matplotlib=3.5.3=pypi_0
+numpy=1.23.2=pypi_0
+openssl=1.1.1q=h2bbff1b_0
+packaging=21.3=pypi_0
+pandas=1.4.4=pypi_0
+pillow=9.2.0=pypi_0
+pip=22.1.2=py38haa95532_0
+pycocotools=2.0.4=pypi_0
+pyparsing=3.0.9=pypi_0
+pyreadline3=3.4.1=pypi_0
+python=3.8.13=h6244533_0
+python-dateutil=2.8.2=pypi_0
+pytz=2022.2.1=pypi_0
+requests=2.28.1=pypi_0
+seaborn=0.12.0=pypi_0
+setuptools=63.4.1=py38haa95532_0
+six=1.16.0=pypi_0
+sqlite=3.39.2=h2bbff1b_0
+torch=1.12.1+cu116=pypi_0
+torchaudio=0.12.1+cu116=pypi_0
+torchvision=0.13.1+cu116=pypi_0
+typing-extensions=4.3.0=pypi_0
+urllib3=1.26.12=pypi_0
+vc=14.2=h21ff451_1
+vs2015_runtime=14.27.29016=h5e58377_2
+wheel=0.37.1=pyhd3eb1b0_0
+wincertstore=0.2=py38haa95532_2

deps/win/pip_requirements.txt

@@ -0,0 +1,30 @@
+certifi==2022.6.15
+charset-normalizer==2.1.1
+coloredlogs==15.0.1
+cycler==0.11.0
+fonttools==4.37.1
+humanfriendly==10.0
+idna==3.3
+kiwisolver==1.4.4
+matplotlib==3.5.3
+numpy==1.23.2
+packaging==21.3
+pandas==1.4.4
+Pillow==9.2.0
+pip==22.1.2
+pycocotools==2.0.4
+pyparsing==3.0.9
+pyreadline3==3.4.1
+python-dateutil==2.8.2
+pytz==2022.2.1
+requests==2.28.1
+seaborn==0.12.0
+setuptools==63.4.1
+six==1.16.0
+torch==1.12.1
+torchaudio==0.12.1
+torchvision==0.13.1
+typing_extensions==4.3.0
+urllib3==1.26.12
+wheel==0.37.1
+wincertstore==0.2

scripts/evaluate_detector.sh

@@ -0,0 +1,3 @@
+#!/bin/sh
+
+python -m  tools.train_detector --model-eval "d://cvcspw/storage/motsynth_output/detection_logs/fasterrcnn_training/checkpoint.pth" --test-only

scripts/inference_detector.sh

@@ -0,0 +1,3 @@
+#!/bin/sh
+
+python -m  tools.inference_detector  --model-path ./storage/motsynth_output/detection_logs/fasterrcnn_training_2/checkpoint.pth

scripts/train_detector.sh

@@ -0,0 +1,3 @@
+#!/bin/sh
+
+python -m  tools.train_detector

src/detection/graph_utils.py

@@ -0,0 +1,87 @@
+import numpy as np
+import torchvision.transforms.functional as F
+from torchvision import transforms
+from typing import DefaultDict
+import matplotlib.pyplot as plt
+import matplotlib
+import torch
+import logging
+from torchvision.utils import draw_bounding_boxes
+matplotlib.style.use('ggplot')
+logging.getLogger('matplotlib').setLevel(logging.CRITICAL)
+logging.getLogger('PIL').setLevel(logging.CRITICAL)
+
+
+def save_plot(train_loss_list, label, output_dir):
+    """
+    Function to save the loss plot to disk.
+    """
+    # Loss plots.
+    plt.figure(figsize=(10, 7))
+    plt.plot(
+        train_loss_list, linestyle='-',
+        label=label
+    )
+    plt.xlabel('Epochs')
+    plt.ylabel('Loss')
+    plt.legend()
+    plt.savefig(f"{output_dir}/{label}.png")
+
+
+def save_train_loss_plot(train_loss_dict: DefaultDict, output_dir):
+    """
+    Function to save the loss plots to disk.
+    """
+    for key in train_loss_dict.keys():
+        save_plot(train_loss_dict[key], key, output_dir)
+
+
+def show(imgs):
+    if not isinstance(imgs, list):
+        imgs = [imgs]
+    fig, axs = plt.subplots(nrows=len(imgs), ncols=1,
+                            figsize=(45, 21), squeeze=False)
+    for i, img in enumerate(imgs):
+        img = img.detach()
+        img = F.to_pil_image(img)
+        img = np.asarray(img)
+        axs[i, 0].imshow(img)
+        axs[i, 0].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])
+    plt.show()
+
+
+def plot_img_tensor(img_tensor):
+    transforms.ToPILImage()(img_tensor).show()
+
+
+def show_img(data_loader, model, device, th=0.7):
+    for imgs, target in data_loader:
+        with torch.no_grad():
+            prediction = model([imgs[0].to(device)])[0]
+        plot_img_tensor(add_bbox(imgs[0], prediction, th))
+        plot_img_tensor(add_bbox(imgs[0], target[0]['boxes']))
+        break
+
+
+def add_bbox(img, output, th=None):
+    img_canvas = img.clone()
+    img_canvas = torch.clip(img*255, 0, 255)
+    img_canvas = img_canvas.type(torch.uint8)
+
+    if th == None:
+        img_with_bbbox = draw_bounding_boxes(
+            img_canvas, boxes=output, width=4)
+    else:
+        mask = (output["scores"] > th) & (output["labels"] == 1)
+        scores_list = [score for score in (
+            output["scores"][mask]).tolist()]
+        labels_list = [str(label) for label in (
+            output["labels"][mask]).tolist()]
+        labels = ["person" for label in labels_list if label == "1"]
+        assert len(labels) == len(scores_list) == len(labels_list)
+
+        for i in range(0, len(labels)):
+            labels[i] = f"{labels[i]}:{scores_list[i]:.3f}"
+        img_with_bbbox = draw_bounding_boxes(
+            img_canvas, boxes=output["boxes"][mask], labels=labels, width=4)
+    return img_with_bbbox

src/detection/model_factory.py

@@ -0,0 +1,54 @@
+import logging
+import torch
+from torchvision.models.detection.faster_rcnn import fasterrcnn_resnet50_fpn_v2, FasterRCNN, FastRCNNPredictor
+from torchvision.models.detection import FasterRCNN_ResNet50_FPN_V2_Weights
+from torchvision.models.resnet import ResNet50_Weights
+
+logger = logging.getLogger(__name__)
+
+
+def set_seeds(seed: int = 42):
+    """Sets random sets for torch operations.
+
+    Args:
+        seed (int, optional): Random seed to set. Defaults to 42.
+    """
+    # Set the seed for general torch operations
+    torch.manual_seed(seed)
+    # Set the seed for CUDA torch operations (ones that happen on the GPU)
+    torch.cuda.manual_seed(seed)
+
+
+class ModelFactory:
+    @staticmethod
+    def get_model(name, weights, backbone, backbone_weights, trainable_backbone_layers):
+        logger.debug(f"get_model -> model:{name}")
+
+        if name == "fasterrcnn_resnet50_fpn":
+            # backbone = backbone
+            model_weights = FasterRCNN_ResNet50_FPN_V2_Weights[weights]
+            model_backbone_weights = ResNet50_Weights[backbone_weights]
+            # trainable_backbone_layers = 1
+            model: FasterRCNN = fasterrcnn_resnet50_fpn_v2(
+                weights=model_weights, backbone_name=backbone, weights_backbone=model_backbone_weights, trainable_backbone_layers=trainable_backbone_layers)
+
+            # for param in model.rpn.parameters():
+            #     param.requires_grad = False
+            # for param in model.roi_heads.parameters():
+            #     param.requires_grad = False
+            # for param in model.backbone.fpn.parameters():
+            #     param.requires_grad = False
+
+            set_seeds()
+
+            num_classes = 2  # 1 class (person) + background
+            in_features = model.roi_heads.box_predictor.cls_score.in_features
+            model.roi_heads.box_predictor = FastRCNNPredictor(
+                in_features, num_classes)
+
+        else:
+            logger.error(
+                "Please, provide a valid model as argument. Select one of the following: fasterrcnn_resnet50_fpn.")
+            raise ValueError(name)
+
+        return model

src/detection/mot_dataset.py

@@ -0,0 +1,48 @@
+
+import numpy as np
+from src.detection.vision.coco_utils import ConvertCocoPolysToMask, CocoDetection, _coco_remove_images_without_annotations
+from src.detection.vision.transforms import Compose
+
+
+class UpdateIsCrowd(object):
+    def __init__(self, min_size, min_vis=0.2):
+        self.min_size = min_size
+        self.min_vis = min_vis
+
+    def __call__(self, image, target):
+        for i, ann in enumerate(target['annotations']):
+            bbox = ann['bbox']
+            bbox_too_small = max(bbox[-1], bbox[-2]) < self.min_size
+
+            if 'vis' in ann:
+                vis = ann['vis']
+
+            elif 'keypoints' in ann:
+                vis = (np.array(ann['keypoints'])[2::3] == 2).mean().round(2)
+
+            else:
+                raise RuntimeError(
+                    "The given annotations have no visibility measure. Are you sure you want to proceed?")
+
+            not_vis = vis < self.min_vis
+            target['annotations'][i]['iscrowd'] = max(
+                ann['iscrowd'], int(bbox_too_small), int(not_vis))
+
+        return image, target
+
+
+def get_mot_dataset(img_folder, ann_file, transforms, min_size=25, min_vis=0.2):
+    t = [UpdateIsCrowd(min_size=min_size, min_vis=min_vis),
+         ConvertCocoPolysToMask()]
+
+    if transforms is not None:
+        t.append(transforms)
+    transforms = Compose(t)
+
+    dataset = CocoDetection(img_folder=img_folder,
+                            ann_file=ann_file,
+                            transforms=transforms)
+
+    dataset = _coco_remove_images_without_annotations(dataset)
+
+    return dataset

src/detection/vision/README.md

@@ -0,0 +1,88 @@
+# Object detection reference training scripts
+
+This folder contains reference training scripts for object detection.
+They serve as a log of how to train specific models, to provide baseline
+training and evaluation scripts to quickly bootstrap research.
+
+To execute the example commands below you must install the following:
+
+```
+cython
+pycocotools
+matplotlib
+```
+
+You must modify the following flags:
+
+`--data-path=/path/to/coco/dataset`
+
+`--nproc_per_node=<number_of_gpus_available>`
+
+Except otherwise noted, all models have been trained on 8x V100 GPUs. 
+
+### Faster R-CNN ResNet-50 FPN
+```
+torchrun --nproc_per_node=8 train.py\
+    --dataset coco --model fasterrcnn_resnet50_fpn --epochs 26\
+    --lr-steps 16 22 --aspect-ratio-group-factor 3
+```
+
+### Faster R-CNN MobileNetV3-Large FPN
+```
+torchrun --nproc_per_node=8 train.py\
+    --dataset coco --model fasterrcnn_mobilenet_v3_large_fpn --epochs 26\
+    --lr-steps 16 22 --aspect-ratio-group-factor 3
+```
+
+### Faster R-CNN MobileNetV3-Large 320 FPN
+```
+torchrun --nproc_per_node=8 train.py\
+    --dataset coco --model fasterrcnn_mobilenet_v3_large_320_fpn --epochs 26\
+    --lr-steps 16 22 --aspect-ratio-group-factor 3
+```
+
+### FCOS ResNet-50 FPN
+```
+torchrun --nproc_per_node=8 train.py\
+    --dataset coco --model fcos_resnet50_fpn --epochs 26\
+    --lr-steps 16 22 --aspect-ratio-group-factor 3  --lr 0.01 --amp
+```
+
+### RetinaNet
+```
+torchrun --nproc_per_node=8 train.py\
+    --dataset coco --model retinanet_resnet50_fpn --epochs 26\
+    --lr-steps 16 22 --aspect-ratio-group-factor 3 --lr 0.01
+```
+
+### SSD300 VGG16
+```
+torchrun --nproc_per_node=8 train.py\
+    --dataset coco --model ssd300_vgg16 --epochs 120\
+    --lr-steps 80 110 --aspect-ratio-group-factor 3 --lr 0.002 --batch-size 4\
+    --weight-decay 0.0005 --data-augmentation ssd
+```
+
+### SSDlite320 MobileNetV3-Large
+```
+torchrun --nproc_per_node=8 train.py\
+    --dataset coco --model ssdlite320_mobilenet_v3_large --epochs 660\
+    --aspect-ratio-group-factor 3 --lr-scheduler cosineannealinglr --lr 0.15 --batch-size 24\
+    --weight-decay 0.00004 --data-augmentation ssdlite
+```
+
+
+### Mask R-CNN
+```
+torchrun --nproc_per_node=8 train.py\
+    --dataset coco --model maskrcnn_resnet50_fpn --epochs 26\
+    --lr-steps 16 22 --aspect-ratio-group-factor 3
+```
+
+
+### Keypoint R-CNN
+```
+torchrun --nproc_per_node=8 train.py\
+    --dataset coco_kp --model keypointrcnn_resnet50_fpn --epochs 46\
+    --lr-steps 36 43 --aspect-ratio-group-factor 3
+```

src/detection/vision/coco_eval.py

@@ -0,0 +1,194 @@
+import copy
+import io
+from contextlib import redirect_stdout
+
+import numpy as np
+import pycocotools.mask as mask_util
+import torch
+from . import utils
+from pycocotools.coco import COCO
+from pycocotools.cocoeval import COCOeval
+
+
+class CocoEvaluator:
+    def __init__(self, coco_gt, iou_types):
+        assert isinstance(iou_types, (list, tuple))
+        coco_gt = copy.deepcopy(coco_gt)
+        self.coco_gt = coco_gt
+
+        self.iou_types = iou_types
+        self.coco_eval = {}
+        for iou_type in iou_types:
+            self.coco_eval[iou_type] = COCOeval(coco_gt, iouType=iou_type)
+
+        self.img_ids = []
+        self.eval_imgs = {k: [] for k in iou_types}
+
+    def update(self, predictions):
+        img_ids = list(np.unique(list(predictions.keys())))
+        self.img_ids.extend(img_ids)
+
+        for iou_type in self.iou_types:
+            results = self.prepare(predictions, iou_type)
+            with redirect_stdout(io.StringIO()):
+                coco_dt = COCO.loadRes(
+                    self.coco_gt, results) if results else COCO()
+            coco_eval = self.coco_eval[iou_type]
+
+            coco_eval.cocoDt = coco_dt
+            coco_eval.params.imgIds = list(img_ids)
+            img_ids, eval_imgs = evaluate(coco_eval)
+
+            self.eval_imgs[iou_type].append(eval_imgs)
+
+    def synchronize_between_processes(self):
+        for iou_type in self.iou_types:
+            self.eval_imgs[iou_type] = np.concatenate(
+                self.eval_imgs[iou_type], 2)
+            create_common_coco_eval(
+                self.coco_eval[iou_type], self.img_ids, self.eval_imgs[iou_type])
+
+    def accumulate(self):
+        for coco_eval in self.coco_eval.values():
+            coco_eval.accumulate()
+
+    def summarize(self):
+        for iou_type, coco_eval in self.coco_eval.items():
+            print(f"IoU metric: {iou_type}")
+            coco_eval.summarize()
+
+    def prepare(self, predictions, iou_type):
+        if iou_type == "bbox":
+            return self.prepare_for_coco_detection(predictions)
+        if iou_type == "segm":
+            return self.prepare_for_coco_segmentation(predictions)
+        if iou_type == "keypoints":
+            return self.prepare_for_coco_keypoint(predictions)
+        raise ValueError(f"Unknown iou type {iou_type}")
+
+    def prepare_for_coco_detection(self, predictions):
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            boxes = prediction["boxes"]
+            boxes = convert_to_xywh(boxes).tolist()
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+
+            coco_results.extend(
+                [
+                    {
+                        "image_id": original_id,
+                        "category_id": labels[k],
+                        "bbox": box,
+                        "score": scores[k],
+                    }
+                    for k, box in enumerate(boxes)
+                ]
+            )
+        return coco_results
+
+    def prepare_for_coco_segmentation(self, predictions):
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            scores = prediction["scores"]
+            labels = prediction["labels"]
+            masks = prediction["masks"]
+
+            masks = masks > 0.5
+
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+
+            rles = [
+                mask_util.encode(np.array(mask[0, :, :, np.newaxis], dtype=np.uint8, order="F"))[0] for mask in masks
+            ]
+            for rle in rles:
+                rle["counts"] = rle["counts"].decode("utf-8")
+
+            coco_results.extend(
+                [
+                    {
+                        "image_id": original_id,
+                        "category_id": labels[k],
+                        "segmentation": rle,
+                        "score": scores[k],
+                    }
+                    for k, rle in enumerate(rles)
+                ]
+            )
+        return coco_results
+
+    def prepare_for_coco_keypoint(self, predictions):
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            boxes = prediction["boxes"]
+            boxes = convert_to_xywh(boxes).tolist()
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+            keypoints = prediction["keypoints"]
+            keypoints = keypoints.flatten(start_dim=1).tolist()
+
+            coco_results.extend(
+                [
+                    {
+                        "image_id": original_id,
+                        "category_id": labels[k],
+                        "keypoints": keypoint,
+                        "score": scores[k],
+                    }
+                    for k, keypoint in enumerate(keypoints)
+                ]
+            )
+        return coco_results
+
+
+def convert_to_xywh(boxes):
+    xmin, ymin, xmax, ymax = boxes.unbind(1)
+    return torch.stack((xmin, ymin, xmax - xmin, ymax - ymin), dim=1)
+
+
+def merge(img_ids, eval_imgs):
+    all_img_ids = utils.all_gather(img_ids)
+    all_eval_imgs = utils.all_gather(eval_imgs)
+
+    merged_img_ids = []
+    for p in all_img_ids:
+        merged_img_ids.extend(p)
+
+    merged_eval_imgs = []
+    for p in all_eval_imgs:
+        merged_eval_imgs.append(p)
+
+    merged_img_ids = np.array(merged_img_ids)
+    merged_eval_imgs = np.concatenate(merged_eval_imgs, 2)
+
+    # keep only unique (and in sorted order) images
+    merged_img_ids, idx = np.unique(merged_img_ids, return_index=True)
+    merged_eval_imgs = merged_eval_imgs[..., idx]
+
+    return merged_img_ids, merged_eval_imgs
+
+
+def create_common_coco_eval(coco_eval, img_ids, eval_imgs):
+    img_ids, eval_imgs = merge(img_ids, eval_imgs)
+    img_ids = list(img_ids)
+    eval_imgs = list(eval_imgs.flatten())
+
+    coco_eval.evalImgs = eval_imgs
+    coco_eval.params.imgIds = img_ids
+    coco_eval._paramsEval = copy.deepcopy(coco_eval.params)
+
+
+def evaluate(imgs):
+    with redirect_stdout(io.StringIO()):
+        imgs.evaluate()
+    return imgs.params.imgIds, np.asarray(imgs.evalImgs).reshape(-1, len(imgs.params.areaRng), len(imgs.params.imgIds))

src/detection/vision/coco_utils.py

@@ -0,0 +1,263 @@
+import copy
+import os
+
+import torch
+import torch.utils.data
+import torchvision
+from . import transforms as T
+from pycocotools import mask as coco_mask
+from pycocotools.coco import COCO
+
+
+class FilterAndRemapCocoCategories:
+    def __init__(self, categories, remap=True):
+        self.categories = categories
+        self.remap = remap
+
+    def __call__(self, image, target):
+        anno = target["annotations"]
+        anno = [obj for obj in anno if obj["category_id"] in self.categories]
+        if not self.remap:
+            target["annotations"] = anno
+            return image, target
+        anno = copy.deepcopy(anno)
+        for obj in anno:
+            obj["category_id"] = self.categories.index(obj["category_id"])
+        target["annotations"] = anno
+        return image, target
+
+
+def convert_coco_poly_to_mask(segmentations, height, width):
+    masks = []
+    for polygons in segmentations:
+        if isinstance(polygons['counts'], list):
+            rles = coco_mask.frPyObjects(polygons, height, width)
+
+        else:
+            rles = [polygons]
+
+        mask = coco_mask.decode(rles)
+        if len(mask.shape) < 3:
+            mask = mask[..., None]
+        mask = torch.as_tensor(mask, dtype=torch.uint8)
+        mask = mask.any(dim=2)
+        masks.append(mask)
+    if masks:
+        masks = torch.stack(masks, dim=0)
+    else:
+        masks = torch.zeros((0, height, width), dtype=torch.uint8)
+    return masks
+
+
+class ConvertCocoPolysToMask:
+    def __call__(self, image, target):
+        w, h = image.size
+
+        image_id = target["image_id"]
+        image_id = torch.tensor([image_id])
+
+        anno = target["annotations"]
+        anno = [obj for obj in anno if obj["iscrowd"] == 0]
+
+        boxes = [obj["bbox"] for obj in anno]
+        # guard against no boxes via resizing
+        boxes = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4)
+        boxes[:, 2:] += boxes[:, :2]
+        boxes[:, 0::2].clamp_(min=0, max=w)
+        boxes[:, 1::2].clamp_(min=0, max=h)
+
+        classes = [obj["category_id"] for obj in anno]
+        classes = torch.tensor(classes, dtype=torch.int64)
+
+        # masks=None
+        if anno and 'segmentation' in anno[0]:
+            segmentations = [obj["segmentation"] for obj in anno]
+
+        else:
+            segmentations = []
+
+        masks = convert_coco_poly_to_mask(segmentations, h, w)
+
+        keypoints = None
+        if anno and "keypoints" in anno[0]:
+            keypoints = [obj["keypoints"] for obj in anno]
+            keypoints = torch.as_tensor(keypoints, dtype=torch.float32)
+            num_keypoints = keypoints.shape[0]
+            if num_keypoints:
+                keypoints = keypoints.view(num_keypoints, -1, 3)
+
+        keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
+        boxes = boxes[keep]
+        classes = classes[keep]
+        if masks is not None and masks.shape[0] > 0:
+            masks = masks[keep]
+
+        if keypoints is not None:
+            keypoints = keypoints[keep]
+
+        target = {}
+        target["boxes"] = boxes
+        target["labels"] = classes
+        target["masks"] = masks
+
+        target["image_id"] = image_id
+        if keypoints is not None:
+            target["keypoints"] = keypoints
+
+        # for conversion to coco api
+        area = torch.tensor([obj["area"] for obj in anno])
+        iscrowd = torch.tensor([obj["iscrowd"] for obj in anno])
+        target["area"] = area
+        target["iscrowd"] = iscrowd
+        #target['vis'] = [obj['vis'] for obj in anno]
+
+        return image, target
+
+
+def _coco_remove_images_without_annotations(dataset, cat_list=None):
+    def _has_only_empty_bbox(anno):
+        return all(any(o <= 1 for o in obj["bbox"][2:]) for obj in anno)
+
+    def _count_visible_keypoints(anno):
+        return sum(sum(1 for v in ann["keypoints"][2::3] if v > 0) for ann in anno)
+
+    min_keypoints_per_image = 10
+
+    def _has_valid_annotation(anno):
+        # if it's empty, there is no annotation
+        if len(anno) == 0:
+            return False
+        # if all boxes have close to zero area, there is no annotation
+        if _has_only_empty_bbox(anno):
+            return False
+        # keypoints task have a slight different critera for considering
+        # if an annotation is valid
+        if "keypoints" not in anno[0]:
+            return True
+        # for keypoint detection tasks, only consider valid images those
+        # containing at least min_keypoints_per_image
+        if _count_visible_keypoints(anno) >= min_keypoints_per_image:
+            return True
+        return False
+
+    assert isinstance(dataset, torchvision.datasets.CocoDetection)
+    ids = []
+    for ds_idx, img_id in enumerate(dataset.ids):
+        ann_ids = dataset.coco.getAnnIds(imgIds=img_id, iscrowd=None)
+        anno = dataset.coco.loadAnns(ann_ids)
+        if cat_list:
+            anno = [obj for obj in anno if obj["category_id"] in cat_list]
+        if _has_valid_annotation(anno):
+            ids.append(ds_idx)
+
+    dataset = torch.utils.data.Subset(dataset, ids)
+    return dataset
+
+
+def convert_to_coco_api(ds):
+    coco_ds = COCO()
+    # annotation IDs need to start at 1, not 0, see torchvision issue #1530
+    ann_id = 1
+    dataset = {"images": [], "categories": [], "annotations": []}
+    categories = set()
+    for img_idx in range(len(ds)):
+        # find better way to get target
+        # targets = ds.get_annotations(img_idx)
+        img, targets = ds[img_idx]
+        image_id = targets["image_id"].item()
+        img_dict = {}
+        img_dict["id"] = image_id
+        img_dict["height"] = img.shape[-2]
+        img_dict["width"] = img.shape[-1]
+        dataset["images"].append(img_dict)
+        bboxes = targets["boxes"].clone()
+        bboxes[:, 2:] -= bboxes[:, :2]
+        bboxes = bboxes.tolist()
+        labels = targets["labels"].tolist()
+        areas = targets["area"].tolist()
+        iscrowd = targets["iscrowd"].tolist()
+        if "masks" in targets:
+            masks = targets["masks"]
+            # make masks Fortran contiguous for coco_mask
+            masks = masks.permute(0, 2, 1).contiguous().permute(0, 2, 1)
+        if "keypoints" in targets:
+            keypoints = targets["keypoints"]
+            keypoints = keypoints.reshape(keypoints.shape[0], -1).tolist()
+        num_objs = len(bboxes)
+        for i in range(num_objs):
+            ann = {}
+            ann["image_id"] = image_id
+            ann["bbox"] = bboxes[i]
+            ann["category_id"] = labels[i]
+            categories.add(labels[i])
+            ann["area"] = areas[i]
+            ann["iscrowd"] = iscrowd[i]
+            ann["id"] = ann_id
+            if "masks" in targets:
+                ann["segmentation"] = coco_mask.encode(masks[i].numpy())
+            if "keypoints" in targets:
+                ann["keypoints"] = keypoints[i]
+                ann["num_keypoints"] = sum(k != 0 for k in keypoints[i][2::3])
+            dataset["annotations"].append(ann)
+            ann_id += 1
+    dataset["categories"] = [{"id": i} for i in sorted(categories)]
+    coco_ds.dataset = dataset
+    coco_ds.createIndex()
+    return coco_ds
+
+
+def get_coco_api_from_dataset(dataset):
+    for _ in range(10):
+        if isinstance(dataset, torchvision.datasets.CocoDetection):
+            break
+        if isinstance(dataset, torch.utils.data.Subset):
+            dataset = dataset.dataset
+    if isinstance(dataset, torchvision.datasets.CocoDetection):
+        return dataset.coco
+    return convert_to_coco_api(dataset)
+
+
+class CocoDetection(torchvision.datasets.CocoDetection):
+    def __init__(self, img_folder, ann_file, transforms):
+        super().__init__(img_folder, ann_file)
+        self._transforms = transforms
+
+    def __getitem__(self, idx):
+        img, target = super().__getitem__(idx)
+        image_id = self.ids[idx]
+        target = dict(image_id=image_id, annotations=target)
+        if self._transforms is not None:
+            img, target = self._transforms(img, target)
+        return img, target
+
+
+def get_coco(root, image_set, transforms, mode="instances"):
+    anno_file_template = "{}_{}2017.json"
+    PATHS = {
+        "train": ("train2017", os.path.join("annotations", anno_file_template.format(mode, "train"))),
+        "val": ("val2017", os.path.join("annotations", anno_file_template.format(mode, "val"))),
+        # "train": ("val2017", os.path.join("annotations", anno_file_template.format(mode, "val")))
+    }
+
+    t = [ConvertCocoPolysToMask()]
+
+    if transforms is not None:
+        t.append(transforms)
+    transforms = T.Compose(t)
+
+    img_folder, ann_file = PATHS[image_set]
+    img_folder = os.path.join(root, img_folder)
+    ann_file = os.path.join(root, ann_file)
+
+    dataset = CocoDetection(img_folder, ann_file, transforms=transforms)
+
+    if image_set == "train":
+        dataset = _coco_remove_images_without_annotations(dataset)
+
+    # dataset = torch.utils.data.Subset(dataset, [i for i in range(500)])
+
+    return dataset
+
+
+def get_coco_kp(root, image_set, transforms):
+    return get_coco(root, image_set, transforms, mode="person_keypoints")

src/detection/vision/engine.py

@@ -0,0 +1,137 @@
+import math
+import sys
+import time
+
+import torch
+import torchvision.models.detection.faster_rcnn
+from . import utils
+from . import coco_eval
+from . import coco_utils
+
+
+def train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq, scaler=None):
+    model.train()
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    metric_logger.add_meter("lr", utils.SmoothedValue(
+        window_size=1, fmt="{value:.6f}"))
+    header = f"Epoch: [{epoch}]"
+    lr_scheduler = None
+    if epoch == 0:
+        warmup_factor = 1.0 / 1000
+        warmup_iters = min(1000, len(data_loader) - 1)
+
+        lr_scheduler = torch.optim.lr_scheduler.LinearLR(
+            optimizer, start_factor=warmup_factor, total_iters=warmup_iters
+        )
+
+    losses_dict = {
+        "lr": [],
+        "loss": [],
+        # loss rpn
+        "loss_objectness": [],
+        "loss_rpn_box_reg": [],
+        # roi heads
+        "loss_classifier": [],
+        "loss_box_reg": [],
+    }
+    for images, targets in metric_logger.log_every(data_loader, print_freq, header):
+        try:
+            images = list(image.to(device) for image in images)
+            targets = [{k: v.to(device) for k, v in t.items()}
+                       for t in targets]
+            with torch.cuda.amp.autocast(enabled=scaler is not None):
+                loss_dict = model(images, targets)
+                losses = sum(loss for loss in loss_dict.values())
+
+            # reduce losses over all GPUs for logging purposes
+            loss_dict_reduced = utils.reduce_dict(loss_dict)
+            losses_reduced = sum(loss for loss in loss_dict_reduced.values())
+
+            loss_value = losses_reduced.item()
+
+            # if problem with loss see below
+            if not math.isfinite(loss_value):
+                print(f"Loss is {loss_value}, stopping training")
+                print(loss_dict_reduced)
+                sys.exit(1)
+
+        except Exception as exp:
+            print("ERROR", str(exp))
+            torch.save({'img': images, 'targets': targets},
+                       'error_causing_batch.pth')
+            raise RuntimeError
+
+        optimizer.zero_grad()
+        if scaler is not None:
+            scaler.scale(losses).backward()
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            losses.backward()
+            optimizer.step()
+
+        if lr_scheduler is not None:
+            lr_scheduler.step()
+
+        metric_logger.update(loss=losses_reduced, **loss_dict_reduced)
+        metric_logger.update(lr=optimizer.param_groups[0]["lr"])
+        for name, meter in metric_logger.meters.items():
+            losses_dict[name].append(meter.global_avg)
+    return metric_logger, losses_dict
+
+
+def _get_iou_types(model):
+    model_without_ddp = model
+    if isinstance(model, torch.nn.parallel.DistributedDataParallel):
+        model_without_ddp = model.module
+    iou_types = ["bbox"]
+    if isinstance(model_without_ddp, torchvision.models.detection.MaskRCNN):
+        iou_types.append("segm")
+    if isinstance(model_without_ddp, torchvision.models.detection.KeypointRCNN):
+        iou_types.append("keypoints")
+    return iou_types
+
+
+@ torch.inference_mode()
+def evaluate(model, data_loader, device, iou_types=None):
+    n_threads = torch.get_num_threads()
+    # FIXME remove this and make paste_masks_in_image run on the GPU
+    torch.set_num_threads(1)
+    cpu_device = torch.device("cpu")
+    model.eval()
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    header = "Test:"
+    coco = coco_utils.get_coco_api_from_dataset(data_loader.dataset)
+    if iou_types is None:
+        iou_types = _get_iou_types(model)
+    coco_evaluator = coco_eval.CocoEvaluator(coco, iou_types)
+    for images, targets in metric_logger.log_every(data_loader, 100, header):
+        images = list(img.to(device) for img in images)
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        model_time = time.time()
+        outputs = model(images)
+
+        outputs = [{k: v.to(cpu_device) for k, v in t.items()}
+                   for t in outputs]
+        model_time = time.time() - model_time
+
+        res = {target["image_id"].item(): output for target,
+               output in zip(targets, outputs)}
+        evaluator_time = time.time()
+        coco_evaluator.update(res)
+        evaluator_time = time.time() - evaluator_time
+        metric_logger.update(model_time=model_time,
+                             evaluator_time=evaluator_time)
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    coco_evaluator.synchronize_between_processes()
+
+    # accumulate predictions from all images and print table with results
+    coco_evaluator.accumulate()
+    coco_evaluator.summarize()
+    torch.set_num_threads(n_threads)
+    return coco_evaluator

src/detection/vision/group_by_aspect_ratio.py

@@ -0,0 +1,196 @@
+import bisect
+import copy
+import math
+from collections import defaultdict
+from itertools import repeat, chain
+
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from PIL import Image
+from torch.utils.data.sampler import BatchSampler, Sampler
+from torch.utils.model_zoo import tqdm
+
+
+def _repeat_to_at_least(iterable, n):
+    repeat_times = math.ceil(n / len(iterable))
+    repeated = chain.from_iterable(repeat(iterable, repeat_times))
+    return list(repeated)
+
+
+class GroupedBatchSampler(BatchSampler):
+    """
+    Wraps another sampler to yield a mini-batch of indices.
+    It enforces that the batch only contain elements from the same group.
+    It also tries to provide mini-batches which follows an ordering which is
+    as close as possible to the ordering from the original sampler.
+    Args:
+        sampler (Sampler): Base sampler.
+        group_ids (list[int]): If the sampler produces indices in range [0, N),
+            `group_ids` must be a list of `N` ints which contains the group id of each sample.
+            The group ids must be a continuous set of integers starting from
+            0, i.e. they must be in the range [0, num_groups).
+        batch_size (int): Size of mini-batch.
+    """
+
+    def __init__(self, sampler, group_ids, batch_size):
+        if not isinstance(sampler, Sampler):
+            raise ValueError(f"sampler should be an instance of torch.utils.data.Sampler, but got sampler={sampler}")
+        self.sampler = sampler
+        self.group_ids = group_ids
+        self.batch_size = batch_size
+
+    def __iter__(self):
+        buffer_per_group = defaultdict(list)
+        samples_per_group = defaultdict(list)
+
+        num_batches = 0
+        for idx in self.sampler:
+            group_id = self.group_ids[idx]
+            buffer_per_group[group_id].append(idx)
+            samples_per_group[group_id].append(idx)
+            if len(buffer_per_group[group_id]) == self.batch_size:
+                yield buffer_per_group[group_id]
+                num_batches += 1
+                del buffer_per_group[group_id]
+            assert len(buffer_per_group[group_id]) < self.batch_size
+
+        # now we have run out of elements that satisfy
+        # the group criteria, let's return the remaining
+        # elements so that the size of the sampler is
+        # deterministic
+        expected_num_batches = len(self)
+        num_remaining = expected_num_batches - num_batches
+        if num_remaining > 0:
+            # for the remaining batches, take first the buffers with largest number
+            # of elements
+            for group_id, _ in sorted(buffer_per_group.items(), key=lambda x: len(x[1]), reverse=True):
+                remaining = self.batch_size - len(buffer_per_group[group_id])
+                samples_from_group_id = _repeat_to_at_least(samples_per_group[group_id], remaining)
+                buffer_per_group[group_id].extend(samples_from_group_id[:remaining])
+                assert len(buffer_per_group[group_id]) == self.batch_size
+                yield buffer_per_group[group_id]
+                num_remaining -= 1
+                if num_remaining == 0:
+                    break
+        assert num_remaining == 0
+
+    def __len__(self):
+        return len(self.sampler) // self.batch_size
+
+
+def _compute_aspect_ratios_slow(dataset, indices=None):
+    print(
+        "Your dataset doesn't support the fast path for "
+        "computing the aspect ratios, so will iterate over "
+        "the full dataset and load every image instead. "
+        "This might take some time..."
+    )
+    if indices is None:
+        indices = range(len(dataset))
+
+    class SubsetSampler(Sampler):
+        def __init__(self, indices):
+            self.indices = indices
+
+        def __iter__(self):
+            return iter(self.indices)
+
+        def __len__(self):
+            return len(self.indices)
+
+    sampler = SubsetSampler(indices)
+    data_loader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=1,
+        sampler=sampler,
+        num_workers=14,  # you might want to increase it for faster processing
+        collate_fn=lambda x: x[0],
+    )
+    aspect_ratios = []
+    with tqdm(total=len(dataset)) as pbar:
+        for _i, (img, _) in enumerate(data_loader):
+            pbar.update(1)
+            height, width = img.shape[-2:]
+            aspect_ratio = float(width) / float(height)
+            aspect_ratios.append(aspect_ratio)
+    return aspect_ratios
+
+
+def _compute_aspect_ratios_custom_dataset(dataset, indices=None):
+    if indices is None:
+        indices = range(len(dataset))
+    aspect_ratios = []
+    for i in indices:
+        height, width = dataset.get_height_and_width(i)
+        aspect_ratio = float(width) / float(height)
+        aspect_ratios.append(aspect_ratio)
+    return aspect_ratios
+
+
+def _compute_aspect_ratios_coco_dataset(dataset, indices=None):
+    if indices is None:
+        indices = range(len(dataset))
+    aspect_ratios = []
+    for i in indices:
+        img_info = dataset.coco.imgs[dataset.ids[i]]
+        aspect_ratio = float(img_info["width"]) / float(img_info["height"])
+        aspect_ratios.append(aspect_ratio)
+    return aspect_ratios
+
+
+def _compute_aspect_ratios_voc_dataset(dataset, indices=None):
+    if indices is None:
+        indices = range(len(dataset))
+    aspect_ratios = []
+    for i in indices:
+        # this doesn't load the data into memory, because PIL loads it lazily
+        width, height = Image.open(dataset.images[i]).size
+        aspect_ratio = float(width) / float(height)
+        aspect_ratios.append(aspect_ratio)
+    return aspect_ratios
+
+
+def _compute_aspect_ratios_subset_dataset(dataset, indices=None):
+    if indices is None:
+        indices = range(len(dataset))
+
+    ds_indices = [dataset.indices[i] for i in indices]
+    return compute_aspect_ratios(dataset.dataset, ds_indices)
+
+
+def compute_aspect_ratios(dataset, indices=None):
+    if hasattr(dataset, "get_height_and_width"):
+        return _compute_aspect_ratios_custom_dataset(dataset, indices)
+
+    if isinstance(dataset, torchvision.datasets.CocoDetection):
+        return _compute_aspect_ratios_coco_dataset(dataset, indices)
+
+    if isinstance(dataset, torchvision.datasets.VOCDetection):
+        return _compute_aspect_ratios_voc_dataset(dataset, indices)
+
+    if isinstance(dataset, torch.utils.data.Subset):
+        return _compute_aspect_ratios_subset_dataset(dataset, indices)
+
+    # slow path
+    return _compute_aspect_ratios_slow(dataset, indices)
+
+
+def _quantize(x, bins):
+    bins = copy.deepcopy(bins)
+    bins = sorted(bins)
+    quantized = list(map(lambda y: bisect.bisect_right(bins, y), x))
+    return quantized
+
+
+def create_aspect_ratio_groups(dataset, k=0):
+    aspect_ratios = compute_aspect_ratios(dataset)
+    bins = (2 ** np.linspace(-1, 1, 2 * k + 1)).tolist() if k > 0 else [1.0]
+    groups = _quantize(aspect_ratios, bins)
+    # count number of elements per group
+    counts = np.unique(groups, return_counts=True)[1]
+    fbins = [0] + bins + [np.inf]
+    print(f"Using {fbins} as bins for aspect ratio quantization")
+    print(f"Count of instances per bin: {counts}")
+    return groups

src/detection/vision/mot_data.py

@@ -0,0 +1,370 @@
+import configparser
+import csv
+import os
+import os.path as osp
+import pickle
+
+import numpy as np
+import pycocotools.mask as rletools
+import torch
+from PIL import Image
+
+
+class MOTObjDetect(torch.utils.data.Dataset):
+    """ Data class for the Multiple Object Tracking Dataset
+    """
+
+    def __init__(self, root, transforms=None, vis_threshold=0.25,
+                 split_seqs=None, frame_range_start=0.0, frame_range_end=1.0):
+        self.root = root
+        self.transforms = transforms
+        self._vis_threshold = vis_threshold
+        self._classes = ('background', 'pedestrian')
+        self._img_paths = []
+        self._split_seqs = split_seqs
+
+        self.mots_gts = {}
+        for f in sorted(os.listdir(root)):
+            path = os.path.join(root, f)
+
+            if not os.path.isdir(path):
+                continue
+
+            if split_seqs is not None and f not in split_seqs:
+                continue
+
+            config_file = os.path.join(path, 'seqinfo.ini')
+
+            assert os.path.exists(config_file), \
+                'Path does not exist: {}'.format(config_file)
+
+            config = configparser.ConfigParser()
+            config.read(config_file)
+            seq_len = int(config['Sequence']['seqLength'])
+            im_ext = config['Sequence']['imExt']
+            im_dir = config['Sequence']['imDir']
+
+            img_dir = os.path.join(path, im_dir)
+
+            start_frame = int(frame_range_start * seq_len)
+            end_frame = int(frame_range_end * seq_len)
+
+            # for i in range(seq_len):
+            for i in range(start_frame, end_frame):
+                img_path = os.path.join(img_dir, f"{i + 1:06d}{im_ext}")
+                assert os.path.exists(
+                    img_path), f'Path does not exist: {img_path}'
+                self._img_paths.append(img_path)
+
+            # print(len(self._img_paths))
+
+            if self.has_masks:
+                gt_file = os.path.join(
+                    os.path.dirname(img_dir), 'gt', 'gt.txt')
+                self.mots_gts[gt_file] = load_mots_gt(gt_file)
+
+    def __str__(self):
+        if self._split_seqs is None:
+            return self.root
+        return f"{self.root}/{self._split_seqs}"
+
+    @property
+    def num_classes(self):
+        return len(self._classes)
+
+    def _get_annotation(self, idx):
+        """
+        """
+
+        if 'test' in self.root:
+            num_objs = 0
+            boxes = torch.zeros((num_objs, 4), dtype=torch.float32)
+
+            return {'boxes': boxes,
+                    'labels': torch.ones((num_objs,), dtype=torch.int64),
+                    'image_id': torch.tensor([idx]),
+                    'area': (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0]),
+                    'iscrowd': torch.zeros((num_objs,), dtype=torch.int64),
+                    'visibilities': torch.zeros((num_objs), dtype=torch.float32)}
+
+        img_path = self._img_paths[idx]
+        file_index = int(os.path.basename(img_path).split('.')[0])
+
+        gt_file = os.path.join(os.path.dirname(
+            os.path.dirname(img_path)), 'gt', 'gt.txt')
+
+        assert os.path.exists(gt_file), \
+            'GT file does not exist: {}'.format(gt_file)
+
+        bounding_boxes = []
+
+        if self.has_masks:
+            mask_objects_per_frame = self.mots_gts[gt_file][file_index]
+            masks = []
+            for mask_object in mask_objects_per_frame:
+                # class_id = 1 is car
+                # class_id = 2 is pedestrian
+                # class_id = 10 IGNORE
+                if mask_object.class_id in [1, 10] or not rletools.area(mask_object.mask):
+                    continue
+
+                bbox = rletools.toBbox(mask_object.mask)
+                x1, y1, w, h = [int(c) for c in bbox]
+
+                bb = {}
+                bb['bb_left'] = x1
+                bb['bb_top'] = y1
+                bb['bb_width'] = w
+                bb['bb_height'] = h
+
+                # print(bb, rletools.area(mask_object.mask))
+
+                bb['visibility'] = 1.0
+                bb['track_id'] = mask_object.track_id
+
+                masks.append(rletools.decode(mask_object.mask))
+                bounding_boxes.append(bb)
+        else:
+            with open(gt_file, "r") as inf:
+                reader = csv.reader(inf, delimiter=',')
+                for row in reader:
+                    visibility = float(row[8])
+
+                    if int(row[0]) == file_index and int(row[6]) == 1 and int(row[7]) == 1 and visibility and visibility >= self._vis_threshold:
+                        bb = {}
+                        bb['bb_left'] = int(row[2])
+                        bb['bb_top'] = int(row[3])
+                        bb['bb_width'] = int(row[4])
+                        bb['bb_height'] = int(row[5])
+                        bb['visibility'] = float(row[8])
+                        bb['track_id'] = int(row[1])
+
+                        bounding_boxes.append(bb)
+
+        num_objs = len(bounding_boxes)
+
+        boxes = torch.zeros((num_objs, 4), dtype=torch.float32)
+        visibilities = torch.zeros((num_objs), dtype=torch.float32)
+        track_ids = torch.zeros((num_objs), dtype=torch.long)
+
+        for i, bb in enumerate(bounding_boxes):
+            # Make pixel indexes 0-based, should already be 0-based (or not)
+            x1 = bb['bb_left']  # - 1
+            y1 = bb['bb_top']  # - 1
+            # This -1 accounts for the width (width of 1 x1=x2)
+            x2 = x1 + bb['bb_width']  # - 1
+            y2 = y1 + bb['bb_height']  # - 1
+
+            boxes[i, 0] = x1
+            boxes[i, 1] = y1
+            boxes[i, 2] = x2
+            boxes[i, 3] = y2
+            visibilities[i] = bb['visibility']
+            track_ids[i] = bb['track_id']
+
+        annos = {'boxes': boxes,
+                 'labels': torch.ones((num_objs,), dtype=torch.int64),
+                 'image_id': torch.tensor([idx]),
+                 'area': (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0]),
+                 'iscrowd': torch.zeros((num_objs,), dtype=torch.int64),
+                 'visibilities': visibilities,
+                 'track_ids': track_ids, }
+
+        if self.has_masks:
+            # annos['masks'] = torch.tensor(masks, dtype=torch.uint8)
+            annos['masks'] = torch.from_numpy(np.stack(masks))
+        return annos
+
+    @property
+    def has_masks(self):
+        return '/MOTS20/' in self.root
+
+    def __getitem__(self, idx):
+        # load images ad masks
+        img_path = self._img_paths[idx]
+        # mask_path = os.path.join(self.root, "PedMasks", self.masks[idx])
+        img = Image.open(img_path).convert("RGB")
+
+        target = self._get_annotation(idx)
+
+        if self.transforms is not None:
+            img, target = self.transforms(img, target)
+
+        return img, target
+
+    def __len__(self):
+        return len(self._img_paths)
+
+    def write_results_files(self, results, output_dir):
+        """Write the detections in the format for MOT17Det sumbission
+
+        all_boxes[image] = N x 5 array of detections in (x1, y1, x2, y2, score)
+
+        Each file contains these lines:
+        <frame>, <id>, <bb_left>, <bb_top>, <bb_width>, <bb_height>, <conf>, <x>, <y>, <z>
+
+        Files to sumbit:
+        ./MOT17-01.txt
+        ./MOT17-02.txt
+        ./MOT17-03.txt
+        ./MOT17-04.txt
+        ./MOT17-05.txt
+        ./MOT17-06.txt
+        ./MOT17-07.txt
+        ./MOT17-08.txt
+        ./MOT17-09.txt
+        ./MOT17-10.txt
+        ./MOT17-11.txt
+        ./MOT17-12.txt
+        ./MOT17-13.txt
+        ./MOT17-14.txt
+        """
+
+        #format_str = "{}, -1, {}, {}, {}, {}, {}, -1, -1, -1"
+
+        files = {}
+        for image_id, res in results.items():
+            path = self._img_paths[image_id]
+            img1, name = osp.split(path)
+            # get image number out of name
+            frame = int(name.split('.')[0])
+            # smth like /train/MOT17-09-FRCNN or /train/MOT17-09
+            tmp = osp.dirname(img1)
+            # get the folder name of the sequence and split it
+            tmp = osp.basename(tmp).split('-')
+            # Now get the output name of the file
+            out = tmp[0]+'-'+tmp[1]+'.txt'
+            outfile = osp.join(output_dir, out)
+
+            # check if out in keys and create empty list if not
+            if outfile not in files.keys():
+                files[outfile] = []
+
+            if 'masks' in res:
+                delimiter = ' '
+                # print(torch.unique(res['masks'][0]))
+                # > 0.5 #res['masks'].bool()
+                masks = res['masks'].squeeze(dim=1)
+
+                index_map = torch.arange(masks.size(0))[:, None, None]
+                index_map = index_map.expand_as(masks)
+
+                masks = torch.logical_and(
+                    # remove background
+                    masks > 0.5,
+                    # remove overlapp by largest probablity
+                    index_map == masks.argmax(dim=0)
+                )
+                for res_i in range(len(masks)):
+                    track_id = -1
+                    if 'track_ids' in res:
+                        track_id = res['track_ids'][res_i].item()
+                    mask = masks[res_i]
+                    mask = np.asfortranarray(mask)
+
+                    rle_mask = rletools.encode(mask)
+
+                    files[outfile].append(
+                        [frame,
+                         track_id,
+                         2,  # class pedestrian
+                         mask.shape[0],
+                         mask.shape[1],
+                         rle_mask['counts'].decode(encoding='UTF-8')])
+            else:
+                delimiter = ','
+                for res_i in range(len(res['boxes'])):
+                    track_id = -1
+                    if 'track_ids' in res:
+                        track_id = res['track_ids'][res_i].item()
+                    box = res['boxes'][res_i]
+                    score = res['scores'][res_i]
+
+                    x1 = box[0].item()
+                    y1 = box[1].item()
+                    x2 = box[2].item()
+                    y2 = box[3].item()
+
+                    out = [frame, track_id, x1, y1, x2 - x1,
+                           y2 - y1, score.item(), -1, -1, -1]
+
+                    if 'keypoints' in res:
+                        out.extend(res['keypoints'][res_i]
+                                   [:, :2].flatten().tolist())
+                        out.extend(res['keypoints_scores']
+                                   [res_i].flatten().tolist())
+
+                    files[outfile].append(out)
+
+        for k, v in files.items():
+            with open(k, "w") as of:
+                writer = csv.writer(of, delimiter=delimiter)
+                for d in v:
+                    writer.writerow(d)
+
+
+class SegmentedObject:
+    """
+    Helper class for segmentation objects.
+    """
+
+    def __init__(self, mask: dict, class_id: int, track_id: int, full_bbox=None) -> None:
+        self.mask = mask
+        self.class_id = class_id
+        self.track_id = track_id
+        self.full_bbox = full_bbox
+
+
+def load_mots_gt(path: str) -> dict:
+    """Load MOTS ground truth from path."""
+    objects_per_frame = {}
+    track_ids_per_frame = {}  # Check that no frame contains two objects with same id
+    combined_mask_per_frame = {}  # Check that no frame contains overlapping masks
+
+    with open(path, "r") as gt_file:
+        for line in gt_file:
+            line = line.strip()
+            fields = line.split(" ")
+
+            frame = int(fields[0])
+            if frame not in objects_per_frame:
+                objects_per_frame[frame] = []
+            # if frame not in track_ids_per_frame:
+            #     track_ids_per_frame[frame] = set()
+            # if int(fields[1]) in track_ids_per_frame[frame]:
+            #     assert False, f"Multiple objects with track id {fields[1]} in frame {fields[0]}"
+            # else:
+            #     track_ids_per_frame[frame].add(int(fields[1]))
+
+            class_id = int(fields[2])
+            if not (class_id == 1 or class_id == 2 or class_id == 10):
+                assert False, "Unknown object class " + fields[2]
+
+            mask = {
+                'size': [int(fields[3]), int(fields[4])],
+                'counts': fields[5].encode(encoding='UTF-8')}
+            if frame not in combined_mask_per_frame:
+                combined_mask_per_frame[frame] = mask
+            elif rletools.area(rletools.merge([
+                    combined_mask_per_frame[frame], mask],
+                    intersect=True)):
+                assert False, "Objects with overlapping masks in frame " + \
+                    fields[0]
+            else:
+                combined_mask_per_frame[frame] = rletools.merge(
+                    [combined_mask_per_frame[frame], mask],
+                    intersect=False)
+
+            full_bbox = None
+            if len(fields) == 10:
+                full_bbox = [int(fields[6]), int(fields[7]),
+                             int(fields[8]), int(fields[9])]
+
+            objects_per_frame[frame].append(SegmentedObject(
+                mask,
+                class_id,
+                int(fields[1]),
+                full_bbox
+            ))
+
+    return objects_per_frame

src/detection/vision/presets.py

@@ -0,0 +1,48 @@
+import torch
+from . import transforms as T
+
+
+class DetectionPresetTrain:
+    def __init__(self, data_augmentation, hflip_prob=0.5, mean=(123.0, 117.0, 104.0)):
+        if data_augmentation == "hflip":
+            self.transforms = T.Compose(
+                [
+                    T.RandomHorizontalFlip(p=hflip_prob),
+                    T.PILToTensor(),
+                    T.ConvertImageDtype(torch.float),
+                ]
+            )
+        elif data_augmentation == "ssd":
+            self.transforms = T.Compose(
+                [
+                    T.RandomPhotometricDistort(),
+                    T.RandomZoomOut(fill=list(mean)),
+                    T.RandomIoUCrop(),
+                    T.RandomHorizontalFlip(p=hflip_prob),
+                    T.PILToTensor(),
+                    T.ConvertImageDtype(torch.float),
+                ]
+            )
+        elif data_augmentation == "ssdlite":
+            self.transforms = T.Compose(
+                [
+                    T.RandomIoUCrop(),
+                    T.RandomHorizontalFlip(p=hflip_prob),
+                    T.PILToTensor(),
+                    T.ConvertImageDtype(torch.float),
+                ]
+            )
+        else:
+            raise ValueError(
+                f'Unknown data augmentation policy "{data_augmentation}"')
+
+    def __call__(self, img, target):
+        return self.transforms(img, target)
+
+
+class DetectionPresetEval:
+    def __init__(self):
+        self.transforms = T.ToTensor()
+
+    def __call__(self, img, target):
+        return self.transforms(img, target)

src/detection/vision/transforms.py

@@ -0,0 +1,284 @@
+from typing import List, Tuple, Dict, Optional
+
+import torch
+import torchvision
+from torch import nn, Tensor
+from torchvision.transforms import functional as F
+from torchvision.transforms import transforms as T
+
+
+def _flip_coco_person_keypoints(kps, width):
+    flip_inds = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]
+    flipped_data = kps[:, flip_inds]
+    flipped_data[..., 0] = width - flipped_data[..., 0]
+    # Maintain COCO convention that if visibility == 0, then x, y = 0
+    inds = flipped_data[..., 2] == 0
+    flipped_data[inds] = 0
+    return flipped_data
+
+
+class Compose:
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, image, target):
+        for t in self.transforms:
+            image, target = t(image, target)
+        return image, target
+
+
+class RandomHorizontalFlip(T.RandomHorizontalFlip):
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        if torch.rand(1) < self.p:
+            image = F.hflip(image)
+            if target is not None:
+                width, _ = F.get_image_size(image)
+                target["boxes"][:, [0, 2]] = width - target["boxes"][:, [2, 0]]
+                if "masks" in target:
+                    target["masks"] = target["masks"].flip(-1)
+                if "keypoints" in target:
+                    keypoints = target["keypoints"]
+                    keypoints = _flip_coco_person_keypoints(keypoints, width)
+                    target["keypoints"] = keypoints
+        return image, target
+
+
+class ToTensor(nn.Module):
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        image = F.pil_to_tensor(image)
+        image = F.convert_image_dtype(image)
+        return image, target
+
+
+class PILToTensor(nn.Module):
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        image = F.pil_to_tensor(image)
+        return image, target
+
+
+class ConvertImageDtype(nn.Module):
+    def __init__(self, dtype: torch.dtype) -> None:
+        super().__init__()
+        self.dtype = dtype
+
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        image = F.convert_image_dtype(image, self.dtype)
+        return image, target
+
+
+class RandomIoUCrop(nn.Module):
+    def __init__(
+        self,
+        min_scale: float = 0.3,
+        max_scale: float = 1.0,
+        min_aspect_ratio: float = 0.5,
+        max_aspect_ratio: float = 2.0,
+        sampler_options: Optional[List[float]] = None,
+        trials: int = 40,
+    ):
+        super().__init__()
+        # Configuration similar to https://github.com/weiliu89/caffe/blob/ssd/examples/ssd/ssd_coco.py#L89-L174
+        self.min_scale = min_scale
+        self.max_scale = max_scale
+        self.min_aspect_ratio = min_aspect_ratio
+        self.max_aspect_ratio = max_aspect_ratio
+        if sampler_options is None:
+            sampler_options = [0.0, 0.1, 0.3, 0.5, 0.7, 0.9, 1.0]
+        self.options = sampler_options
+        self.trials = trials
+
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        if target is None:
+            raise ValueError("The targets can't be None for this transform.")
+
+        if isinstance(image, torch.Tensor):
+            if image.ndimension() not in {2, 3}:
+                raise ValueError(f"image should be 2/3 dimensional. Got {image.ndimension()} dimensions.")
+            elif image.ndimension() == 2:
+                image = image.unsqueeze(0)
+
+        orig_w, orig_h = F.get_image_size(image)
+
+        while True:
+            # sample an option
+            idx = int(torch.randint(low=0, high=len(self.options), size=(1,)))
+            min_jaccard_overlap = self.options[idx]
+            if min_jaccard_overlap >= 1.0:  # a value larger than 1 encodes the leave as-is option
+                return image, target
+
+            for _ in range(self.trials):
+                # check the aspect ratio limitations
+                r = self.min_scale + (self.max_scale - self.min_scale) * torch.rand(2)
+                new_w = int(orig_w * r[0])
+                new_h = int(orig_h * r[1])
+                aspect_ratio = new_w / new_h
+                if not (self.min_aspect_ratio <= aspect_ratio <= self.max_aspect_ratio):
+                    continue
+
+                # check for 0 area crops
+                r = torch.rand(2)
+                left = int((orig_w - new_w) * r[0])
+                top = int((orig_h - new_h) * r[1])
+                right = left + new_w
+                bottom = top + new_h
+                if left == right or top == bottom:
+                    continue
+
+                # check for any valid boxes with centers within the crop area
+                cx = 0.5 * (target["boxes"][:, 0] + target["boxes"][:, 2])
+                cy = 0.5 * (target["boxes"][:, 1] + target["boxes"][:, 3])
+                is_within_crop_area = (left < cx) & (cx < right) & (top < cy) & (cy < bottom)
+                if not is_within_crop_area.any():
+                    continue
+
+                # check at least 1 box with jaccard limitations
+                boxes = target["boxes"][is_within_crop_area]
+                ious = torchvision.ops.boxes.box_iou(
+                    boxes, torch.tensor([[left, top, right, bottom]], dtype=boxes.dtype, device=boxes.device)
+                )
+                if ious.max() < min_jaccard_overlap:
+                    continue
+
+                # keep only valid boxes and perform cropping
+                target["boxes"] = boxes
+                target["labels"] = target["labels"][is_within_crop_area]
+                target["boxes"][:, 0::2] -= left
+                target["boxes"][:, 1::2] -= top
+                target["boxes"][:, 0::2].clamp_(min=0, max=new_w)
+                target["boxes"][:, 1::2].clamp_(min=0, max=new_h)
+                image = F.crop(image, top, left, new_h, new_w)
+
+                return image, target
+
+
+class RandomZoomOut(nn.Module):
+    def __init__(
+        self, fill: Optional[List[float]] = None, side_range: Tuple[float, float] = (1.0, 4.0), p: float = 0.5
+    ):
+        super().__init__()
+        if fill is None:
+            fill = [0.0, 0.0, 0.0]
+        self.fill = fill
+        self.side_range = side_range
+        if side_range[0] < 1.0 or side_range[0] > side_range[1]:
+            raise ValueError(f"Invalid canvas side range provided {side_range}.")
+        self.p = p
+
+    @torch.jit.unused
+    def _get_fill_value(self, is_pil):
+        # type: (bool) -> int
+        # We fake the type to make it work on JIT
+        return tuple(int(x) for x in self.fill) if is_pil else 0
+
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        if isinstance(image, torch.Tensor):
+            if image.ndimension() not in {2, 3}:
+                raise ValueError(f"image should be 2/3 dimensional. Got {image.ndimension()} dimensions.")
+            elif image.ndimension() == 2:
+                image = image.unsqueeze(0)
+
+        if torch.rand(1) >= self.p:
+            return image, target
+
+        orig_w, orig_h = F.get_image_size(image)
+
+        r = self.side_range[0] + torch.rand(1) * (self.side_range[1] - self.side_range[0])
+        canvas_width = int(orig_w * r)
+        canvas_height = int(orig_h * r)
+
+        r = torch.rand(2)
+        left = int((canvas_width - orig_w) * r[0])
+        top = int((canvas_height - orig_h) * r[1])
+        right = canvas_width - (left + orig_w)
+        bottom = canvas_height - (top + orig_h)
+
+        if torch.jit.is_scripting():
+            fill = 0
+        else:
+            fill = self._get_fill_value(F._is_pil_image(image))
+
+        image = F.pad(image, [left, top, right, bottom], fill=fill)
+        if isinstance(image, torch.Tensor):
+            # PyTorch's pad supports only integers on fill. So we need to overwrite the colour
+            v = torch.tensor(self.fill, device=image.device, dtype=image.dtype).view(-1, 1, 1)
+            image[..., :top, :] = image[..., :, :left] = image[..., (top + orig_h) :, :] = image[
+                ..., :, (left + orig_w) :
+            ] = v
+
+        if target is not None:
+            target["boxes"][:, 0::2] += left
+            target["boxes"][:, 1::2] += top
+
+        return image, target
+
+
+class RandomPhotometricDistort(nn.Module):
+    def __init__(
+        self,
+        contrast: Tuple[float] = (0.5, 1.5),
+        saturation: Tuple[float] = (0.5, 1.5),
+        hue: Tuple[float] = (-0.05, 0.05),
+        brightness: Tuple[float] = (0.875, 1.125),
+        p: float = 0.5,
+    ):
+        super().__init__()
+        self._brightness = T.ColorJitter(brightness=brightness)
+        self._contrast = T.ColorJitter(contrast=contrast)
+        self._hue = T.ColorJitter(hue=hue)
+        self._saturation = T.ColorJitter(saturation=saturation)
+        self.p = p
+
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        if isinstance(image, torch.Tensor):
+            if image.ndimension() not in {2, 3}:
+                raise ValueError(f"image should be 2/3 dimensional. Got {image.ndimension()} dimensions.")
+            elif image.ndimension() == 2:
+                image = image.unsqueeze(0)
+
+        r = torch.rand(7)
+
+        if r[0] < self.p:
+            image = self._brightness(image)
+
+        contrast_before = r[1] < 0.5
+        if contrast_before:
+            if r[2] < self.p:
+                image = self._contrast(image)
+
+        if r[3] < self.p:
+            image = self._saturation(image)
+
+        if r[4] < self.p:
+            image = self._hue(image)
+
+        if not contrast_before:
+            if r[5] < self.p:
+                image = self._contrast(image)
+
+        if r[6] < self.p:
+            channels = F.get_image_num_channels(image)
+            permutation = torch.randperm(channels)
+
+            is_pil = F._is_pil_image(image)
+            if is_pil:
+                image = F.pil_to_tensor(image)
+                image = F.convert_image_dtype(image)
+            image = image[..., permutation, :, :]
+            if is_pil:
+                image = F.to_pil_image(image)
+
+        return image, target

src/detection/vision/utils.py

@@ -0,0 +1,282 @@
+import datetime
+import errno
+import os
+import time
+from collections import defaultdict, deque
+
+import torch
+import torch.distributed as dist
+
+
+class SmoothedValue:
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median, avg=self.avg, global_avg=self.global_avg, max=self.max, value=self.value
+        )
+
+
+def all_gather(data):
+    """
+    Run all_gather on arbitrary picklable data (not necessarily tensors)
+    Args:
+        data: any picklable object
+    Returns:
+        list[data]: list of data gathered from each rank
+    """
+    world_size = get_world_size()
+    if world_size == 1:
+        return [data]
+    data_list = [None] * world_size
+    dist.all_gather_object(data_list, data)
+    return data_list
+
+
+def reduce_dict(input_dict, average=True):
+    """
+    Args:
+        input_dict (dict): all the values will be reduced
+        average (bool): whether to do average or sum
+    Reduce the values in the dictionary from all processes so that all processes
+    have the averaged results. Returns a dict with the same fields as
+    input_dict, after reduction.
+    """
+    world_size = get_world_size()
+    if world_size < 2:
+        return input_dict
+    with torch.inference_mode():
+        names = []
+        values = []
+        # sort the keys so that they are consistent across processes
+        for k in sorted(input_dict.keys()):
+            names.append(k)
+            values.append(input_dict[k])
+        values = torch.stack(values, dim=0)
+        dist.all_reduce(values)
+        if average:
+            values /= world_size
+        reduced_dict = {k: v for k, v in zip(names, values)}
+    return reduced_dict
+
+
+class MetricLogger:
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{attr}'")
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append(f"{name}: {str(meter)}")
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None):
+        i = 0
+        if not header:
+            header = ""
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt="{avg:.4f}")
+        data_time = SmoothedValue(fmt="{avg:.4f}")
+        space_fmt = ":" + str(len(str(len(iterable)))) + "d"
+        if torch.cuda.is_available():
+            log_msg = self.delimiter.join(
+                [
+                    header,
+                    "[{0" + space_fmt + "}/{1}]",
+                    "eta: {eta}",
+                    "{meters}",
+                    "time: {time}",
+                    "data: {data}",
+                    "max mem: {memory:.0f}",
+                ]
+            )
+        else:
+            log_msg = self.delimiter.join(
+                [header, "[{0" + space_fmt + "}/{1}]", "eta: {eta}", "{meters}", "time: {time}", "data: {data}"]
+            )
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == len(iterable) - 1:
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    print(
+                        log_msg.format(
+                            i,
+                            len(iterable),
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                            memory=torch.cuda.max_memory_allocated() / MB,
+                        )
+                    )
+                else:
+                    print(
+                        log_msg.format(
+                            i, len(iterable), eta=eta_string, meters=str(self), time=str(iter_time), data=str(data_time)
+                        )
+                    )
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print(f"{header} Total time: {total_time_str} ({total_time / len(iterable):.4f} s / it)")
+
+
+def collate_fn(batch):
+    return tuple(zip(*batch))
+
+
+def mkdir(path):
+    try:
+        os.makedirs(path, exist_ok=True)
+    except OSError as e:
+        if e.errno != errno.EEXIST:
+            raise
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    import builtins as __builtin__
+
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop("force", False)
+        if is_master or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+
+def init_distributed_mode(args):
+    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ["WORLD_SIZE"])
+        args.gpu = int(os.environ["LOCAL_RANK"])
+    elif "SLURM_PROCID" in os.environ:
+        args.rank = int(os.environ["SLURM_PROCID"])
+        args.gpu = args.rank % torch.cuda.device_count()
+    else:
+        print("Not using distributed mode")
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = "nccl"
+    print(f"| distributed init (rank {args.rank}): {args.dist_url}", flush=True)
+    torch.distributed.init_process_group(
+        backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank
+    )
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)

tools/anns/combine_anns.py

@@ -0,0 +1,87 @@
+import os
+import os.path as osp
+import json
+import argparse
+import tqdm
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--motsynth-path', help="Directory path containing the 'annotations' directory with .json files")
+    parser.add_argument(
+        '--save-path', help='Root file in which the new annoation files will be stored. If not provided, motsynth-root will be used')
+    parser.add_argument('--save-dir', default='comb_annotations',
+                        help="name of directory within 'save-path'in which MOTS annotation files will be stored")
+    parser.add_argument('--subsample', default=20, type=int,
+                        help="Frame subsampling rate. If e.g. 10 is selected, then we will select 1 in 10 frames")
+    parser.add_argument('--split', default='train',
+                        help="Name of split (i.e. set of sequences being merged) being used. A file named '{args.split}.txt needs to exist in the splits dir")
+    parser.add_argument(
+        '--name', help="Name of the split that file that will be generated. If not provided, the split name will be used")
+
+    args = parser.parse_args()
+
+    if args.save_path is None:
+        args.save_path = args.motsynth_path
+
+    if args.name is None:
+        args.name = args.split
+
+    assert args.subsample > 0, "Argument '--subsample' needs to be a positive integer. Set it to 1 to use every frame"
+
+    return args
+
+
+def read_split_file(path):
+    with open(path, 'r') as file:
+        seq_list = file.read().splitlines()
+
+    return seq_list
+
+
+def main(args):
+    # Determine which sequences to use
+    seqs = [seq.zfill(3) for seq in read_split_file(osp.join(
+        osp.dirname(os.path.abspath(__file__)), 'splits', f'{args.split}.txt'))]
+    comb_anns = {'images': [], 'annotations': [],
+                 'categories': None, 'info': {}}
+
+    for seq in tqdm.tqdm(seqs):
+        ann_path = osp.join(args.motsynth_path, 'annotations',  f'{seq}.json')
+        with open(ann_path) as f:
+            seq_ann = json.load(f)
+
+        # Subsample images and annotations if needed
+        if args.subsample > 1:
+            seq_ann['images'] = [{**img, **seq_ann['info']} for img in seq_ann['images'] if (
+                (img['frame_n'] - 1) % args.subsample) == 0]  # -1 bc in the paper this was 0-based
+            img_ids = [img['id'] for img in seq_ann['images']]
+            seq_ann['annotations'] = [
+                ann for ann in seq_ann['annotations'] if ann['image_id'] in img_ids]
+
+        comb_anns['images'].extend(seq_ann['images'])
+        comb_anns['annotations'].extend(seq_ann['annotations'])
+        comb_anns['info'][seq] = seq_ann['info']
+
+    if len(seqs) > 0:
+        comb_anns['categories'] = seq_ann['categories']
+        comb_anns['licenses'] = seq_ann['categories']
+
+    # Sanity check:
+    img_ids = [img['id'] for img in comb_anns['images']]
+    ann_ids = [ann['id'] for ann in comb_anns['annotations']]
+    assert len(img_ids) == len(set(img_ids))
+    assert len(ann_ids) == len(set(ann_ids))
+
+    # Save the new annotations file
+    comb_anns_dir = osp.join(args.save_path, args.save_dir)
+    os.makedirs(comb_anns_dir, exist_ok=True)
+    comb_anns_path = osp.join(comb_anns_dir, f"{args.name}.json")
+    with open(comb_anns_path, 'w') as json_file:
+        json.dump(comb_anns, json_file)
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    main(args)

tools/anns/generate_mot_format_files.py

@@ -0,0 +1,73 @@
+import pandas as pd
+import numpy as np
+
+import os.path as osp
+import os
+import json
+
+import tqdm
+import argparse
+
+from generate_mots_format_files import save_seqinfo
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--motsynth-path', help="Directory path containing the 'annotations' directory with .json files")
+    parser.add_argument('--save-path', help='Root file in which the new annoation files will be stored. If not provided, motsynth-root will be used')
+    parser.add_argument('--save-dir', default='mot_annotations', help="name of directory within 'save-path'in which MOTS annotation files will be stored")
+    args = parser.parse_args()
+
+    if args.save_path is None:
+        args.save_path = args.motsynth_path
+
+    return args
+
+def main(args):
+    ann_dir = osp.join(args.motsynth_path, 'annotations')
+    mot_ann_dir = osp.join(args.save_path, args.save_dir)
+    seqs = [f'{seq_num:03}' for seq_num in range(768) if seq_num not in (629, 757, 524, 652)]
+
+    for seq  in tqdm.tqdm(seqs):
+        ann_path = osp.join(ann_dir, f'{seq}.json')
+        with open(ann_path) as f:
+            seq_ann = json.load(f)
+
+        rows = []
+        img_id2frame = {im['id']: im['frame_n'] for im in seq_ann['images']}    
+
+        for ann in seq_ann['annotations']:
+            # We compute the 3D location as the mid point between both feet keypoints in 3D
+            kps = np.array(ann['keypoints_3d']).reshape(-1, 4)
+            feet_pos_3d = kps[[-1, -4], :3].mean(axis = 0).round(4)
+
+            row = {'frame': img_id2frame[ann['image_id']],# STARTS AT 0!!!
+                   'id': ann['ped_id'],
+                   'bb_left': ann['bbox'][0] + 1, # Make it 1-based??
+                   'bb_top': ann['bbox'][1] + 1,
+                   'bb_width': ann['bbox'][2],
+                   'bb_height': ann['bbox'][3],
+                   'conf': 1 - ann['iscrowd'],
+                   'class': 1 if ann['iscrowd'] == 0 else 8, # Class 8 means distractor. It is the one used by Trackeval as 'iscrowd' 
+                   # We compute visibility as the proportion of visible keypoints
+                   'vis': (np.array(ann['keypoints'])[2::3] ==2).mean().round(2),
+                   'x': feet_pos_3d[0],
+                   'y': feet_pos_3d[1],
+                   'z': feet_pos_3d[2]}
+
+            rows.append(row)
+            
+        # Save gt.txt file
+        # Format in https://github.com/dendorferpatrick/MOTChallengeEvalKit/tree/master/MOT
+        mot_ann = pd.DataFrame(rows, columns = ['frame', 'id', 'bb_left', 'bb_top', 'bb_width', 'bb_height', 'conf','class', 'vis', 'x', 'y', 'z'])
+        gt_dir = osp.join(mot_ann_dir, seq, 'gt')
+        os.makedirs(gt_dir, exist_ok=True)
+        mot_ann.to_csv(osp.join(gt_dir, 'gt.txt'), header=None, index=None, sep=',')
+        
+
+        # Save seqinfo.ini
+        seqinfo_path = osp.join(mot_ann_dir, seq, 'seqinfo.ini')
+        save_seqinfo(seqinfo_path, info = seq_ann['info'])
+
+if __name__ =='__main__':
+    args = parse_args()
+    main(args)

tools/anns/generate_mots_format_files.py

@@ -0,0 +1,102 @@
+"""
+TODOs: 
+- argparse
+- List sequences by number
+- Get rid of asserts
+"""
+
+import pandas as pd
+
+import os.path as osp
+import os
+import json
+
+import configparser
+
+import tqdm
+import argparse
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--motsynth-path', help="Directory path containing the 'annotations' directory with .json files")
+    parser.add_argument('--save-path', help='Root file in which the new annoation files will be stored. If not provided, motsynth-root will be used')
+    parser.add_argument('--save-dir', default='mots_annotations', help="name of directory within 'save-path'in which MOTS annotation files will be stored")
+    args = parser.parse_args()
+
+    if args.save_path is None:
+        args.save_path = args.motsynth_path
+
+    return args
+
+def save_seqinfo(seqinfo_path, info):
+    seqinfo = configparser.ConfigParser()
+    seqinfo.optionxform = str # Otherwise capital letters are ignored in keys
+
+    seqinfo['Sequence'] = dict(name=info['seq_name'],
+                            frameRate=info['fps'],
+                            seqLength=info['sequence_length'],
+                            imWidth= info['img_width'],
+                            imHeight= info['img_height'],
+                            weather=info['weather'],
+                            time=info['time'],
+                            isNight=info['is_night'],
+                            isMoving=info['is_moving'],
+                            FOV=info['cam_fov'],
+                            imExt='.jpg',
+                            fx=1158,
+                            fy=1158,
+                            cx=960,
+                            cy=540)
+        
+    with open(seqinfo_path, 'w') as configfile:    # save
+        seqinfo.write(configfile, space_around_delimiters=False)
+
+
+def main(args):
+    ann_dir = osp.join(args.motsynth_path, 'annotations')
+    mots_ann_dir = osp.join(args.save_path, args.save_dir)
+    
+    seqs = [f'{seq_num:03}' for seq_num in range(768) if seq_num not in (629, 757, 524, 652)]
+
+    for seq  in tqdm.tqdm(seqs):
+        ann_path = osp.join(ann_dir, f'{seq}.json')
+        with open(ann_path) as f:
+            seq_ann = json.load(f)
+
+        rows = []
+        img_id2frame = {im['id']: im['frame_n'] for im in seq_ann['images']}    
+
+        for ann in seq_ann['annotations']:
+            assert ann['category_id'] == 1
+            if ann['area']: # Include only objects with non-empty masks
+                if not ann['iscrowd']:
+                    mots_id = 2000 + ann['ped_id']
+                
+                else: # ID = 10000 means that the instance should be ignored during eval.
+                    mots_id = 10000
+
+                row = {'time_frame': img_id2frame[ann['image_id']],# STARTS AT 0!!!
+                    'id': mots_id,
+                    'class_id': 2,
+                    'img_height': ann['segmentation']['size'][0],
+                    'img_width': ann['segmentation']['size'][1],
+                    'rle': ann['segmentation']['counts']}
+
+                rows.append(row)
+                
+        # Save gt.txt file
+        # Format in https://www.vision.rwth-aachen.de/page/mots        
+        mots_ann = pd.DataFrame(rows, columns = ['time_frame', 'id', 'class_id', 'img_height', 'img_width', 'rle'])
+        gt_dir = osp.join(mots_ann_dir, seq, 'gt')
+        os.makedirs(gt_dir, exist_ok=True)
+        mots_ann.to_csv(osp.join(gt_dir, 'gt.txt'), header=None, index=None, sep=' ')
+        
+
+        # Save seqinfo.ini
+        seqinfo_path = osp.join(mots_ann_dir, seq, 'seqinfo.ini')
+        save_seqinfo(seqinfo_path, info = seq_ann['info'])
+
+
+if __name__ =='__main__':
+    args = parse_args()
+    main(args)

tools/anns/motcha_to_coco.py

@@ -0,0 +1,145 @@
+import os
+import os.path as osp
+import numpy as np
+import json
+
+import argparse
+import configparser
+import datetime
+
+import tqdm
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--data-root', help="Path containing the dataset in a folder")
+    parser.add_argument('--dataset', default='MOT17',
+                        help='Name of the dataset to be used. Should be either MOT17 or MOT20')
+    parser.add_argument(
+        '--save-dir', help='Root file in which the new annoation files will be stored. If not provided, data-root will be used')
+    parser.add_argument('--split', default='train',
+                        help="Split processed within the dataset. Should be either 'train' or 'test'")
+    parser.add_argument('--save-combined', default=True, action='store_true',
+                        help="Determines whether a separate .json file containing all sequence annotations will be created")
+    parser.add_argument('--subsample', default=20, type=int,
+                        help="Frame subsampling rate. If e.g. 10 is selected, then we will select 1 in 20 frames")
+
+    args = parser.parse_args()
+
+    if args.save_dir is None:
+        args.save_dir = osp.join(args.data_root, 'motcha_coco_annotations')
+
+    return args
+
+
+def get_img_id(dataset, seq, fname):
+    # Dataset num, seq num, frame num
+    return int(f"{dataset[3:5]}{seq.split('-')[1]}{int(fname.split('.')[0]):06}")
+
+
+def read_seqinfo(path):
+    cp = configparser.ConfigParser()
+    cp.read(path)
+    return {'height': int(cp.get('Sequence', 'imHeight')),
+            'width': int(cp.get('Sequence', 'imWidth')),
+            'fps': int(cp.get('Sequence', 'frameRate'))}
+
+
+def main(args):
+    data_path = osp.join(args.data_root, args.dataset, args.split)
+    seqs = os.listdir(data_path)
+
+    if args.save_combined:
+        comb_data = {'info': {'dataset': args.dataset,
+                              'split': args.split,
+                              'creation_date': datetime.datetime.today().strftime('%Y-%m-%d-%H-%M')},
+                     'images': [],
+                     'annotations': [],
+                     'categories': [{'id': 1, 'name': 'person', 'supercategory': 'person'}]}
+
+    for seq in tqdm.tqdm(seqs):
+        if args.dataset.lower() == 'mot17':
+            # Choose an arbitrary set of detections for MOT17, annotations are the same
+            if not seq.endswith('FRCNN'):
+                continue
+
+        print(f"Processing sequence {seq} in dataset {args.dataset}")
+
+        seq_path = osp.join(data_path, seq)
+        seqinfo_path = osp.join(seq_path, 'seqinfo.ini')
+        gt_path = osp.join(seq_path, 'gt/gt.txt')
+        im_dir = osp.join(seq_path, 'img1')
+
+        if args.dataset.lower() == 'mot17':
+            seq_ = '-'.join(seq.split('-')[:-1])  # Get rid of detector string
+
+        else:
+            seq_ = seq.copy()
+
+        seqinfo = read_seqinfo(seqinfo_path)
+        data = {'info': {'sequence': seq_,
+                         'dataset': args.dataset,
+                         'split': args.split,
+                         'creation_date': datetime.datetime.today().strftime('%Y-%m-%d-%H-%M'),
+                         **seqinfo},
+                'images': [],
+                'annotations': [],
+                'categories': [{'id': 1, 'name': 'person', 'supercategory': 'person'}]}
+
+        # Load Bounding Box annotations
+        gt = np.loadtxt(gt_path, dtype=np.float32, delimiter=',')
+        keep_classes = [1, 2, 7, 8, 12]
+        mask = np.isin(gt[:, 7], keep_classes)
+        gt = gt[mask]
+        anns = [{'ped_id': row[1],
+                'frame_n': row[0],
+                 'category_id': 1,
+                 'id': f"{get_img_id(args.dataset, seq, f'{int(row[0]):06}.jpg')}{int(row_i):010}",
+                 'image_id': get_img_id(args.dataset, seq, f'{int(row[0]):06}.jpg'),
+                 # 'bbox': row[2:6].tolist(),
+                 # MOTCha annotations are 1-based
+                 'bbox': [row[2] - 1, row[3] - 1, row[4], row[5]],
+                 'area': row[4]*row[5],
+                 'vis': row[8],
+                 'iscrowd': 1 - row[6]}
+                for row_i, row in enumerate(gt.astype(float)) if row[0] % args.subsample == 0]
+
+        # Load Image information
+        all_img_ids = list(set([aa['image_id'] for aa in anns]))
+        imgs = [{'file_name': osp.join(args.dataset, args.split, seq, 'img1', fname),
+                 'height': seqinfo['height'],
+                 'width': seqinfo['width'],
+                 'frame_n': int(fname.split('.')[0]),
+                 'id': get_img_id(args.dataset, seq, fname)}
+                for fname in os.listdir(im_dir) if get_img_id(args.dataset, seq, fname) in all_img_ids]
+        assert len(set([im['id'] for im in imgs])) == len(imgs)
+        data['images'].extend(imgs)
+        assert len(str(imgs[0]['id'])) == len(str(anns[0]['image_id']))
+
+        data['annotations'].extend(anns)
+
+        os.makedirs(args.save_dir, exist_ok=True)
+        fname = f"{args.dataset}_{seq_}.json" if args.dataset not in seq_ else f"{seq_}.json"
+        save_path = osp.join(args.save_dir, fname)
+        with open(save_path, 'w') as f:
+            json.dump(data, f)
+
+        print(f"Saved result at {save_path}")
+
+        if args.save_combined:
+            comb_data['annotations'].extend(anns)
+            comb_data['images'].extend(imgs)
+
+    if args.save_combined:
+        save_path = osp.join(
+            args.save_dir, f"{args.dataset}_{args.split}.json")
+        with open(save_path, 'w') as f:
+            json.dump(comb_data, f)
+
+        print(f"Saved combined result at {save_path}")
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    main(args)

tools/anns/splits/motsynth_split1.txt

@@ -0,0 +1,16 @@
+000
+008
+016
+024
+032
+040
+048
+056
+064
+072
+080
+088
+096
+104
+112
+120

tools/anns/splits/motsynth_split2.txt

@@ -0,0 +1,31 @@
+000
+004
+008
+012
+016
+020
+024
+028
+032
+036
+040
+044
+048
+052
+056
+060
+064
+068
+072
+076
+080
+084
+088
+092
+096
+100
+104
+108
+112
+116
+120

tools/anns/splits/motsynth_split3.txt

@@ -0,0 +1,62 @@
+000
+002
+004
+006
+008
+010
+012
+014
+016
+018
+020
+022
+024
+026
+028
+030
+032
+034
+036
+038
+040
+042
+044
+046
+048
+050
+052
+054
+056
+058
+060
+062
+064
+066
+068
+070
+072
+074
+076
+078
+080
+082
+084
+086
+088
+090
+092
+094
+096
+098
+100
+102
+104
+106
+108
+110
+112
+114
+116
+118
+120
+122

tools/anns/splits/motsynth_split4.txt

@@ -0,0 +1,123 @@
+000
+001
+002
+003
+004
+005
+006
+007
+008
+009
+010
+011
+012
+013
+014
+015
+016
+017
+018
+019
+020
+021
+022
+023
+024
+025
+026
+027
+028
+029
+030
+031
+032
+033
+034
+035
+036
+037
+038
+039
+040
+041
+042
+043
+044
+045
+046
+047
+048
+049
+050
+051
+052
+053
+054
+055
+056
+057
+058
+059
+060
+061
+062
+063
+064
+065
+066
+067
+068
+069
+070
+071
+072
+073
+074
+075
+076
+077
+078
+079
+080
+081
+082
+083
+084
+085
+086
+087
+088
+089
+090
+091
+092
+093
+094
+095
+096
+097
+098
+099
+100
+101
+102
+103
+104
+105
+106
+107
+108
+109
+110
+111
+112
+113
+114
+115
+116
+117
+118
+119
+120
+121
+122

tools/anns/store_reid_imgs.py

@@ -0,0 +1,84 @@
+import os
+import os.path as osp
+import numpy as np
+import json
+
+import argparse
+
+import tqdm
+from PIL import Image
+from collections import defaultdict
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--ann-path', help=".JSON annotations file path in COCO format")
+    parser.add_argument('--frames-path', help="Root directory containing images")
+    parser.add_argument('--save-dir', help='Root file in which the new annoation files will be stored. If not provided, data-root will be used')
+    parser.add_argument('--start-iter', default=0, type=int)
+    
+    #args = parser.parse_args(['--ann-path', '/storage/user/brasoand/MOTSynth/comb_annotations/train_mini.json'])
+    args = parser.parse_args()
+
+    if args.frames_path is None:
+        args.frames_path = osp.dirname(osp.dirname(args.ann_path))
+
+    if args.save_dir is None:
+        #args.save_dir = osp.join(osp.dirname(osp.dirname(args.ann_path)), 'reid_images')
+        args.save_dir = osp.join(osp.dirname(osp.dirname(args.ann_path)), 'reid')
+
+
+    return args
+
+def crop_box(im, bbox):
+    x1, y1, w, h = bbox
+    x2, y2 = x1+ w, y1+ h
+    return im.crop((x1, y1, x2, y2))
+
+
+def main(args):
+    os.makedirs(args.save_dir, exist_ok=True)
+
+    # Read annotations
+    with open(args.ann_path) as f:
+        anns = json.load(f)
+
+
+    # Annotation ids are used as file names to store boxes. 
+    # Therefore they need to be unique
+    ann_ids = [ann['id'] for ann in anns['annotations']]
+    assert len(ann_ids) == len(set(ann_ids))    
+    imgid2file = {img['id']: img['file_name'] for img in anns['images']}
+
+    # TODO: This needs to go!!!!!!!    
+    anns['annotations'] = [ann for ann in anns['annotations'] if not osp.exists(osp.join(args.save_dir, f"{ann['id']}.png"))]
+    len(anns['annotations'])
+    im2anns = defaultdict(list)
+    for ann in anns['annotations']:
+        im2anns[imgid2file[ann['image_id']]].append(ann)
+
+    for img_file, im_anns  in tqdm.tqdm(im2anns.items()):
+        #break
+        # Read Image
+        im_path = osp.join(args.frames_path, img_file)
+        if not osp.exists(im_path):
+            im_path = osp.join(args.frames_path, img_file.replace('rgb/', ''))
+
+        assert osp.exists(im_path)
+        im = Image.open(im_path)
+
+        for ann in im_anns:
+            box_path = osp.join(args.save_dir, f"{ann['id']}.png")
+            
+            if osp.exists(box_path):
+                continue
+
+            #if ann['bbox'][-2] > 2000 or ann['bbox'][-1] > 2000:
+            #    continue
+
+            box_im = crop_box(im, ann['bbox'])
+            box_im.save(box_path)
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    main(args)

tools/anns/to_frames.py

@@ -0,0 +1,56 @@
+import argparse
+import os
+import os.path as osp
+import cv2
+import glob
+
+import sys
+
+import tqdm
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Get frames from a video')
+    parser.add_argument(
+        '--motsynth-root', help='Directory hosting MOTSYnth part directories')
+    args = parser.parse_args()
+
+    video_paths = glob.glob(
+        osp.join(args.motsynth_root, 'MOTSynth_[0-9]/[0-9][0-9][0-9].mp4'))
+
+    frames_dir = os.path.join(args.motsynth_root, "frames")
+    os.makedirs(frames_dir, exist_ok=True)
+
+    print("Start extracting frames...")
+
+    for video_path in tqdm.tqdm(video_paths):
+        vidcap = cv2.VideoCapture(video_path)
+
+        seq_name = osp.basename(video_path).split(".")[0].zfill(3)
+        out_dir = os.path.join(frames_dir, seq_name, 'rgb')
+        os.makedirs(out_dir, exist_ok=True)
+
+        count = 1
+        success = True
+
+        #print("Unpacking video...")
+
+        while success:
+            success, image = vidcap.read()
+            if count < 3:
+                count += 1
+                continue
+            if not success or count == 1803:
+                break
+            if count % 200 == 0:
+                print("Extract frames until: " +
+                      str(count - 3).zfill(4) + ".jpg")
+            filename = os.path.join(out_dir, str(count - 3).zfill(4) + ".jpg")
+            cv2.imwrite(filename, image)     # save frame as JPEG file
+            count += 1
+
+    print("Done!")
+
+
+if __name__ == '__main__':
+    main()

tools/inference_detector.py

@@ -0,0 +1,46 @@
+import torch
+from torchvision.models.detection.faster_rcnn import fasterrcnn_resnet50_fpn_v2, FastRCNNPredictor
+from configs.path_cfg import OUTPUT_DIR
+from src.detection.vision.engine import evaluate
+from tools.train_detector import create_dataset, create_data_loader, get_transform
+from src.detection.graph_utils import add_bbox, show_img
+import os.path as osp
+import argparse
+
+
+def parse_args(add_help=True):
+    parser = argparse.ArgumentParser(
+        description="Detector inference", add_help=add_help)
+
+    # path to model used for inference
+    parser.add_argument("--model-path", type=str,
+                        help="Path with model checkpoint used for inference")
+
+    args = parser.parse_args()
+
+    if args.model_path is None:
+        args.model_path = osp.join(
+            OUTPUT_DIR, "detection_logs", "fasterrcnn_training", "checkpoint.pth")
+    return args
+
+
+def main(args):
+    ds_val = create_dataset(
+        "motsynth_val", get_transform(False, "hflip"), "test")
+    data_loader_val = create_data_loader(ds_val, "test", 1, 0)
+
+    device = torch.device("cuda")
+    model = fasterrcnn_resnet50_fpn_v2()
+    in_features = model.roi_heads.box_predictor.cls_score.in_features
+    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 2)
+    checkpoint = torch.load(
+        args.model_path, map_location="cpu")
+    model.load_state_dict(checkpoint["model"])
+    model.eval()
+    model.to(device)
+    show_img(data_loader_val, model, device, 0.8)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

tools/train_detector.py

@@ -0,0 +1,408 @@
+from typing import List
+from configs.path_cfg import MOTSYNTH_ROOT, MOTCHA_ROOT, OUTPUT_DIR
+import datetime
+import os.path as osp
+import os
+import time
+import coloredlogs
+import logging
+from torchinfo import summary
+import torch
+import torch.utils.data
+from src.detection.vision.mot_data import MOTObjDetect
+from src.detection.model_factory import ModelFactory
+from src.detection.graph_utils import save_train_loss_plot
+import src.detection.vision.presets as presets
+import src.detection.vision.utils as utils
+from src.detection.vision.engine import train_one_epoch, evaluate
+from src.detection.vision.group_by_aspect_ratio import GroupedBatchSampler, create_aspect_ratio_groups
+from src.detection.mot_dataset import get_mot_dataset
+import torchvision
+
+from torchvision.models.detection.faster_rcnn import fasterrcnn_resnet50_fpn_v2
+from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
+coloredlogs.install(level='DEBUG')
+logger = logging.getLogger(__name__)
+
+
+def get_args_parser(add_help=True):
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="PyTorch Detection Training", add_help=add_help)
+
+    # Output directory used to save model, plots and summary
+    parser.add_argument("--output-dir", default='fasterrcnn_training',
+                        type=str, help="Path to save outputs (default: fasterrcnn_training)")
+
+    # Dataset params
+    parser.add_argument("--train-dataset", default="motsynth_split1",
+                        type=str, help="Dataset name. Please select one of the following:  motsynth_split1, motsynth_split2, motsynth_split3, motsynth_split4, MOT17 (default: motsynth_split1)")
+    parser.add_argument("--val-dataset", default="MOT17",
+                        type=str, help="Dataset name. Please select one of the following: MOT17 (default: MOT17)")
+
+    # Transforms params
+    parser.add_argument(
+        "--data-augmentation", default="hflip", type=str, help="Data augmentation policy (default: hflip)"
+    )
+
+    # Data Loaders params
+    parser.add_argument(
+        "-b", "--batch-size", default=3, type=int, help="Images per gpu (default: 3)"
+    )
+    parser.add_argument(
+        "-j", "--workers", default=0, type=int, metavar="N", help="Number of data loading workers (default: 0)"
+    )
+    parser.add_argument("--aspect-ratio-group-factor", default=3,
+                        type=int, help="Aspect ration group factor (default:3)")
+
+    # Model param
+    parser.add_argument(
+        "--model", default="fasterrcnn_resnet50_fpn", type=str, help="Model name (default: fasterrcnn_resnet50_fpn)")
+    parser.add_argument(
+        "--weights", default="DEFAULT", type=str, help="Model weights (default: DEFAULT)"
+    )
+    parser.add_argument(
+        "--backbone", default='resnet50', type=str, help="Type of backbone (default: resnet50)"
+    )
+    parser.add_argument(
+        "--trainable-backbone-layers", default=3, type=int, help="Number of trainable layers of backbone (default: 3)"
+    )
+    parser.add_argument(
+        "--backbone-weights", default="DEFAULT", type=str, help="Backbone weights (default: DEFAULT)"
+    )
+
+    # Device param
+    parser.add_argument("--device", default="cuda", type=str,
+                        help="device (default: cuda)")
+
+    # Test mode param
+    parser.add_argument(
+        "--test-only",
+        dest="test_only",
+        help="Only test the model",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--model-eval", type=str, help="model path for test only mode"
+    )
+
+    # Optimizer params
+    parser.add_argument(
+        "--lr",
+        default=0.0025,
+        type=float,
+        help="Learning rate (default: 0.0025)",
+    )
+    parser.add_argument("--momentum", default=0.9,
+                        type=float, metavar="M", help="Momentum (default: 0.9")
+    parser.add_argument(
+        "--wd",
+        "--weight-decay",
+        default=1e-4,
+        type=float,
+        metavar="W",
+        help="Weight decay (default: 1e-4)",
+        dest="weight_decay",
+    )
+
+    # Lr Scheduler params
+    parser.add_argument(
+        "--lr-scheduler", default="multisteplr", type=str, help="Name of lr scheduler (default: multisteplr)"
+    )
+    parser.add_argument(
+        "--lr-steps",
+        default=[16, 22],
+        nargs="+",
+        type=int,
+        help="Decrease lr every step-size epochs (multisteplr scheduler only)",
+    )
+    parser.add_argument(
+        "--lr-gamma", default=0.1, type=float, help="Decrease lr by a factor of lr-gamma (multisteplr scheduler only)"
+    )
+
+    # Mixed precision training params
+    parser.add_argument("--amp", action="store_true",
+                        help="Use torch.cuda.amp for mixed precision training")
+
+    # Resume training params
+    parser.add_argument("--resume", default="", type=str,
+                        help="path of checkpoint")
+
+    # training param
+    parser.add_argument("--start_epoch", default=0,
+                        type=int, help="start epoch")
+    parser.add_argument("--epochs", default=30, type=int,
+                        metavar="N", help="number of total epochs to run")
+    parser.add_argument("--print-freq", default=20,
+                        type=int, help="print frequency")
+
+    return parser
+
+
+def get_transform(train, data_augmentation):
+    if train:
+        return presets.DetectionPresetTrain(data_augmentation)
+    else:
+        return presets.DetectionPresetEval()
+
+
+def get_motsynth_dataset(ds_name: str, transforms):
+    data_path = osp.join(MOTSYNTH_ROOT, 'comb_annotations', f"{ds_name}.json")
+    dataset = get_mot_dataset(MOTSYNTH_ROOT, data_path, transforms=transforms)
+    return dataset
+
+
+def get_MOT17_dataset(split: str, split_seqs: List, transforms):
+    data_path = osp.join(MOTCHA_ROOT, "MOT17", "train")
+    dataset = MOTObjDetect(
+        data_path, transforms=transforms, split_seqs=split_seqs)
+    return dataset
+
+
+def create_dataset(ds_name: str, transforms, split=None):
+    if (ds_name.startswith("motsynth")):
+        return get_motsynth_dataset(ds_name, transforms)
+
+    elif (ds_name.startswith("MOT17")):
+        if split == "train":
+            split_seqs = ['MOT17-02-FRCNN', 'MOT17-04-FRCNN',
+                          'MOT17-11-FRCNN', 'MOT17-13-FRCNN']
+        elif split == "test":
+            split_seqs = ['MOT17-09-FRCNN', 'MOT17-10-FRCNN', 'MOT17-05-FRCNN']
+        return get_MOT17_dataset(split, split_seqs, transforms)
+
+    else:
+        logger.error(
+            "Please, provide a valid dataset as argument. Select one of the following:  motsynth_split1, motsynth_split2, motsynth_split3, motsynth_split4, MOT17.")
+        raise ValueError(ds_name)
+
+
+def create_data_loader(dataset, split: str, batch_size, workers, aspect_ratio_group_factor=-1):
+    data_loader = None
+    if split == "train":
+        # random sampling on training dataset
+        train_sampler = torch.utils.data.RandomSampler(dataset)
+        if aspect_ratio_group_factor >= 0:
+            group_ids = create_aspect_ratio_groups(
+                dataset, k=aspect_ratio_group_factor)
+            train_batch_sampler = GroupedBatchSampler(
+                train_sampler, group_ids, batch_size)
+        else:
+            train_batch_sampler = torch.utils.data.BatchSampler(
+                train_sampler, batch_size, drop_last=True)
+        data_loader = torch.utils.data.DataLoader(
+            dataset, batch_sampler=train_batch_sampler, num_workers=workers, collate_fn=utils.collate_fn
+        )
+    elif split == "test":
+        # sequential sampling on eval dataset
+        test_sampler = torch.utils.data.SequentialSampler(dataset)
+        data_loader = torch.utils.data.DataLoader(
+            dataset, batch_size=1, sampler=test_sampler, num_workers=workers, collate_fn=utils.collate_fn
+        )
+    return data_loader
+
+
+def create_optimizer(model, lr, momentum, weight_decay):
+    params = [p for p in model.parameters() if p.requires_grad]
+    optimizer = torch.optim.SGD(
+        params, lr=lr, momentum=momentum, weight_decay=weight_decay)
+    return optimizer
+
+
+def create_lr_scheduler(optimizer, lr_scheduler_type, lr_steps, lr_gamma, epochs):
+    if lr_scheduler_type == "multisteplr":
+        lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer, milestones=lr_steps, gamma=lr_gamma)
+        logger.debug(
+            f"lr_scheduler: {lr_scheduler_type}, milestones: {lr_steps}, gamma: {lr_gamma}")
+
+    elif lr_scheduler_type == "cosineannealinglr":
+        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer, T_max=epochs)
+        logger.debug(
+            f"lr_scheduler: {lr_scheduler_type}, T_max: {epochs}")
+    else:
+        raise RuntimeError(
+            f"Invalid lr scheduler '{lr_scheduler_type}'. Only MultiStepLR and CosineAnnealingLR are supported."
+        )
+    return lr_scheduler
+
+
+def resume_training(model, optimizer, lr_scheduler, scaler, args):
+    checkpoint = torch.load(args.resume, map_location="cpu")
+
+    model.load_state_dict(checkpoint["model"])
+    optimizer.load_state_dict(checkpoint["optimizer"])
+    lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
+    args.start_epoch = checkpoint["epoch"] + 1
+    if args.amp:
+        scaler.load_state_dict(checkpoint["scaler"])
+
+
+def save_model_checkpoint(model, optimizer, lr_scheduler, epoch, scaler, output_dir, args):
+    if output_dir:
+        checkpoint = {
+            "model": model.state_dict(),
+            "optimizer": optimizer.state_dict(),
+            "lr_scheduler": lr_scheduler.state_dict(),
+            "args": args,
+            "epoch": epoch,
+        }
+        if args.amp:
+            checkpoint["scaler"] = scaler.state_dict()
+        utils.save_on_master(checkpoint, os.path.join(
+            output_dir, f"model_{epoch}.pth"))
+        utils.save_on_master(checkpoint, os.path.join(
+            output_dir, "checkpoint.pth"))
+
+
+def save_plots(losses_dict, batch_loss_dict, output_dir):
+    if not losses_dict:
+        for name, metric in batch_loss_dict.items():
+            losses_dict[name] = []
+
+    for name, metric in batch_loss_dict.items():
+        losses_dict[name].extend(metric)
+    save_train_loss_plot(losses_dict, output_dir)
+
+
+def save_model_summary(model, output_dir, batch_size):
+    with open(osp.join(output_dir, "summary.txt"), 'w', encoding="utf-8") as f:
+        print(summary(model,
+                      # (batch_size, color_channels, height, width)
+                      input_size=(batch_size, 3, 1080, 1920),
+                      verbose=0,
+                      col_names=["input_size", "output_size",
+                                 "num_params", "kernel_size", "trainable"],
+                      col_width=20,
+                      row_settings=["var_names"]), file=f)
+
+
+def save_args(output_dir, args):
+    with open(osp.join(output_dir, "args.txt"), 'w', encoding="utf-8") as f:
+        print(args, file=f)
+
+
+def save_evaluate_summary(stats, output_dir):
+    metrics = ["AP", "AP50", "AP75", "APs", "APm", "APl"]
+    # the standard metrics
+    results = {
+        metric: float(stats[idx] *
+                      100 if stats[idx] >= 0 else "nan")
+        for idx, metric in enumerate(metrics)
+    }
+    with open(osp.join(output_dir, "evaluate.txt"), 'w', encoding="utf-8") as f:
+        print(results, file=f)
+
+
+def main(args):
+
+    output_dir = None
+    if args.output_dir:
+        output_dir = osp.join(
+            OUTPUT_DIR, 'detection_logs', args.output_dir)
+        utils.mkdir(output_dir)
+    output_plots_dir = osp.join(output_dir, "plots")
+    utils.mkdir(output_plots_dir)
+
+    logger.debug("COMMAND LINE ARGUMENTS")
+    logger.debug(args)
+    save_args(output_dir, args)
+
+    device = torch.device(args.device)
+    logger.debug(f"DEVICE: {device}")
+
+    logger.debug("CREATE DATASETS")
+    ds_train_name = args.train_dataset
+    ds_val_name = args.val_dataset
+    data_augmentation = args.data_augmentation
+
+    dataset_train = create_dataset(
+        ds_train_name, get_transform(True, data_augmentation), "train")
+    dataset_test = create_dataset(
+        ds_val_name, get_transform(False, data_augmentation), "test")
+
+    logger.debug("CREATE DATA LOADERS")
+    batch_size = args.batch_size
+    workers = args.workers
+    aspect_ratio_group_factor = args.aspect_ratio_group_factor
+    data_loader_train = create_data_loader(
+        dataset_train, "train", batch_size, workers, aspect_ratio_group_factor)
+    data_loader_test = create_data_loader(
+        dataset_test, "test", batch_size, workers)
+
+    if args.test_only:
+        logger.debug("TEST ONLY")
+        model = fasterrcnn_resnet50_fpn_v2()
+        in_features = model.roi_heads.box_predictor.cls_score.in_features
+        model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 2)
+        checkpoint = torch.load(args.model_eval, map_location="cuda")
+        model.load_state_dict(checkpoint["model"])
+        model.to(device)
+        coco_evaluator = evaluate(model, data_loader_test,
+                                  device=device, iou_types=['bbox'])
+        save_evaluate_summary(
+            coco_evaluator.coco_eval['bbox'].stats, output_dir)
+        return
+
+    logger.debug("CREATE MODEL")
+    model_name = args.model
+    weights = args.weights
+    backbone = args.backbone
+    backbone_weights = args.backbone_weights
+    trainable_backbone_layers = args.trainable_backbone_layers
+    model = ModelFactory.get_model(
+        model_name, weights, backbone, backbone_weights, trainable_backbone_layers)
+    save_model_summary(model, output_dir, batch_size)
+
+    logger.debug("CREATE OPTIMIZER")
+    lr = args.lr
+    momentum = args.momentum
+    weight_decay = args.weight_decay
+    optimizer = create_optimizer(
+        model, lr, momentum, weight_decay)
+
+    logger.debug("CREATE LR SCHEDULER")
+    epochs = args.epochs
+    lr_scheduler_type = args.lr_scheduler.lower()
+    lr_steps = args.lr_steps
+    lr_gamma = args.lr_gamma
+    lr_scheduler = create_lr_scheduler(
+        optimizer, lr_scheduler_type, lr_steps, lr_gamma, epochs)
+
+    logger.debug("CONFIGURE SCALER FOR amp")
+    scaler = torch.cuda.amp.GradScaler() if args.amp else None
+
+    if args.resume:
+        logger.debug("RESUME TRAINING")
+        resume_training(model, optimizer, lr_scheduler,
+                        scaler, args)
+
+    logger.debug("START TRAINING")
+    print_freq = args.print_freq
+    start_epoch = args.start_epoch
+    losses_dict = {}
+    start_time = time.time()
+    for epoch in range(start_epoch, epochs):
+        _, batch_loss_dict = train_one_epoch(model, optimizer, data_loader_train, device,
+                                             epoch, print_freq, scaler)
+        lr_scheduler.step()
+        save_plots(losses_dict, batch_loss_dict,
+                   output_dir=output_plots_dir)
+
+        coco_evaluator = evaluate(model, data_loader_test,
+                                  device=device, iou_types=['bbox'])
+        save_evaluate_summary(
+            coco_evaluator.coco_eval['bbox'].stats, output_dir)
+
+    save_model_checkpoint(
+        model, optimizer, lr_scheduler, epoch, scaler, output_dir, args)
+    total_time = time.time() - start_time
+    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    logger.debug(f"TRAINING TIME: {total_time_str}")
+
+
+if __name__ == "__main__":
+    args = get_args_parser().parse_args()
+    main(args)