Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion .github/workflows/pre_commit.yml
Original file line number Diff line number Diff line change
Expand Up @@ -152,9 +152,12 @@ jobs:

- *install-dependencies

- name: Run python unit tests
- name: Run model_api unit tests
run: uv --directory model_api run pytest tests/unit --cov

- name: Run model_converter unit tests
run: uv --directory model_converter run --group tests pytest tests/unit --cov

- &prepare-test-data
name: Prepare test data
run: |
Expand Down
29 changes: 29 additions & 0 deletions model_converter/examples/config.json
Original file line number Diff line number Diff line change
Expand Up @@ -618,6 +618,35 @@
"scale_values": "58.395 57.12 57.375",
"license": "apache-2.0",
"license_link": "https://spdx.org/licenses/Apache-2.0.html"
},
{
"model_short_name": "maskrcnn_resnet50_fpn",
"model_class_name": "torchvision.models.detection.maskrcnn_resnet50_fpn",
"model_library": "torchvision",
"model_full_name": "Mask R-CNN ResNet-50 FPN",
"description": "Mask R-CNN with a ResNet-50-FPN backbone trained on COCO for object detection and instance segmentation",
"docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.detection.maskrcnn_resnet50_fpn.html",
"weights_url": "https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth",
"input_shape": [1, 3, 800, 800],
"input_names": ["image"],
"output_names": ["boxes", "labels", "masks"],
"model_params": null,
"model_type": "MaskRCNN",
"reverse_input_channels": true,
"mean_values": "0 0 0",
"scale_values": "255 255 255",
"resize_type": "fit_to_window_letterbox",
"pad_value": 0,
"input_dtype": "u8",
"confidence_threshold": 0.5,
"postprocess_semantic_masks": true,
"nms_execute": false,
"iou_threshold": 0.5,
"agnostic_nms": false,
"nms_max_predictions": 200,
"license": "bsd-3-clause",
"license_link": "https://spdx.org/licenses/BSD-3-Clause.html",
"labels": "COCO_V1"
}
]
}
15 changes: 14 additions & 1 deletion model_converter/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,9 @@ fixable = ["ALL"]
unfixable = []
dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"

[tool.ruff.lint.per-file-ignores]
"**/tests/**/*.py" = ["SLF001", "FBT003"]

[tool.ruff.lint.mccabe]
max-complexity = 15

Expand All @@ -193,6 +196,16 @@ notice-rgx = """
[tool.bandit]
skips = ["B101", "B310"]

[tool.pytest.ini_options]
pythonpath = ["src"]

[tool.coverage.run]
source = ["model_converter"]

[tool.coverage.report]
fail_under = 45
fail_under = 100
show_missing = true
exclude_lines = [
"if __name__ == .__main__.",
"if TYPE_CHECKING:",
]
2 changes: 1 addition & 1 deletion model_converter/src/model_converter/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,6 @@

"""Tools for converting models to OpenVINO IR."""

from .model_converter import ModelConverter, list_models, main
from .cli import ModelConverter, list_models, main

__all__ = ["ModelConverter", "list_models", "main"]
7 changes: 4 additions & 3 deletions model_converter/src/model_converter/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,8 @@

"""Run the model converter with ``python -m model_converter``."""

from .model_converter import main
import sys

if __name__ == "__main__":
raise SystemExit(main())
from model_converter.cli import main

sys.exit(main())
42 changes: 42 additions & 0 deletions model_converter/src/model_converter/adapters/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
#
# Copyright (C) 2026 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
#

"""Export adapters for different model types."""

from __future__ import annotations

from typing import TYPE_CHECKING

from model_converter.adapters.base import ExportAdapter
from model_converter.adapters.maskrcnn import TorchvisionMaskRCNNExportAdapter

if TYPE_CHECKING:
import torch

_ADAPTER_REGISTRY: dict[str, type[ExportAdapter]] = {
"maskrcnn": TorchvisionMaskRCNNExportAdapter,
}


def get_adapter(model_type: str, model: torch.nn.Module) -> torch.nn.Module:
"""
Get the appropriate export adapter for a model type.

If no adapter is registered for the model type, returns the model unchanged.

Args:
model_type: Model type string (e.g., "MaskRCNN")
model: The PyTorch model to adapt

Returns:
Adapted model (or original model if no adapter needed)
"""
adapter_class = _ADAPTER_REGISTRY.get(model_type.lower())
if adapter_class is not None:
return adapter_class(model)
return model


__all__ = ["ExportAdapter", "TorchvisionMaskRCNNExportAdapter", "get_adapter"]
16 changes: 16 additions & 0 deletions model_converter/src/model_converter/adapters/base.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
#
# Copyright (C) 2026 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
#

"""Base export adapter interface."""

import torch.nn as nn


class ExportAdapter(nn.Module):
"""Base class for export adapters that reshape model outputs for Model API."""

def __init__(self, model: nn.Module):
super().__init__()
self.model = model
31 changes: 31 additions & 0 deletions model_converter/src/model_converter/adapters/maskrcnn.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
#
# Copyright (C) 2026 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
#

"""Mask R-CNN export adapter for TorchVision models."""

from collections import OrderedDict

import torch

from model_converter.adapters.base import ExportAdapter


class TorchvisionMaskRCNNExportAdapter(ExportAdapter):
"""Adapt TorchVision Mask R-CNN to the Model API MaskRCNN output contract."""

def forward(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Return boxes-with-scores, shifted labels, and raw masks for one image."""
image_list = [images[0]]
transformed_images, _ = self.model.transform(image_list, None)
features = self.model.backbone(transformed_images.tensors)
if isinstance(features, torch.Tensor):
features = OrderedDict([("0", features)])
proposals, _ = self.model.rpn(transformed_images, features, None)
predictions, _ = self.model.roi_heads(features, proposals, transformed_images.image_sizes, None)
prediction = predictions[0]
boxes = torch.cat((prediction["boxes"], prediction["scores"].unsqueeze(1)), dim=1)
labels = prediction["labels"] - 1
masks = prediction["masks"].squeeze(1)
return boxes, labels, masks
Loading