feat: enable to install a subset of vision models only

Author: fracape

.gitignore

@@ -2,7 +2,6 @@
 *.bin
 *.inc
 *.tar.gz
-*.sh
 .DS_Store
 builds
 compressai_vision/version.py

compressai_vision/datasets/utils.py

@@ -35,9 +35,6 @@
 import numpy as np
 import torch
 
-from jde.utils.datasets import letterbox
-from mmpose.structures.bbox import get_warp_matrix
-from segment_anything.utils.transforms import ResizeLongestSide
 from torch.nn import functional as F
 from torchvision import transforms
 
@@ -130,6 +127,7 @@ def __call__(self, dataset_dict):
         Returns:
             dict: a format that compressai-vision pipelines accept
         """
+        from mmpose.structures.bbox import get_warp_matrix
 
         dataset_dict = copy.deepcopy(dataset_dict)
         # the copied dictionary will be modified by code below
@@ -284,6 +282,7 @@ def __call__(self, dataset_dict):
         Returns:
             dict: a format that compressai-vision pipelines accept
         """
+        from jde.utils.datasets import letterbox
 
         dataset_dict = copy.deepcopy(dataset_dict)
         # the copied dictionary will be modified by code below
@@ -313,6 +312,8 @@ def __init__(self, img_size=1024):
         Args:
             img_size: single value - target size to SAM as input
         """
+        from segment_anything.utils.transforms import ResizeLongestSide
+
         self.target_size = img_size
         self.transform = ResizeLongestSide(img_size)
 

compressai_vision/evaluators/evaluators.py

@@ -41,19 +41,8 @@
 import pandas as pd
 import torch
 
-from detectron2.data import MetadataCatalog
-from detectron2.evaluation import COCOEvaluator
-from detectron2.utils.visualizer import Visualizer
-from jde.utils.io import unzip_objs
-from mmpose.datasets.datasets import BaseCocoStyleDataset
-from mmpose.datasets.transforms import PackPoseInputs
-from mmpose.evaluation.metrics import CocoMetric
-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 yolox.utils import xyxy2xywh
 
 from compressai_vision.datasets import deccode_compressed_rle
 from compressai_vision.registry import register_evaluator
@@ -132,6 +121,8 @@ def __init__(
 
         self.set_annotation_info(dataset)
 
+        from detectron2.evaluation import COCOEvaluator
+
         self._evaluator = COCOEvaluator(
             dataset_name, False, output_dir=output_dir, use_fast_impl=False
         )
@@ -156,6 +147,9 @@ def save_visualization(self, gt, pred, output_dir, threshold):
         gt_image = cv2.resize(gt_image, (gt[0]["width"], gt[0]["height"]))
 
         img_id = gt[0]["image_id"]
+        from detectron2.data import MetadataCatalog
+        from detectron2.utils.visualizer import Visualizer
+
         metadata = MetadataCatalog.get(self.dataset_name)
         instances = pred[0]["instances"].to("cpu")
         if threshold:
@@ -343,6 +337,9 @@ def save_visualization(self, gt, pred, output_dir, threshold):
         gt_image = cv2.resize(gt_image, (gt[0]["width"], gt[0]["height"]))
 
         img_id = gt[0]["image_id"]
+        from detectron2.data import MetadataCatalog
+        from detectron2.utils.visualizer import Visualizer
+
         metadata = MetadataCatalog.get(self.dataset_name)
         instances = pred[0]["instances"].to("cpu")
 
@@ -575,6 +572,10 @@ def __init__(
 
         self.set_annotation_info(dataset)
 
+        from jde.utils.io import unzip_objs
+
+        self.unzip_objs = unzip_objs
+
         mm.lap.default_solver = "lap"
         self.dataset = dataset.dataset
         self.eval_info_file_name = self.get_jde_eval_info_name(self.dataset_name)
@@ -734,13 +735,13 @@ def mot_eval(self):
             frm_id = int(gt_frame["image_id"])
 
             pred_objs = self._predictions[frm_id].copy()
-            pred_tlwhs, pred_ids, _ = unzip_objs(pred_objs)
+            pred_tlwhs, pred_ids, _ = self.unzip_objs(pred_objs)
 
             gt_objs = gt_frame["annotations"]["gt"].copy()
-            gt_tlwhs, gt_ids, _ = unzip_objs(gt_objs)
+            gt_tlwhs, gt_ids, _ = self.unzip_objs(gt_objs)
 
             gt_ignore = gt_frame["annotations"]["gt_ignore"].copy()
-            gt_ignore_tlwhs, _, _ = unzip_objs(gt_ignore)
+            gt_ignore_tlwhs, _, _ = self.unzip_objs(gt_ignore)
 
             # remove ignored results
             keep = np.ones(len(pred_tlwhs), dtype=bool)
@@ -913,6 +914,10 @@ def __init__(
             datacatalog_name, dataset_name, dataset, output_dir, eval_criteria
         )
 
+        from pycocotools.coco import COCO
+        from yolox.data.datasets.coco import remove_useless_info
+        from yolox.evaluators import COCOEvaluator as YOLOX_COCOEvaluator
+
         self.set_annotation_info(dataset)
 
         cocoapi = COCO(self.annotation_path)
@@ -1055,6 +1060,16 @@ def __init__(
             dataset.get_org_mapper_func().compute_scale_and_center
         )
 
+        try:
+            from mmpose.datasets.datasets import BaseCocoStyleDataset
+            from mmpose.datasets.transforms import PackPoseInputs
+            from mmpose.evaluation.metrics import CocoMetric
+        except ImportError:
+            self._logger.error(
+                "Failed to import mmpose. Please install it, e.g. with 'pip install mmpose'."
+            )
+            raise
+
         if "metainfo" in args:
             metainfo = args["metainfo"]
         else:

compressai_vision/model_wrappers/__init__.py

@@ -27,34 +27,7 @@
 # OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-from .detectron2 import (
-    BaseWrapper,
-    faster_rcnn_R_50_FPN_3x,
-    faster_rcnn_X_101_32x8d_FPN_3x,
-    mask_rcnn_R_50_FPN_3x,
-    mask_rcnn_X_101_32x8d_FPN_3x,
-    panoptic_rcnn_R_101_FPN_3x,
-)
-from .jde import jde_1088x608
-from .rtmo import rtmo_multi_person_pose_estimation
-from .sam import (
-    sam_vit_b_01ec64,
-    sam_vit_h_4b8939,
-    sam_vit_l_0b3195,
-)
-from .yolox import yolox_darknet53
+from .base_wrapper import BaseWrapper
+from . import detectron2, jde, rtmo, sam, yolox
 
-__all__ = [
-    "BaseWrapper",
-    "faster_rcnn_X_101_32x8d_FPN_3x",
-    "mask_rcnn_X_101_32x8d_FPN_3x",
-    "faster_rcnn_R_50_FPN_3x",
-    "mask_rcnn_R_50_FPN_3x",
-    "panoptic_rcnn_R_101_FPN_3x",
-    "jde_1088x608",
-    "yolox_darknet53",
-    "rtmo_multi_person_pose_estimation",
-    "sam_vit_h_4b8939",
-    "sam_vit_b_01ec64",
-    "sam_vit_l_0b3195",
-]
+__all__ = ["BaseWrapper"]

compressai_vision/model_wrappers/detectron2.py

@@ -34,15 +34,6 @@
 
 import torch
 
-from detectron2.checkpoint import DetectionCheckpointer
-from detectron2.config import get_cfg
-from detectron2.modeling import build_model
-from detectron2.modeling.meta_arch.panoptic_fpn import (
-    combine_semantic_and_instance_outputs,
-    detector_postprocess,
-    sem_seg_postprocess,
-)
-from detectron2.structures import ImageList
 
 from compressai_vision.registry import register_vision_model
 
@@ -159,10 +150,16 @@ def __str__(self):
 
 class Rcnn_R_50_X_101_FPN(BaseWrapper):
     def __init__(self, device: str, **kwargs):
+        from detectron2.checkpoint import DetectionCheckpointer
+        from detectron2.config import get_cfg
+        from detectron2.modeling import build_model
+
         super().__init__(device)
 
         self._cfg = get_cfg()
         self._cfg.MODEL.DEVICE = device
+        self.DetectionCheckpointer = DetectionCheckpointer
+        self.build_model = build_model
         _path_prefix = (
             f"{root_path}"
             if kwargs["model_path_prefix"] == "default"
@@ -171,11 +168,11 @@ def __init__(self, device: str, **kwargs):
         self._cfg.merge_from_file(f"{_path_prefix}/{kwargs['cfg']}")
         _integer_conv_weight = bool(kwargs["integer_conv_weight"])
 
-        self.model = build_model(self._cfg)
+        self.model = self.build_model(self._cfg)
         self.replace_conv2d_modules(self.model)
         self.model = self.model.to(device).eval()
 
-        DetectionCheckpointer(self.model).load(f"{_path_prefix}/{kwargs['weights']}")
+        self.DetectionCheckpointer(self.model).load(f"{_path_prefix}/{kwargs['weights']}")
 
         for param in self.model.parameters():
             param.requires_grad = False
@@ -271,6 +268,8 @@ def quantize_weights(model):
         return model
 
     def input_resize(self, images: List):
+        from detectron2.structures import ImageList
+
         return ImageList.from_tensors(images, self.size_divisibility)
 
     def input_to_features(self, x, device: str) -> Dict:
@@ -540,7 +539,18 @@ def __init__(self, device: str, **kwargs):
 @register_vision_model("panoptic_rcnn_R_101_FPN_3x")
 class panoptic_rcnn_R_101_FPN_3x(Rcnn_R_50_X_101_FPN):
     def __init__(self, device="cpu", **kwargs):
+        from detectron2.modeling.meta_arch.panoptic_fpn import (
+            combine_semantic_and_instance_outputs,
+            detector_postprocess,
+            sem_seg_postprocess,
+        )
+
         super().__init__(device, **kwargs)
+        self.sem_seg_postprocess = sem_seg_postprocess
+        self.detector_postprocess = detector_postprocess
+        self.combine_semantic_and_instance_outputs = (
+            combine_semantic_and_instance_outputs
+        )
         self.sem_seg_head = self.model.sem_seg_head
 
         combine_overlap_thresh = 0.5
@@ -590,12 +600,12 @@ def __init__(self, img_size: list):
         ):
             height = input_per_image["height"]
             width = input_per_image["width"]
-            sem_seg_r = sem_seg_postprocess(sem_seg_result, image_size, height, width)
-            detector_r = detector_postprocess(detector_result, height, width)
+            sem_seg_r = self.sem_seg_postprocess(sem_seg_result, image_size, height, width)
+            detector_r = self.detector_postprocess(detector_result, height, width)
 
             processed_results.append({"sem_seg": sem_seg_r, "instances": detector_r})
 
-            panoptic_r = combine_semantic_and_instance_outputs(
+            panoptic_r = self.combine_semantic_and_instance_outputs(
                 detector_r,
                 sem_seg_r.argmax(dim=0),
                 self.combine_overlap_thresh,

compressai_vision/model_wrappers/jde.py

@@ -32,30 +32,12 @@
 from pathlib import Path
 from typing import Dict, List
 
-import jde
 import torch
 
-from jde.models import Darknet
-from jde.tracker import matching
-from jde.tracker.basetrack import TrackState
-from jde.tracker.multitracker import (
-    STrack,
-    joint_stracks,
-    remove_duplicate_stracks,
-    sub_stracks,
-)
-from jde.utils.kalman_filter import KalmanFilter
-from jde.utils.utils import non_max_suppression, scale_coords
-
 from compressai_vision.registry import register_vision_model
 
 from .base_wrapper import BaseWrapper
 
-# Patch in modified create_modules
-from .jde_lowlevel import create_modules
-
-jde.models.create_modules = create_modules
-
 __all__ = [
     "jde_1088x608",
 ]
@@ -67,8 +49,19 @@
 @register_vision_model("jde_1088x608")
 class jde_1088x608(BaseWrapper):
     def __init__(self, device: str, **kwargs):
+        import jde
+        from jde.models import Darknet
+        from jde.utils.kalman_filter import KalmanFilter
+
+        from .jde_lowlevel import create_modules
+
+        jde.models.create_modules = create_modules
+
         super().__init__(device)
 
+        self.Darknet = Darknet
+        self.KalmanFilter = KalmanFilter
+
         _path_prefix = (
             f"{root_path}"
             if kwargs["model_path_prefix"] == "default"
@@ -99,7 +92,7 @@ def __init__(self, device: str, **kwargs):
             zip(self.split_layer_list, [None] * len(self.split_layer_list))
         )
 
-        self.darknet = Darknet(self.model_info["cfg"], device, nID=14455)
+        self.darknet = self.Darknet(self.model_info["cfg"], device, nID=14455)
         self.darknet.load_state_dict(
             torch.load(self.model_info["weights"], map_location="cpu")["model"],
             strict=False,
@@ -112,7 +105,7 @@ def __init__(self, device: str, **kwargs):
         if _integer_conv_weight:
             self.darknet = self.quantize_weights(self.darknet)
 
-        self.kalman_filter = KalmanFilter()
+        self.kalman_filter = self.KalmanFilter()
 
         if "logging_level" in kwargs:
             self.logger.level = kwargs["logging_level"]
@@ -219,6 +212,15 @@ def deeper_features_for_accuracy_proxy(self, x: Dict):
         # return x_deeper
 
     def _jde_process(self, pred, org_img_size: tuple, input_img_size: tuple):
+        from jde.tracker import matching
+        from jde.tracker.basetrack import TrackState
+        from jde.tracker.multitracker import (
+            STrack,
+            joint_stracks,
+            remove_duplicate_stracks,
+            sub_stracks,
+        )
+        from jde.utils.utils import non_max_suppression, scale_coords
         r"""Re-implementation of JDE from Z. Wang, L. Zheng, Y. Liu, and S. Wang:
         : `"Towards Real-Time Multi-Object Tracking"`_,
         The European Conference on Computer Vision (ECCV), 2020