fix: install mmpose in both cpu and gpu modes

Author: fracape

.gitlab-ci.yml

@@ -17,7 +17,7 @@ stages:
   - bash scripts/install_uv.sh --no-install --no-weights --cpu
 
 uv-install:
-  image: python:3.8-buster
+  image: python:3.8-bullseye
   stage: setup
   script:
     - curl -LsSf https://astral.sh/uv/install.sh | sh
@@ -30,7 +30,7 @@ uv-install:
     - docker
 ruff-lint:
   stage: static-analysis
-  image: python:3.8-buster
+  image: python:3.8-bullseye
   before_script:
     - *default-before-script
     - uv pip install -U pip wheel setuptools
@@ -42,10 +42,15 @@ ruff-lint:
 
 pytest:
   stage: test
-  image: python:3.8-buster
+  image: python:3.8-bullseye
   before_script:
+    - apt-get update -y
+    - apt-get install -y libgl1-mesa-glx build-essential
     - *default-before-script
     - uv pip install -U pip wheel setuptools
+    - uv pip install torch==1.8.0+cpu torchvision==0.9.0+cpu -f https://download.pytorch.org/whl/torch_stable.html
+    - bash scripts/install_uv.sh --no-weights --cpu
+    - uv pip install pycocotools "fvcore<0.1.6" "iopath<0.1.10" omegaconf "hydra-core" yacs PyYAML termcolor tabulate cloudpickle future pydot protobuf==3.20.0
     - uv sync --group=dev
   script:
     - uv run --no-sync pytest
@@ -54,7 +59,7 @@ pytest:
 
 ruff-format:
   stage: static-analysis
-  image: python:3.8-buster
+  image: python:3.8-bullseye
   before_script:
     - *default-before-script
     - uv pip install -U pip wheel setuptools
@@ -66,7 +71,7 @@ ruff-format:
 
 ruff-organize-imports:
   stage: static-analysis
-  image: python:3.8-buster
+  image: python:3.8-bullseye
   before_script:
     - *default-before-script
     - uv pip install -U pip wheel setuptools
@@ -78,7 +83,7 @@ ruff-organize-imports:
 
 doc:
   stage: doc
-  image: python:3.8-buster
+  image: python:3.8-bullseye
   before_script:
     - *default-before-script
     - uv pip install -U pip wheel setuptools

pyproject.toml

@@ -74,19 +74,7 @@ cpu = [
     "torch>=2.0.0",
     "torchvision>=0.15.1",
 ]
