[feat] added support for split inference with yolox-darknet53

Author: chyomin06

README.MD

@@ -20,6 +20,8 @@ It currently focuses on two types of pipeline:
 
 - [JDE](https://github.com/Zhongdao/Towards-Realtime-MOT) is used for Object Tracking
 
+- [YOLOX-Darknet53](https://github.com/Megvii-BaseDetection/YOLOX) is used for object detection
+
 ## Documentation
 
 A complete documentation is provided [here](https://interdigitalinc.github.io/CompressAI-Vision/index.html), including [installation](https://interdigitalinc.github.io/CompressAI-Vision/installation), [CLI usage](https://interdigitalinc.github.io/CompressAI-Vision/cli_usage.html), as well as [tutorials](https://interdigitalinc.github.io/CompressAI-Vision/tutorials).
@@ -136,3 +138,4 @@ Fabien Racapé, Hyomin Choi, Eimran Eimon, Sampsa Riikonen, Jacky Yat-Hong Lam
  * [VVC VTM reference software](https://vcgit.hhi.fraunhofer.de/jvet/VVCSoftware_VTM)
  * [Detectron2](https://detectron2.readthedocs.io/en/latest/index.html)
  * [JDE](https://github.com/Zhongdao/Towards-Realtime-MOT.git)
+ * [YOLOX](https://github.com/Megvii-BaseDetection/YOLOX)

cfgs/vision_model/default.yaml

@@ -37,4 +37,13 @@ jde_1088x608:
   track_buffer: 30
   frame_rate: 30 # It is odd to consider this at here but following original code.
   splits : [36, 61, 74] # MPEG FCM TEST with JDE on TVD
-  #splits : [105, 90, 75] # MPEG FCM TEST with JDE on HiEve
+  #splits : [105, 90, 75] # MPEG FCM TEST with JDE on HiEve
+
+yolox_darknet53:
+  model_path_prefix: ${..model_root_path}
+  cfg: "Built-in configurations"
+  num_classes: 80
+  conf_thres: 0.001
+  nms_thres: 0.65
+  weights: "weights/yolox/darknet53/yolox_darknet.pth"
+  splits: "l13"

compressai_vision/datasets/image.py

@@ -47,7 +47,7 @@
 
 from compressai_vision.registry import register_datacatalog, register_dataset
 
-from .utils import JDECustomMapper, LinearMapper
+from .utils import JDECustomMapper, LinearMapper, YOLOXCustomMapper
 
 
 def manual_load_data(path, ext):
@@ -295,6 +295,48 @@ def __len__(self):
         return len(self.mapDataset)
 
 
+@register_dataset("YOLOXDataset")
+class YOLOXDataset(BaseDataset):
+    def __init__(self, root, dataset_name, imgs_folder, **kwargs):
+        super().__init__(root, dataset_name, imgs_folder, **kwargs)
+
+        self.dataset = kwargs["dataset"].dataset
+
+        self.sampler = InferenceSampler(len(kwargs["dataset"]))
+        self.collate_fn = bypass_collator
+
+        _dataset = DatasetFromList(self.dataset, copy=False)
+
+        if kwargs["linear_mapper"] is True:
+            mapper = LinearMapper()
+        else:
+            mapper = YOLOXCustomMapper(kwargs["patch_size"])
+
+        self.input_size = kwargs["patch_size"]
+        self.mapDataset = MapDataset(_dataset, mapper)
+        self._org_mapper_func = PicklableWrapper(
+            YOLOXCustomMapper(kwargs["patch_size"])
+        )
+
+        metaData = MetadataCatalog.get(dataset_name)
+        try:
+            self.thing_classes = metaData.thing_classes
+            self.thing_dataset_id_to_contiguous_id = (
+                metaData.thing_dataset_id_to_contiguous_id
+            )
+        except AttributeError:
+            self.logger.warning("No attribute: thing_classes")
+
+    def get_org_mapper_func(self):
+        return self._org_mapper_func
+
+    def __getitem__(self, idx):
+        return self.mapDataset[idx]
+
+    def __len__(self):
+        return len(self.mapDataset)
+
+
 class DataCatalog:
     def __init__(
         self,

compressai_vision/datasets/utils.py

@@ -35,8 +35,105 @@
 import numpy as np
 import torch
 from jde.utils.datasets import letterbox
+from torchvision import transforms
 
-__all__ = ["JDECustomMapper", "LinearMapper"]
+__all__ = ["YOLOXCustomMapper", "JDECustomMapper", "LinearMapper"]
+
+
+def yolox_style_scaling(img, input_size, padding=False):
+    r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
+
+    resized_img = cv2.resize(
+        img,
+        (int(img.shape[1] * r), int(img.shape[0] * r)),
+        interpolation=cv2.INTER_LINEAR,
+    ).astype(np.uint8)
+
+    if padding:
+        padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114
+        padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img
+
+        return padded_img
+
+    return resized_img
+
+
+class YOLOXCustomMapper:
+    """
+    A callable which takes a dataset dict in CompressAI-Vision generic dataset format, but for YOLOX evaluation,
+    and map it into a format used by the model.
+
+    This is the default callable to be used to map your dataset dict into inference data.
+
+    This callable function refers to
+        preproc function at
+        <https://github.com/Megvii-BaseDetection/YOLOX/yolox/data/data_augment.py>
+
+        Full license statement can be found at
+        <https://github.com/Megvii-BaseDetection/YOLOX?tab=Apache-2.0-1-ov-file#readme>
+
+    """
+
+    def __init__(self, img_size=[640, 640], aug_transforms=None):
+        """
+        Args:
+            img_size: expected input size (Height, Width)
+        """
+
+        self.input_img_size = img_size
+
+        if aug_transforms != None:
+            self.aug_transforms = aug_transforms
+        else:
+            self.aug_transforms = transforms.Compose([transforms.ToTensor()])
+
+    def __call__(self, dataset_dict):
+        """
+        Args:
+            dataset_dict (dict): Metadata of one image.
+
+        Returns:
+            dict: a format that compressai-vision pipelines accept
+        """
+
+        dataset_dict = copy.deepcopy(dataset_dict)
+        # the copied dictionary will be modified by code below
+
+        dataset_dict.pop("annotations", None)
+
+        # replicate the implemetation of the original codes
+        # Read image
+        org_img = cv2.imread(dataset_dict["file_name"])  # return img in BGR by default
+
+        assert (
+            len(org_img.shape) == 3
+        ), f"detect an input image with 2 chs, {dataset_dict['file_name']}"
+
+        dataset_dict["height"], dataset_dict["width"], _ = org_img.shape
+
+        # yolox style input scaling
+        # 1st scaling
+        resized_img = yolox_style_scaling(org_img, self.input_img_size)
+        # 2nd scaling & padding
+        resized_img = yolox_style_scaling(
+            resized_img, self.input_img_size, padding=True
+        )
+
+        tensor_image = self.aug_transforms(
+            np.ascontiguousarray(resized_img, dtype=np.float32)
+        )
+
+        # old way
+        # kept BGR & swap axis
+        # image = resized_img.transpose(2, 0, 1)
+        # normalize contiguous array of image
+        # image = np.ascontiguousarray(image, dtype=np.float32)
+        # to tensor
+        # tensor_image = torch.as_tensor(image)
+
+        dataset_dict["image"] = tensor_image
+
+        return dataset_dict
 
 
 class JDECustomMapper:

compressai_vision/evaluators/__init__.py

@@ -35,14 +35,14 @@
     MOT_TVD_Eval,
     OpenImagesChallengeEval,
     VisualQualityEval,
-    YOLOEval,
+    YOLOXCOCOEval,
 )
 
 __all__ = [
     "BaseEvaluator",
     "COCOEVal",
     "OpenImagesChallengeEval",
-    "YOLOEval",
+    "YOLOXCOCOEval",
     "MOT_JDE_Eval",
     "MOT_HiEve_Eval",
     "MOT_TVD_Eval",

compressai_vision/evaluators/evaluators.py

@@ -29,6 +29,7 @@
 
 import json
 import math
+from collections import defaultdict
 from pathlib import Path
 
 import motmetrics as mm
@@ -37,8 +38,11 @@
 import torch
 from detectron2.evaluation import COCOEvaluator
 from jde.utils.io import unzip_objs
+from pycocotools.coco import COCO
 from pytorch_msssim import ms_ssim
 from tqdm import tqdm
+from yolox.data.datasets.coco import remove_useless_info
+from yolox.evaluators import COCOEvaluator as YOLOX_COCOEvaluator
 
 from compressai_vision.datasets import deccode_compressed_rle
 from compressai_vision.registry import register_evaluator
@@ -627,20 +631,76 @@ def mot_eval(self):
         return self.digest_summary(summary)
 
 
-@register_evaluator("YOLO-EVAL")
-class YOLOEval(BaseEvaluator):
+@register_evaluator("YOLOX-COCO-EVAL")
+class YOLOXCOCOEval(BaseEvaluator):
     def __init__(
         self,
         datacatalog_name,
         dataset_name,
         dataset,
         output_dir="./vision_output/",
-        criteria="AP50",
+        criteria="AP",
     ):
         super().__init__(datacatalog_name, dataset_name, dataset, output_dir, criteria)
 
         self.set_annotation_info(dataset)
 
+        cocoapi = COCO(self.annotation_path)
+        remove_useless_info(cocoapi)
+        class_ids = sorted(cocoapi.getCatIds())
+        cats = cocoapi.loadCats(cocoapi.getCatIds())
+
+        class dummy_dataloader:
+            def __init__(self):
+                class dummy_dataset:
+                    def __init__(self):
+                        self.coco = cocoapi
+                        self.class_ids = class_ids
+                        self.cats = cats
+
+                self.dataset = dummy_dataset()
+                self.batch_size = 1
+
+        dataloader = dummy_dataloader()
+        self._evaluator = YOLOX_COCOEvaluator(
+            dataloader, dataset.input_size, -1, -1, -1
+        )
+        self.reset()
+
+    def reset(self):
+        self.data_list = []
+        self.output_data = defaultdict()
+
+    def digest(self, gt, pred):
+        assert len(gt) == 1
+
+        img_heights = [gt[0]["height"]]
+        img_widths = [gt[0]["width"]]
+        img_ids = [gt[0]["image_id"]]
+
+        data_list_elem, image_wise_data = self._evaluator.convert_to_coco_format(
+            pred, [img_heights, img_widths], img_ids, return_outputs=True
+        )
+        self.data_list.extend(data_list_elem)
+        self.output_data.update(image_wise_data)
+
+    def results(self, save_path: str = None):
+        dummy_statistics = torch.FloatTensor([0, 0, len(self.output_data)])
+        eval_results = self._evaluator.evaluate_prediction(
+            self.data_list, dummy_statistics
+        )
+
+        if save_path:
+            self.write_results(eval_results, save_path)
+
+        self.write_results(eval_results)
+
+        *listed_items, summary = eval_results
+
+        self._logger.info("\n" + summary)
+
+        return {"AP": listed_items[0] * 100, "AP50": listed_items[1] * 100}
+
 
 @register_evaluator("VISUAL-QUALITY-EVAL")
 class VisualQualityEval(BaseEvaluator):

compressai_vision/model_wrappers/__init__.py

@@ -35,6 +35,7 @@
     mask_rcnn_X_101_32x8d_FPN_3x,
 )
 from .jde import jde_1088x608
+from .yolox import yolox_darknet53
 
 __all__ = [
     "BaseWrapper",
@@ -43,4 +44,5 @@
     "faster_rcnn_R_50_FPN_3x",
     "mask_rcnn_R_50_FPN_3x",
     "jde_1088x608",
+    "yolox_darknet53",
 ]