[feat] support semantic segmentation (#17)

Author: jianwensong

cfgs/vision_model/default.yaml

@@ -26,6 +26,12 @@ mask_rcnn_X_101_32x8d_FPN_3x:
   weights: "weights/detectron2/COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x/139653917/model_final_2d9806.pkl"
   splits : "fpn" #, "c2" or "r2"
 
+panoptic_rcnn_R_101_FPN_3x:
+  model_path_prefix: ${..model_root_path}
+  cfg: "models/detectron2/configs/COCO-PanopticSegmentation/panoptic_fpn_R_101_3x.yaml"
+  weights: "weights/detectron2/COCO-PanopticSegmentation/panoptic_fpn_R_101_3x/139514519/model_final_cafdb1.pkl"
+  splits : "fpn"
+
 jde_1088x608:
   model_path_prefix: ${..model_root_path}
   cfg: "models/Towards-Realtime-MOT/cfg/yolov3_1088x608.cfg"

compressai_vision/codecs/encdec_utils/png_yuv.py

@@ -79,11 +79,13 @@ def __call__(self, input: Dict, file_prefix: str):
         if len(file_names) > 1:  # video
             # NOTE: using glob for now, should be more robust and look at skipped
             # NOTE: somewhat rigid pattern (lowercase png)
-            filename_pattern = f"{str(Path(file_names[0]).parent)}/*.png"
+
+            parent = Path(file_names[0]).parent
+            ext = next((e for e in ["*.png", "*.jpg"] if list(parent.glob(e))), None)
+            filename_pattern = f"{parent}/{ext}"
+            images_in_folder = len(list(parent.glob(ext)))
             nb_frames = input["last_frame"] - input["frame_skip"]
-            images_in_folder = len(
-                [file for file in Path(file_names[0]).parent.glob("*.png")]
-            )
+
             assert (
                 images_in_folder == nb_frames
             ), f"input folder contains {images_in_folder} images, {nb_frames} were expected"
@@ -302,7 +304,7 @@ def _determine_output_filename(self, output_file_prefix: str):
             cmd_suffix = ["-start_number", "0"]
             prefix = output_file_prefix.split("qp")[0]
             filename = f"{prefix}%03d.png"
-        elif datacatalog in ["MPEGHIEVE"]:
+        elif datacatalog in ["MPEGHIEVE", "PANDASET"]:
             cmd_suffix = ["-start_number", "0"]
             filename = f"%06d.png"
         elif datacatalog in ["MPEGTVDTRACKING"]:

compressai_vision/codecs/utils.py

@@ -48,6 +48,7 @@
     "MPEGTVDTRACKING": (-4.722218990325928, 48.58344268798828),
     "MPEGHIEVE": (-1.0795, 11.8232),
     "SFUHW": (-17.8848, 16.69418),
+    "PANDASET": (-30.0, 30.0),
 }
 
 

compressai_vision/datasets/image.py

@@ -35,6 +35,7 @@
 from pathlib import Path
 from typing import Dict, List
 
+import numpy as np
 from detectron2.data import DatasetCatalog, MetadataCatalog
 from detectron2.data.common import DatasetFromList, MapDataset
 from detectron2.data.dataset_mapper import DatasetMapper
@@ -681,6 +682,80 @@ def get_min_max_across_tensors(self):
         return (minv, maxv)
 
 
+@register_datacatalog("PANDASET")
+class PANDASET(DataCatalog):
+    """Load an image folder database with Detectron2 Cfg. testing image samples
+    and annotations are respectively stored in separate directories
+    (Currently this class supports none of training related operation ):
+
+    .. code-block:: none
+        - rootdir/
+            - camera
+                - front_camera
+                    - 00.jpg
+                    - 01.jpg
+                    - xx.jpg
+            - annotations
+                - xxxx.json TODO
+    Args:
+        root (string): root directory of the dataset
+
+    """
+
+    def __init__(
+        self,
+        root,
+        imgs_folder="camera/front_camera",
+        annotation_file=None,
+        seqinfo="seqinfo.ini",
+        dataset_name="pandaset",
+        ext="jpg",
+    ):
+        super().__init__(
+            root,
+            imgs_folder=imgs_folder,
+            annotation_file=annotation_file,
+            seqinfo=seqinfo,
+            dataset_name=dataset_name,
+            ext=ext,
+        )
+
+        img_lists = sorted(self.imgs_folder_path.glob(f"*.{ext}"))
+
+        # self.data_type = "mot"
+        # print(annotation_file)
+        # seq_id = os.path.splitext(os.path.split(annotation_file)[1])[0]
+        gt_frame_list = np.load(self.annotation_path, allow_pickle=True)[
+            "gt"
+        ]  # read_results(
+        gt_frame_dict = {k: v for k, v in enumerate(gt_frame_list)}
+        #    str(self.annotation_path), self.data_type, is_gt=True
+        # )
+
+        self._dataset = []
+        self._gt_labels = gt_frame_dict
+        # self._gt_ignore_labels = gt_ignore_frame_dict
+
+        for file_name in img_lists:
+            img_id = file_name.name.split(f".{ext}")[0]
+
+            new_d = {
+                "file_name": str(file_name),
+                "image_id": img_id,
+                "annotations": {
+                    "gt": gt_frame_dict.get(int(img_id), []),
+                    # "gt_ignore": gt_ignore_frame_dict.get(int(img_id), []),
+                },
+            }
+            self._dataset.append(new_d)
+
+    def get_min_max_across_tensors(self):
+        # FIXME
+        minv = -30.0
+        maxv = 30.0
+        return (minv, maxv)
+
+
 @register_datacatalog("COCO")
 class COCO(DataCatalog):
     """Load an image folder database with Detectron2 Cfg. testing image samples

compressai_vision/evaluators/base_evaluator.py

@@ -73,6 +73,10 @@ def get_coco_eval_info_name(name=None):
         # the not used input is intended to interface with the function computing class-wise peformance
         return "coco_instances_results.json"
 
+    @staticmethod
+    def get_miou_eval_info_name(name):
+        return f"SemanticSegmentationEval_on_PANDASET_{name}.json"
+
     def reset(self):
         raise NotImplementedError