-cu118 = [
-    "torch>=2.0.0",
-    "torchvision>=0.15.1",
-]
-cu128 = [
-    "torch>=2.7.0 ; python_version >= '3.9'",
-    "torchvision>=0.22.0 ; python_version >= '3.9'",
-]
-models-detectron2 = [
-    "detectron2==0.6",
-]
 models-jde = [
-    "cython-bbox",  # Encourages local cython-bbox source.
     "jde",
 ]
 models-mmpose = [
@@ -129,11 +117,7 @@ tutorials = [
     "ipywidgets",
     "jupyter",
 ]
-models-detectron2 = [
-    "detectron2==0.6",
-]
 models-jde = [
-    "cython-bbox",  # Encourages local cython-bbox source.
     "jde",
 ]
 models-mmpose = [
@@ -162,16 +146,8 @@ skip_gitignore = true
 known_third_party = "PIL,torchvision,torch"
 
 [tool.uv]
-conflicts = [
-    [
-        { extra = "cpu" },
-        { extra = "cu118" },
-        { extra = "cu128" },
-    ],
-]
 no-build-isolation-package = [
     "chumpy",  # build-system.requires: pip
-    "detectron2",  # build-system.requires: Cython, numpy, setuptools
     "mmpose",  # build-system.requires: numpy, torch
     "xtcocotools",  # build-system.requires: Cython, numpy
     "yolox",  # build-system.requires: pycocotools, torch, yolox, ...
@@ -187,29 +163,6 @@ requires-dist = [
     "six",
 ]
 
-[[tool.uv.dependency-metadata]]
-name = "detectron2"
-version = "0.6"
-requires-dist = [
-    # From https://github.com/facebookresearch/detectron2/blob/main/setup.py
-    "Pillow>=7.1",
-    "matplotlib",
-    "pycocotools>=2.0.2",
-    "termcolor>=1.1",
-    "yacs>=0.1.8",
-    "tabulate",
-    "cloudpickle",
-    "tqdm>4.29.0",
-    "tensorboard",
-    "fvcore>=0.1.5,<0.1.6",
-    "iopath>=0.1.7,<0.1.10",
-    "future",
-    "pydot",
-    "dataclasses; python_version<'3.7'",
-    "omegaconf>=2.1",
-    "hydra-core>=1.1",
-    # "black==21.4b2",
-]
 
 [[tool.uv.dependency-metadata]]
 name = "mmpose"
@@ -273,34 +226,17 @@ name = "pytorch-cpu"
 url = "https://download.pytorch.org/whl/cpu"
 explicit = true
 
-[[tool.uv.index]]
-name = "pytorch-cu118"
-url = "https://download.pytorch.org/whl/cu118"
-explicit = true
+[tool.uv.overrides]
+jde = { path = "models/Towards-Realtime-MOT" }
+yolox = { path = "models/yolox" }
+segment_anything = { path = "models/segment_anything" }
 
-[[tool.uv.index]]
-name = "pytorch-cu128"
-url = "https://download.pytorch.org/whl/cu128"
-explicit = true
+torch = { index = "pytorch-cpu" }
+torchvision = { index = "pytorch-cpu" }
 
 [tool.uv.sources]
 compressai = { workspace = true }
-cython-bbox = { path = "scripts/cython_bbox" }
-detectron2 = { path = "models/detectron2" }
-jde = { path = "models/Towards-Realtime-MOT" }
-mmpose = { path = "models/mmpose" }
-yolox = { path = "models/yolox" }
-segment_anything = { path = "models/segment_anything" }
-torch = [
-    { index = "pytorch-cpu", extra = "cpu" },
-    { index = "pytorch-cu118", extra = "cu118" },
-    { index = "pytorch-cu128", extra = "cu128" },
-]
-torchvision = [
-    { index = "pytorch-cpu", extra = "cpu" },
-    { index = "pytorch-cu118", extra = "cu118" },
-    { index = "pytorch-cu128", extra = "cu128" },
-]
+
 
 [tool.uv.workspace]
 members = [

scripts/install.sh

@@ -33,21 +33,16 @@ $ python3 -m venv venv
 $ source venv/bin/activate
 
 RUN OPTIONS:
-                [-m|--model, vision models to install, (detectron2/jde/yolox/mmpose/segment-anything/all) default=all]
+                [-m|--model, vision models to install, (jde/yolox/mmpose/segment-anything/all) default=all]
                 [-t|--torch torch version, default="2.0.0"]
                 [--torchvision torchvision version, default="0.15.1"]
                 [--cpu) build for cpu only)]
                 [--cuda_version) provide cuda version e.g. "11.8", default: check nvcc output)]
-                [--detectron2_url use this if you want to specify a pre-built detectron2 (find at
-                    "https://detectron2.readthedocs.io/en/latest/tutorials/install.html#install-pre-built-detectron2-linux-only"),
-                    not required for regular versions derived from cuda and torch versions above.
-                    default:"https://dl.fbaipublicfiles.com/detectron2/wheels/cu102/torch1.9/index.html"]
                 [--models_dir directory to install vision models to, default: compressai_vision_root]
                 [--no-install) do not install (i.e. useful for only preparing source code by downloading and patching 
                 [--no-weights) prevents the installation script from downloading vision model parameters]
 
 
-EXAMPLE         [bash install.sh -m detectron2 -t "1.9.1" --cuda_version "11.8" --compressai /path/to/compressai]
 
 _EOF_
             exit;
@@ -67,11 +62,6 @@ done;
 
 
 WEIGHTS="
-3c25caca37baabbff3e22cc9eb0923db165a0c18b867871a3bf3570bac9b7ef0  detectron2/COCO-Detection/faster_rcnn_R_50_FPN_3x/137849458/model_final_280758.pkl                  https://dl.fbaipublicfiles.com/detectron2/COCO-Detection/faster_rcnn_R_50_FPN_3x/137849458/model_final_280758.pkl
-fe5ad56ff746aa55c5f453b01f8395134e9281d240dbeb473411d4a6b262c9dc  detectron2/COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x/139173657/model_final_68b088.pkl           https://dl.fbaipublicfiles.com/detectron2/COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x/139173657/model_final_68b088.pkl
-9a737e290372f1f70994ebcbd89d8004dbb3ae30a605fd915a190fa4a782dd66  detectron2/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl         https://dl.fbaipublicfiles.com/detectron2/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl
-12f6e1811baf1b4d329c3f5ac5ec52d8f634d3cedc82a13fff55d0c05d84f442  detectron2/COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x/139653917/model_final_2d9806.pkl  https://dl.fbaipublicfiles.com/detectron2/COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x/139653917/model_final_2d9806.pkl
-808c675e647298688589c895c9581f7f3963995c5708bc53f66449200321d147  detectron2/COCO-PanopticSegmentation/panoptic_fpn_R_101_3x/139514519/model_final_cafdb1.pkl         https://dl.fbaipublicfiles.com/detectron2/COCO-PanopticSegmentation/panoptic_fpn_R_101_3x/139514519/model_final_cafdb1.pkl
 6b135b0affa38899b607010c86c2f8dbc1c06956bad9ca1edd45b01e626933f1  jde/jde.1088x608.uncertainty.pt                                                                     https://drive.usercontent.google.com/download?export=download&confirm=t&id=1nlnuYfGNuHWZztQHXwVZSL_FvfE551pA
 516a421f8717548300c3ee6356a3444ac539083d4a9912f8ca1619ee63d0986d  mmpose/rtmo_coco/rtmo-l_16xb16-600e_coco-640x640-516a421f_20231211.pth                              https://download.openmmlab.com/mmpose/v1/projects/rtmo/rtmo-l_16xb16-600e_coco-640x640-516a421f_20231211.pth
 b5905e9faf500a2608c93991f91a41a6150bcd2dd30986865a73becd94542fa1  yolox/darknet53/yolox_darknet.pth                                                                   https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_darknet.pth
@@ -181,8 +171,7 @@ run_install () {
     if [[ "${PACKAGE_MANAGER}" == "uv" ]]; then
         echo
         echo "Detect differences from uv.lock:"
-        echo "uv sync --inexact --extra=${BUILD_SUFFIX} --dry-run"
-        uv sync --inexact --extra="${BUILD_SUFFIX}" --dry-run
+        echo "uv sync --inexact --extra=${BUILD_SUFFIX}"
         uv sync --inexact --extra="${BUILD_SUFFIX}"
     fi
 }
@@ -205,15 +194,6 @@ detect_cuda_version () {
     echo "Detected CUDA version: ${CUDA_VERSION}"
 }
 
-install_torch () {
-    if [ "${CPU}" == "True" ]; then
-        "${PIP[@]}" install "torch==${TORCH_VERSION}" "torchvision==${TORCHVISION_VERSION}" --index-url "https://download.pytorch.org/whl/${BUILD_SUFFIX}"
-    else
-        # TODO(refactor): If we use --index-url instead, the "+${BUILD_SUFFIX}" does not need to be explicitly specified.
-        "${PIP[@]}" install "torch==${TORCH_VERSION}+${BUILD_SUFFIX}" "torchvision==${TORCHVISION_VERSION}+${BUILD_SUFFIX}" --extra-index-url "https://download.pytorch.org/whl/${BUILD_SUFFIX}"
-    fi
-}
-
 prepare_detectron2 () {
     echo
     echo "Preparing detectron2 for installation"
@@ -226,28 +206,37 @@ prepare_detectron2 () {
 
     git clone --single-branch --branch main https://github.com/facebookresearch/detectron2.git "${MODELS_SOURCE_DIR}/detectron2"
     cd "${MODELS_SOURCE_DIR}/detectron2"
-    git -c advice.detachedHead=false  checkout 175b2453c2bc4227b8039118c01494ee75b08136
+    git -c advice.detachedHead=false checkout c554533136231a94971a995185501539f5c808f9
     git apply "${SCRIPT_DIR}/patches/0001-detectron2-fpn-bottom-up-separate.patch" || echo "Patch could not be applied. Possibly already applied."
     cd "${COMPRESSAI_VISION_ROOT_DIR}"
 }
 
 install_detectron2 () {
     echo
-    echo "Installing detectron2"
+    echo "Installing detectron2 (reference: https://github.com/facebookresearch/detectron2)"
     echo
 
     cd "${MODELS_SOURCE_DIR}/detectron2"
 
     if [[ "${PACKAGE_MANAGER}" == "pip3" ]]; then
         "${PIP[@]}" install -e .
     elif [[ "${PACKAGE_MANAGER}" == "uv" ]]; then
-        cd "${COMPRESSAI_VISION_ROOT_DIR}"
-        uv sync --inexact --group=models-detectron2
+        "${PIP[@]}" install -e .
     fi
 
     cd "${COMPRESSAI_VISION_ROOT_DIR}"
 }
 
+install_torch () {
+    if [ "${CPU}" == "True" ]; then
+        "${PIP[@]}" install "torch==${TORCH_VERSION}" "torchvision==${TORCHVISION_VERSION}" --index-url "https://download.pytorch.org/whl/${BUILD_SUFFIX}"
+    else
+        # TODO(refactor): If we use --index-url instead, the "+${BUILD_SUFFIX}" does not need to be explicitly specified.
+        "${PIP[@]}" install "torch==${TORCH_VERSION}+${BUILD_SUFFIX}" "torchvision==${TORCHVISION_VERSION}+${BUILD_SUFFIX}" --extra-index-url "https://download.pytorch.org/whl/${BUILD_SUFFIX}"
+    fi
+}
+
+
 prepare_cython_bbox () {
     echo
     echo "Preparing cython_bbox for installation"
@@ -339,6 +328,15 @@ prepare_yolox () {
     cd "${COMPRESSAI_VISION_ROOT_DIR}"
 }
 
+install_mmcv () {
+    echo
+    echo "Installing mmcv"
+    echo
+    local torch_version_url_part="torch${TORCH_VERSION%.*}"
+    local find_links_url="https://download.openmmlab.com/mmcv/dist/${BUILD_SUFFIX}/${torch_version_url_part}/index.html"
+    "${PIP[@]}" install "mmcv==2.0.1" --find-links "${find_links_url}"
+}
+
 install_yolox () {
     echo
     echo "Installing YOLOX (reference: https://github.com/Megvii-BaseDetection/YOLOX)"
@@ -375,6 +373,7 @@ prepare_mmpose () {
 }
 
 install_mmpose () {
+
     echo
     echo "Installing MMPOSE (reference: https://github.com/open-mmlab/mmpose/tree/main)"
     echo
@@ -394,7 +393,8 @@ install_mmpose () {
         uv sync --inexact --group=models-mmpose
     fi
 
-    "${MIM[@]}" install "mmcv==2.0.1"
+    install_mmcv
+
     "${MIM[@]}" install "mmdet==3.1.0"
 
     cd "${COMPRESSAI_VISION_ROOT_DIR}"
@@ -439,7 +439,7 @@ download_weights () {
     mkdir -p "${MODELS_WEIGHT_DIR}"
     cd "${MODELS_WEIGHT_DIR}/"
 
-    for model in detectron2 jde mmpose yolox segment_anything; do
+    for model in jde mmpose yolox segment_anything; do
         if ! [[ ",${MODEL,,}," == *",${model},"* ]] && [[ ",${MODEL,,}," != *",all,"* ]]; then
             continue
         fi

tests/evaluators/test_evaluators.py

@@ -36,28 +36,36 @@
 
 from compressai_vision.evaluators import COCOEVal
 
+
 def mock_coco_api(annotation_file):
     # Create a mock COCO API object
     coco_api = MagicMock()
-    coco_api.anns = {1: {'image_id': 1, 'category_id': 1, 'bbox': [10, 10, 50, 50]}}
+    coco_api.anns = {1: {"image_id": 1, "category_id": 1, "bbox": [10, 10, 50, 50]}}
     return coco_api
 
+
 @pytest.fixture
 def mock_dataset():
     dataset = MagicMock()
     dataset.dataset_name = "mock_coco_dataset"
     dataset.annotation_path = "mock_annotations.json"
-    with open(dataset.annotation_path, 'w') as f:
-        json.dump({
-            'images': [{'id': 1, 'width': 640, 'height': 480}],
-            'annotations': [{'id': 1, 'image_id': 1, 'category_id': 1, 'bbox': [10, 10, 50, 50]}],
-            'categories': [{'id': 1, 'name': 'person'}]
-        }, f)
+    with open(dataset.annotation_path, "w") as f:
+        json.dump(
+            {
+                "images": [{"id": 1, "width": 640, "height": 480}],
+                "annotations": [
+                    {"id": 1, "image_id": 1, "category_id": 1, "bbox": [10, 10, 50, 50]}
+                ],
+                "categories": [{"id": 1, "name": "person"}],
+            },
+            f,
+        )
     yield dataset
     os.remove(dataset.annotation_path)
 
-@patch('detectron2.evaluation.COCOEvaluator', autospec=True)
-@patch('compressai_vision.evaluators.evaluators.deccode_compressed_rle')
+
+@patch("detectron2.evaluation.COCOEvaluator", autospec=True)
+@patch("compressai_vision.evaluators.evaluators.deccode_compressed_rle")
 def test_coco_eval(mock_deccode, mock_evaluator, mock_dataset):
     # Configure the mock to have the expected attributes
     mock_coco_api_obj = MagicMock()
@@ -70,11 +78,13 @@ def test_coco_eval(mock_deccode, mock_evaluator, mock_dataset):
         datacatalog_name="MPEGOIV6",
         dataset_name="mock_coco_dataset",
         dataset=mock_dataset,
-        output_dir="./test_output"
+        output_dir="./test_output",
     )
 
-    gt = [{'image_id': 1, 'width': 640, 'height': 480, 'image': torch.zeros(3, 480, 640)}]
-    pred = [{'instances': MagicMock()}]
+    gt = [
+        {"image_id": 1, "width": 640, "height": 480, "image": torch.zeros(3, 480, 640)}
+    ]
+    pred = [{"instances": MagicMock()}]
 
     evaluator.digest(gt, pred)
     mock_evaluator.return_value.process.assert_called_once_with(gt, pred)