diff --git a/model_converter/presets/config.json b/model_converter/presets/config.json index 624ef723..bc7b4be2 100644 --- a/model_converter/presets/config.json +++ b/model_converter/presets/config.json @@ -224,8 +224,8 @@ "model_params": null, "model_type": "Classification", "reverse_input_channels": true, - "mean_values": "123.675 116.28 103.53", - "scale_values": "58.395 57.12 57.375", + "mean_values": "127.5 127.5 127.5", + "scale_values": "127.5 127.5 127.5", "license": "bsd-3-clause", "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", "labels": "IMAGENET1K_V1", @@ -265,7 +265,7 @@ "output_names": ["output1"], "model_params": null, "model_type": "Classification", - "reverse_input_channels": false, + "reverse_input_channels": true, "mean_values": "123.675 116.28 103.53", "scale_values": "58.395 57.12 57.375", "license": "bsd-3-clause", @@ -417,7 +417,7 @@ "resize_type": "fit_to_window_letterbox", "pad_value": 0, "input_dtype": "u8", - "confidence_threshold": 0.5, + "confidence_threshold": 0.05, "postprocess_semantic_masks": true, "nms_execute": false, "iou_threshold": 0.5, @@ -425,7 +425,7 @@ "nms_max_predictions": 200, "license": "bsd-3-clause", "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "COCO_V1", + "labels": "COCO_80", "dataset_type": "coco-detection" }, { diff --git a/model_converter/src/model_converter/adapters/maskrcnn.py b/model_converter/src/model_converter/adapters/maskrcnn.py index 2a0e16c1..d473e086 100644 --- a/model_converter/src/model_converter/adapters/maskrcnn.py +++ b/model_converter/src/model_converter/adapters/maskrcnn.py @@ -10,13 +10,26 @@ import torch from model_converter.adapters.base import ExportAdapter +from model_converter.metrics.coco_detection import COCO91_TO_COCO80 + +# Lookup tensor mapping a torchvision COCO 91-class category ID to the value the +# exported model must emit so that, after the Model API ``MaskRCNN`` wrapper adds +# ``+1`` to every label, the reported label equals the contiguous 80-class index +# (0-79) expected by the labels metadata and the COCO evaluator. We therefore +# store ``COCO80_index - 1`` here. Category IDs that COCO dropped (and the +# background ID 0) are unreachable for real predictions; they map to ``-1`` purely +# to keep the lookup well-defined. +_COCO91_TO_COCO80_LUT = torch.tensor( + [idx - 1 if idx is not None else -1 for idx in COCO91_TO_COCO80], + dtype=torch.int64, +) 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.""" + """Return boxes-with-scores, wrapper-compensated 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) @@ -26,6 +39,7 @@ def forward(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, tor 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 + lut = _COCO91_TO_COCO80_LUT.to(prediction["labels"].device) + labels = lut[prediction["labels"]] masks = prediction["masks"].squeeze(1) return boxes, labels, masks diff --git a/model_converter/src/model_converter/converters/base.py b/model_converter/src/model_converter/converters/base.py index 7232e1b3..f08a17e0 100644 --- a/model_converter/src/model_converter/converters/base.py +++ b/model_converter/src/model_converter/converters/base.py @@ -119,6 +119,128 @@ def _build_result(self, config: dict[str, Any]) -> ConversionResult: original_url=original_url_for_config(config), ) + def _skip_if_already_converted( + self, + config: dict[str, Any], + model_short_name: str, + xml_stem: str | None = None, + ) -> bool: + """Record a skip and return ``True`` when FP16 and INT8 models both exist. + + ``xml_stem`` overrides the ``.xml`` file name inside the variant folders + (defaults to ``model_short_name``); YOLO uses the Ultralytics + ``yolo_version`` here. + """ + xml_stem = xml_stem or model_short_name + fp16_model_path = self.output_dir / f"{model_short_name}-fp16-ov" / f"{xml_stem}.xml" + int8_model_path = self.output_dir / f"{model_short_name}-int8-ov" / f"{xml_stem}.xml" + if fp16_model_path.exists() and int8_model_path.exists(): + self.logger.info(f"Skipping {model_short_name}: FP16 and INT8 models already exist") + self._record_result(self._build_result(config), converted=False, quantized=False, skipped=True) + return True + return False + + def _validate_license(self, config: dict[str, Any], model_short_name: str) -> None: + """Raise ``ValueError`` when ``license`` or ``license_link`` is missing.""" + if not config.get("license"): + error_msg = f"Model '{model_short_name}' must define 'license' in configuration" + raise ValueError(error_msg) + if not config.get("license_link"): + error_msg = f"Model '{model_short_name}' must define 'license_link' in configuration" + raise ValueError(error_msg) + + def _log_model_banner(self, config: dict[str, Any], model_short_name: str, *, label: str = "model") -> None: + """Log the ``"=" * 80`` processing banner shared by the converters.""" + self.logger.info("=" * 80) + self.logger.info(f"Processing {label}: {config.get('model_full_name', model_short_name)}") + self.logger.info(f"Short name: {model_short_name}") + if "description" in config: + self.logger.info(f"Description: {config['description']}") + self.logger.info("=" * 80) + + def _finalize_success( + self, + config: dict[str, Any], + model_short_name: str, + *, + accuracy: "AccuracyResults | None", + quantization_attempted: bool, + ) -> bool: + """Record a successful conversion result and log it; always returns ``True``.""" + quantized = accuracy.int8_succeeded if quantization_attempted and accuracy is not None else True + self._record_result( + self._build_result(config), + converted=True, + quantized=quantized, + accuracy=accuracy, + ) + self.logger.info(f"✓ Successfully converted {model_short_name}") + return True + + def _record_failure( + self, + config: dict[str, Any], + model_short_name: str, + error: Exception, + *, + label: str = "model", + ) -> bool: + """Log a conversion failure, record a failed result, and return ``False``.""" + import traceback + + self.logger.error(f"✗ Failed to process {label} {model_short_name}: {error}") + self._record_result(self._build_result(config), converted=False, quantized=False) + self.logger.debug(traceback.format_exc()) + return False + + def _select_accuracy_metric( + self, + config: dict[str, Any], + dataset_path: Path | None, + accuracy: "AccuracyResults", + ) -> "tuple[Metric | None, bool]": + """Resolve the metric for a config and flag whether it is top-1. + + Returns ``(metric, is_top1)``. ``metric`` is ``None`` when accuracy + measurement is disabled or no metric applies. When a metric is found its + ``name`` is copied onto ``accuracy``. + """ + from model_converter.metrics import TopOneAccuracy + + metric = self._metric_for_config(config, dataset_path) if self.measure_accuracy else None + if metric is not None: + accuracy.metric_name = metric.name + return metric, isinstance(metric, TopOneAccuracy) + + def _collect_metric_validation_samples( + self, + metric: "Metric | None", + is_top1: bool, + dataset_path: Path | None, + dataset_type: str | None, + ) -> "list[CalibrationSample] | None": + """Collect raw validation samples for Model API-based metrics, else ``None``. + + Top-1 classification and the no-metric case use the preprocessed-tensor + path and do not need raw samples. + """ + if metric is None or is_top1: + return None + return self._collect_validation_samples(dataset_path, dataset_type, subset_size=500) or None + + def _cleanup_fp32(self, fp32_model_path: Path) -> None: + """Remove the temporary FP32 ``.xml``/``.bin`` pair used for quantization.""" + try: + if fp32_model_path.exists(): + fp32_model_path.unlink() + self.logger.debug(f"Removed temporary FP32 model: {fp32_model_path}") + fp32_bin_path = fp32_model_path.with_suffix(".bin") + if fp32_bin_path.exists(): + fp32_bin_path.unlink() + self.logger.debug(f"Removed temporary FP32 weights: {fp32_bin_path}") + except OSError as e: + self.logger.warning(f"Failed to remove temporary FP32 files: {e}") + def _metric_for_config( self, config: dict[str, Any], @@ -314,6 +436,15 @@ def get_labels(self, label_set: str) -> str | None: categories = [label.replace(" ", "_") for label in categories] return " ".join(categories) + if label_set == "COCO_80": + from torchvision.models.detection import MaskRCNN_ResNet50_FPN_Weights + + from model_converter.metrics.coco_detection import COCO80_TO_COCO91 + + categories = MaskRCNN_ResNet50_FPN_Weights.COCO_V1.meta["categories"] + categories = [categories[cat_id].replace(" ", "_") for cat_id in COCO80_TO_COCO91] + return " ".join(categories) + return None def _collect_dataset_entries( @@ -642,6 +773,30 @@ def _update_metric_with_result( ) metric.update(pred_mask, gt_mask) + @staticmethod + def _build_coco_prediction( + image_id: Any, + label: Any, + bbox_xyxy: Any, + score: Any, + ) -> dict[str, Any]: + """Build one COCO-format prediction dict from an xyxy box and 80-class label. + + The 0-79 ``label`` index is mapped to the original COCO 91-class category + ID via :data:`COCO80_TO_COCO91`; out-of-range indices fall back to + ``label + 1``. The ``[x_min, y_min, x_max, y_max]`` box is converted to + COCO ``[x, y, w, h]`` format. + """ + x_min, y_min, x_max, y_max = (float(v) for v in bbox_xyxy) + n = int(label) + coco_cat_id = COCO80_TO_COCO91[n] if n < len(COCO80_TO_COCO91) else n + 1 + return { + "image_id": int(image_id), + "category_id": coco_cat_id, + "bbox": [x_min, y_min, x_max - x_min, y_max - y_min], + "score": float(score), + } + @staticmethod def _feed_bbox_predictions( metric: "CocoDetectionMAP", @@ -653,19 +808,11 @@ def _feed_bbox_predictions( scores = getattr(result, "scores", None) if bboxes is None or labels is None or scores is None: return - preds: list[dict[str, Any]] = [] - for bbox, label, score in zip(bboxes, labels, scores): - x_min, y_min, x_max, y_max = (float(v) for v in bbox) - n = int(label) - coco_cat_id = COCO80_TO_COCO91[n] if n < len(COCO80_TO_COCO91) else n + 1 - preds.append( - { - "image_id": int(sample.image_id) if sample.image_id is not None else 0, - "category_id": coco_cat_id, - "bbox": [x_min, y_min, x_max - x_min, y_max - y_min], - "score": float(score), - }, - ) + image_id = sample.image_id if sample.image_id is not None else 0 + preds = [ + BaseConverter._build_coco_prediction(image_id, label, bbox, score) + for bbox, label, score in zip(bboxes, labels, scores) + ] metric.update(predictions=preds) def quantize_model( @@ -753,8 +900,7 @@ def quantize_model( accuracy_results.int8_succeeded = True # Save model_info as config.json to track downloads - with (output_folder / "config.json").open("w") as f: - json.dump(quantized_model.get_rt_info(["model_info"]).value, f, indent=4) + self._write_config_json(output_folder, quantized_model.get_rt_info(["model_info"]).value) # Validate accuracy if validation data provided. # Two paths: (1) Top-1 classification uses preprocessed validation_data + labels via raw @@ -818,10 +964,7 @@ def quantize_model( accuracy_results.measured = True # Copy .gitattributes file - gitattributes_template = Path(__file__).parent.parent / "templates" / ".gitattributes" - if gitattributes_template.exists(): - shutil.copy2(gitattributes_template, output_folder / ".gitattributes") - self.logger.debug(f"Copied .gitattributes to: {output_folder}") + self._copy_gitattributes(output_folder) # Copy README for INT8 model self.copy_readme( @@ -842,6 +985,19 @@ def quantize_model( self.logger.debug(traceback.format_exc()) return model_path + def _copy_gitattributes(self, output_folder: Path) -> None: + """Copy the shared ``.gitattributes`` template into ``output_folder`` if present.""" + gitattributes_template = Path(__file__).parent.parent / "templates" / ".gitattributes" + if gitattributes_template.exists(): + shutil.copy2(gitattributes_template, output_folder / ".gitattributes") + self.logger.debug(f"Copied .gitattributes to: {output_folder}") + + @staticmethod + def _write_config_json(output_folder: Path, model_info: Any) -> None: + """Write ``model_info`` rt_info as ``config.json`` to track downloads.""" + with (output_folder / "config.json").open("w") as f: + json.dump(model_info, f, indent=4) + @staticmethod def _metadata_value(value: Any) -> str: """Convert config values to Model API rt_info string values.""" diff --git a/model_converter/src/model_converter/converters/getitune.py b/model_converter/src/model_converter/converters/getitune.py index 8134949a..3c492efb 100644 --- a/model_converter/src/model_converter/converters/getitune.py +++ b/model_converter/src/model_converter/converters/getitune.py @@ -5,20 +5,15 @@ """Getitune (training_extensions) model converter.""" -import json import shutil import subprocess # nosec B404 — fixed-argv invocation of `uv run`, no shell, no untrusted input import tempfile from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import Any from model_converter.converters.base import BaseConverter -from model_converter.metrics import TopOneAccuracy from model_converter.reporting import AccuracyResults -if TYPE_CHECKING: - from model_converter.datasets import CalibrationSample - class GetituneConverter(BaseConverter): """Converter for getitune models from training_extensions. @@ -51,13 +46,7 @@ def process_model_config(self, config: dict[str, Any]) -> bool: """ model_short_name = config.get("model_short_name", "unknown") - # Check if both FP16 and INT8 models already exist - fp16_model_path = self.output_dir / f"{model_short_name}-fp16-ov" / f"{model_short_name}.xml" - int8_model_path = self.output_dir / f"{model_short_name}-int8-ov" / f"{model_short_name}.xml" - - if fp16_model_path.exists() and int8_model_path.exists(): - self.logger.info(f"Skipping {model_short_name}: FP16 and INT8 models already exist") - self._record_result(self._build_result(config), converted=False, quantized=False, skipped=True) + if self._skip_if_already_converted(config, model_short_name): return True try: @@ -72,22 +61,9 @@ def process_model_config(self, config: dict[str, Any]) -> bool: error_msg = f"training_extensions directory not found: {self.training_extensions_dir}" raise FileNotFoundError(error_msg) - model_license = config.get("license") - model_license_link = config.get("license_link") - - if not model_license: - error_msg = f"Model '{model_short_name}' must define 'license' in configuration" - raise ValueError(error_msg) - if not model_license_link: - error_msg = f"Model '{model_short_name}' must define 'license_link' in configuration" - raise ValueError(error_msg) + self._validate_license(config, model_short_name) - self.logger.info("=" * 80) - self.logger.info(f"Processing getitune model: {config.get('model_full_name', model_short_name)}") - self.logger.info(f"Short name: {model_short_name}") - if "description" in config: - self.logger.info(f"Description: {config['description']}") - self.logger.info("=" * 80) + self._log_model_banner(config, model_short_name, label="getitune model") # Export model using training_extensions script exported_model_path = self._run_export(config) @@ -101,24 +77,15 @@ def process_model_config(self, config: dict[str, Any]) -> bool: if quantization_attempted: accuracy = self._quantize_exported_model(config) - quantized = accuracy.int8_succeeded if quantization_attempted and accuracy is not None else True - self._record_result( - self._build_result(config), - converted=True, - quantized=quantized, + return self._finalize_success( + config, + model_short_name, accuracy=accuracy, + quantization_attempted=quantization_attempted, ) - self.logger.info(f"✓ Successfully converted {model_short_name}") - return True - except (ValueError, RuntimeError, FileNotFoundError, OSError, subprocess.CalledProcessError) as e: - self.logger.error(f"✗ Failed to process getitune model {model_short_name}: {e}") - self._record_result(self._build_result(config), converted=False, quantized=False) - import traceback - - self.logger.debug(traceback.format_exc()) - return False + return self._record_failure(config, model_short_name, e, label="getitune model") def _run_export(self, config: dict[str, Any]) -> Path: """Run the export_pretrained_models.py script. @@ -255,15 +222,12 @@ def _repackage_model(self, config: dict[str, Any], exported_model_path: Path) -> core = ov.Core() model = core.read_model(target_xml) model_info = model.get_rt_info(["model_info"]).value - with (output_folder / "config.json").open("w") as f: - json.dump(model_info, f, indent=4) + self._write_config_json(output_folder, model_info) except (ImportError, RuntimeError, KeyError) as e: self.logger.warning(f"Could not extract model_info metadata: {e}") # Copy .gitattributes file - gitattributes_template = Path(__file__).parent.parent / "templates" / ".gitattributes" - if gitattributes_template.exists(): - shutil.copy2(gitattributes_template, output_folder / ".gitattributes") + self._copy_gitattributes(output_folder) # Copy README self.copy_readme(config, output_folder, variant="fp16") @@ -390,10 +354,7 @@ def _quantize_exported_model(self, config: dict[str, Any]) -> AccuracyResults: # (multilabel mAP, COCO mAP, mIoU) flow through Model API via # :meth:`_measure_metric` with raw image samples. dataset_path = self._resolve_dataset_path(config) - metric = self._metric_for_config(config, dataset_path) if self.measure_accuracy else None - is_top1 = isinstance(metric, TopOneAccuracy) - if metric is not None: - accuracy.metric_name = metric.name + metric, is_top1 = self._select_accuracy_metric(config, dataset_path, accuracy) self.logger.info("Creating calibration dataset for INT8 quantization") if is_top1: @@ -410,16 +371,12 @@ def _quantize_exported_model(self, config: dict[str, Any]) -> AccuracyResults: dataset_type=config.get("dataset_type"), ) - validation_samples: "list[CalibrationSample] | None" = None - if metric is not None and not is_top1: - validation_samples = ( - self._collect_validation_samples( - dataset_path, - config.get("dataset_type"), - subset_size=500, - ) - or None - ) + validation_samples = self._collect_metric_validation_samples( + metric, + is_top1, + dataset_path, + config.get("dataset_type"), + ) if calibration_data: self.quantize_model( @@ -435,16 +392,7 @@ def _quantize_exported_model(self, config: dict[str, Any]) -> AccuracyResults: ) # Clean up temporary FP32 model after quantization - try: - fp32_path = self.output_dir / f"{model_short_name}-fp16-ov" / f"{model_short_name}_fp32.xml" - if fp32_path.exists(): - fp32_path.unlink() - self.logger.debug(f"Removed temporary FP32 model: {fp32_path}") - fp32_bin_path = fp32_path.with_suffix(".bin") - if fp32_bin_path.exists(): - fp32_bin_path.unlink() - self.logger.debug(f"Removed temporary FP32 weights: {fp32_bin_path}") - except OSError as e: - self.logger.warning(f"Failed to remove temporary FP32 files: {e}") + fp32_path = self.output_dir / f"{model_short_name}-fp16-ov" / f"{model_short_name}_fp32.xml" + self._cleanup_fp32(fp32_path) return accuracy diff --git a/model_converter/src/model_converter/converters/pytorch.py b/model_converter/src/model_converter/converters/pytorch.py index 03f92b2a..74546057 100644 --- a/model_converter/src/model_converter/converters/pytorch.py +++ b/model_converter/src/model_converter/converters/pytorch.py @@ -6,22 +6,17 @@ """PyTorch-based converter shared by torchvision and timm converters.""" import importlib -import json -import shutil +from dataclasses import dataclass from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import Any import torch import torch.nn as nn from model_converter.adapters import get_adapter from model_converter.converters.base import BaseConverter -from model_converter.metrics import TopOneAccuracy from model_converter.reporting import AccuracyResults -if TYPE_CHECKING: - from model_converter.datasets import CalibrationSample - _MODEL_API_METADATA_FIELDS = ( "resize_type", "pad_value", @@ -35,6 +30,19 @@ ) +@dataclass(frozen=True) +class ExportParams: + """Export/quantization parameters resolved from a model configuration.""" + + input_shape: list[int] + input_names: list[str] + output_names: list[str] + reverse_input_channels: bool + mean_values: str + scale_values: str + model_type: str + + class PyTorchConverter(BaseConverter): """Shared converter for PyTorch-based models (torchvision, timm). @@ -202,14 +210,10 @@ def export_to_openvino( self.logger.info(f"✓ Model saved: {xml_path}") # Save model_info as config.json to track downloads - with (output_folder / "config.json").open("w") as f: - json.dump(ov_model.get_rt_info(["model_info"]).value, f, indent=4) + self._write_config_json(output_folder, ov_model.get_rt_info(["model_info"]).value) # Copy .gitattributes file - gitattributes_template = Path(__file__).parent.parent / "templates" / ".gitattributes" - if gitattributes_template.exists(): - shutil.copy2(gitattributes_template, output_folder / ".gitattributes") - self.logger.debug(f"Copied .gitattributes to: {output_folder}") + self._copy_gitattributes(output_folder) # Copy README for FP16 model self.copy_readme( @@ -307,6 +311,19 @@ def _build_metadata(self, config: dict[str, Any]) -> dict[tuple[str, str], str]: return metadata + @staticmethod + def _extract_export_params(config: dict[str, Any]) -> ExportParams: + """Resolve the shared export/quantization parameters from ``config``.""" + return ExportParams( + input_shape=config.get("input_shape", [1, 3, 224, 224]), + input_names=config.get("input_names", ["input"]), + output_names=config.get("output_names", ["result"]), + reverse_input_channels=config.get("reverse_input_channels", True), + mean_values=config.get("mean_values", "123.675 116.28 103.53"), + scale_values=config.get("scale_values", "58.395 57.12 57.375"), + model_type=config.get("model_type", ""), + ) + def validate_torch_model( self, model: nn.Module, @@ -357,10 +374,7 @@ def _quantize_and_cleanup(self, config: dict[str, Any], fp32_model_path: Path, * # multilabel/COCO/mIoU flow through Model API via :meth:`_measure_metric`. dataset_path = self._resolve_dataset_path(config) config_with_type = {**config, "model_type": model_type} - metric = self._metric_for_config(config_with_type, dataset_path) if self.measure_accuracy else None - is_top1 = isinstance(metric, TopOneAccuracy) - if metric is not None: - accuracy.metric_name = metric.name + metric, is_top1 = self._select_accuracy_metric(config_with_type, dataset_path, accuracy) resize_type = config.get("resize_type", "standard") if is_top1: @@ -377,16 +391,12 @@ def _quantize_and_cleanup(self, config: dict[str, Any], fp32_model_path: Path, * dataset_type=config.get("dataset_type"), ) - validation_samples: "list[CalibrationSample] | None" = None - if metric is not None and not is_top1: - validation_samples = ( - self._collect_validation_samples( - dataset_path, - config.get("dataset_type"), - subset_size=500, - ) - or None - ) + validation_samples = self._collect_metric_validation_samples( + metric, + is_top1, + dataset_path, + config.get("dataset_type"), + ) if validation_data: torch_model = kwargs.get("torch_model") @@ -411,15 +421,6 @@ def _quantize_and_cleanup(self, config: dict[str, Any], fp32_model_path: Path, * ) # Clean up temporary FP32 model after quantization - try: - if fp32_model_path.exists(): - fp32_model_path.unlink() - self.logger.debug(f"Removed temporary FP32 model: {fp32_model_path}") - fp32_bin_path = fp32_model_path.with_suffix(".bin") - if fp32_bin_path.exists(): - fp32_bin_path.unlink() - self.logger.debug(f"Removed temporary FP32 weights: {fp32_bin_path}") - except OSError as e: - self.logger.warning(f"Failed to remove temporary FP32 files: {e}") + self._cleanup_fp32(fp32_model_path) return accuracy diff --git a/model_converter/src/model_converter/converters/timm.py b/model_converter/src/model_converter/converters/timm.py index 0cd0e306..889ffedd 100644 --- a/model_converter/src/model_converter/converters/timm.py +++ b/model_converter/src/model_converter/converters/timm.py @@ -139,32 +139,13 @@ def process_model_config(self, config: dict[str, Any]) -> bool: """ model_short_name = config.get("model_short_name", "unknown") - # Check if both FP16 and INT8 models already exist - fp16_model_path = self.output_dir / f"{model_short_name}-fp16-ov" / f"{model_short_name}.xml" - int8_model_path = self.output_dir / f"{model_short_name}-int8-ov" / f"{model_short_name}.xml" - - if fp16_model_path.exists() and int8_model_path.exists(): - self.logger.info(f"Skipping {model_short_name}: FP16 and INT8 models already exist") - self._record_result(self._build_result(config), converted=False, quantized=False, skipped=True) + if self._skip_if_already_converted(config, model_short_name): return True try: - model_license = config.get("license") - model_license_link = config.get("license_link") - - if not model_license: - error_msg = f"Model '{model_short_name}' must define 'license' in configuration" - raise ValueError(error_msg) - if not model_license_link: - error_msg = f"Model '{model_short_name}' must define 'license_link' in configuration" - raise ValueError(error_msg) + self._validate_license(config, model_short_name) - self.logger.info("=" * 80) - self.logger.info(f"Processing model: {config.get('model_full_name', model_short_name)}") - self.logger.info(f"Short name: {model_short_name}") - if "description" in config: - self.logger.info(f"Description: {config['description']}") - self.logger.info("=" * 80) + self._log_model_banner(config, model_short_name) # Load model from HuggingFace huggingface_repo = config.get("huggingface_repo") @@ -195,24 +176,18 @@ def process_model_config(self, config: dict[str, Any]) -> bool: self._apply_timm_data_config(model, config) # Prepare export parameters - input_shape = config.get("input_shape", [1, 3, 224, 224]) - input_names = config.get("input_names", ["input"]) - output_names = config.get("output_names", ["result"]) - reverse_input_channels = config.get("reverse_input_channels", True) - mean_values = config.get("mean_values", "123.675 116.28 103.53") - scale_values = config.get("scale_values", "58.395 57.12 57.375") - model_type = config.get("model_type", "") + params = self._extract_export_params(config) metadata = self._build_metadata(config) output_path = self.output_dir / model_short_name - fp16_model_path, fp32_model_path = self.export_to_openvino( + _fp16_model_path, fp32_model_path = self.export_to_openvino( model=model, - input_shape=input_shape, + input_shape=params.input_shape, output_path=output_path, model_config=config, - input_names=input_names, - output_names=output_names, + input_names=params.input_names, + output_names=params.output_names, metadata=metadata, ) @@ -223,29 +198,20 @@ def process_model_config(self, config: dict[str, Any]) -> bool: accuracy = self._quantize_and_cleanup( config, fp32_model_path, - model_type=model_type, - input_shape=input_shape, - mean_values=mean_values, - scale_values=scale_values, - reverse_input_channels=reverse_input_channels, + model_type=params.model_type, + input_shape=params.input_shape, + mean_values=params.mean_values, + scale_values=params.scale_values, + reverse_input_channels=params.reverse_input_channels, torch_model=model, ) - quantized = accuracy.int8_succeeded if quantization_attempted and accuracy is not None else True - self._record_result( - self._build_result(config), - converted=True, - quantized=quantized, + return self._finalize_success( + config, + model_short_name, accuracy=accuracy, + quantization_attempted=quantization_attempted, ) - self.logger.info(f"✓ Successfully converted {model_short_name}") - return True - except (ValueError, RuntimeError, ImportError, FileNotFoundError) as e: - self.logger.error(f"✗ Failed to process model {model_short_name}: {e}") - self._record_result(self._build_result(config), converted=False, quantized=False) - import traceback - - self.logger.debug(traceback.format_exc()) - return False + return self._record_failure(config, model_short_name, e) diff --git a/model_converter/src/model_converter/converters/torchvision.py b/model_converter/src/model_converter/converters/torchvision.py index d351544f..4512ebcc 100644 --- a/model_converter/src/model_converter/converters/torchvision.py +++ b/model_converter/src/model_converter/converters/torchvision.py @@ -37,32 +37,13 @@ def process_model_config(self, config: dict[str, Any]) -> bool: """ model_short_name = config.get("model_short_name", "unknown") - # Check if both FP16 and INT8 models already exist - fp16_model_path = self.output_dir / f"{model_short_name}-fp16-ov" / f"{model_short_name}.xml" - int8_model_path = self.output_dir / f"{model_short_name}-int8-ov" / f"{model_short_name}.xml" - - if fp16_model_path.exists() and int8_model_path.exists(): - self.logger.info(f"Skipping {model_short_name}: FP16 and INT8 models already exist") - self._record_result(self._build_result(config), converted=False, quantized=False, skipped=True) + if self._skip_if_already_converted(config, model_short_name): return True try: - model_license = config.get("license") - model_license_link = config.get("license_link") - - if not model_license: - error_msg = f"Model '{model_short_name}' must define 'license' in configuration" - raise ValueError(error_msg) - if not model_license_link: - error_msg = f"Model '{model_short_name}' must define 'license_link' in configuration" - raise ValueError(error_msg) - - self.logger.info("=" * 80) - self.logger.info(f"Processing model: {config.get('model_full_name', model_short_name)}") - self.logger.info(f"Short name: {model_short_name}") - if "description" in config: - self.logger.info(f"Description: {config['description']}") - self.logger.info("=" * 80) + self._validate_license(config, model_short_name) + + self._log_model_banner(config, model_short_name) # Download weights and load model weights_url = config["weights_url"] @@ -77,24 +58,18 @@ def process_model_config(self, config: dict[str, Any]) -> bool: model = self.create_model(model_class, checkpoint, model_params) # Prepare export parameters - input_shape = config.get("input_shape", [1, 3, 224, 224]) - input_names = config.get("input_names", ["input"]) - output_names = config.get("output_names", ["result"]) - reverse_input_channels = config.get("reverse_input_channels", True) - mean_values = config.get("mean_values", "123.675 116.28 103.53") - scale_values = config.get("scale_values", "58.395 57.12 57.375") - model_type = config.get("model_type", "") + params = self._extract_export_params(config) metadata = self._build_metadata(config) output_path = self.output_dir / model_short_name - fp16_model_path, fp32_model_path = self.export_to_openvino( + _fp16_model_path, fp32_model_path = self.export_to_openvino( model=model, - input_shape=input_shape, + input_shape=params.input_shape, output_path=output_path, model_config=config, - input_names=input_names, - output_names=output_names, + input_names=params.input_names, + output_names=params.output_names, metadata=metadata, ) @@ -105,29 +80,20 @@ def process_model_config(self, config: dict[str, Any]) -> bool: accuracy = self._quantize_and_cleanup( config, fp32_model_path, - model_type=model_type, - input_shape=input_shape, - mean_values=mean_values, - scale_values=scale_values, - reverse_input_channels=reverse_input_channels, + model_type=params.model_type, + input_shape=params.input_shape, + mean_values=params.mean_values, + scale_values=params.scale_values, + reverse_input_channels=params.reverse_input_channels, torch_model=model, ) - quantized = accuracy.int8_succeeded if quantization_attempted and accuracy is not None else True - self._record_result( - self._build_result(config), - converted=True, - quantized=quantized, + return self._finalize_success( + config, + model_short_name, accuracy=accuracy, + quantization_attempted=quantization_attempted, ) - self.logger.info(f"✓ Successfully converted {model_short_name}") - return True - except (ValueError, RuntimeError, ImportError, FileNotFoundError) as e: - self.logger.error(f"✗ Failed to process model {model_short_name}: {e}") - self._record_result(self._build_result(config), converted=False, quantized=False) - import traceback - - self.logger.debug(traceback.format_exc()) - return False + return self._record_failure(config, model_short_name, e) diff --git a/model_converter/src/model_converter/converters/yolo.py b/model_converter/src/model_converter/converters/yolo.py index 1b801006..fabc531d 100644 --- a/model_converter/src/model_converter/converters/yolo.py +++ b/model_converter/src/model_converter/converters/yolo.py @@ -50,9 +50,7 @@ def process_model_config(self, config: dict[str, Any]) -> bool: fp16_folder = self.output_dir / f"{model_short_name}-fp16-ov" int8_folder = self.output_dir / f"{model_short_name}-int8-ov" - if (fp16_folder / f"{yolo_version}.xml").exists() and (int8_folder / f"{yolo_version}.xml").exists(): - self.logger.info(f"Skipping {model_short_name}: FP16 and INT8 models already exist") - self._record_result(self._build_result(config), converted=False, quantized=False, skipped=True) + if self._skip_if_already_converted(config, model_short_name, xml_stem=yolo_version): return True try: @@ -116,12 +114,7 @@ def process_model_config(self, config: dict[str, Any]) -> bool: return True except (ValueError, RuntimeError, ImportError, FileNotFoundError, OSError) as e: - self.logger.error(f"✗ Failed to process YOLO model {model_short_name}: {e}") - self._record_result(self._build_result(config), converted=False, quantized=False) - import traceback - - self.logger.debug(traceback.format_exc()) - return False + return self._record_failure(config, model_short_name, e, label="YOLO model") def _measure_yolo_accuracy( self, @@ -245,15 +238,8 @@ class indices to COCO 91-class category IDs via :data:`COCO80_TO_COCO91`, scores = results.boxes.conf.cpu().tolist() for (x1, y1, x2, y2), cls_idx, score in zip(boxes_xyxy, cls_ids, scores): - n = int(cls_idx) - coco_cat_id = _COCO80_TO_COCO91[n] if n < len(_COCO80_TO_COCO91) else n + 1 predictions.append( - { - "image_id": int(sample.image_id), - "category_id": coco_cat_id, - "bbox": [x1, y1, x2 - x1, y2 - y1], - "score": float(score), - }, + self._build_coco_prediction(sample.image_id, cls_idx, (x1, y1, x2, y2), score), ) metric.update(predictions=predictions) diff --git a/model_converter/src/model_converter/metrics/coco_detection.py b/model_converter/src/model_converter/metrics/coco_detection.py index e2226fe6..cdf9bea9 100644 --- a/model_converter/src/model_converter/metrics/coco_detection.py +++ b/model_converter/src/model_converter/metrics/coco_detection.py @@ -103,6 +103,14 @@ 90, ] +# Inverse of :data:`COCO80_TO_COCO91`: maps an original COCO 91-class category ID +# to the contiguous 80-class index (0-79) used by OpenVINO detection models. +# Category IDs that COCO dropped (and the background ID 0) map to ``None``. +COCO91_TO_COCO80: list[int | None] = [None] * 91 +for _idx, _cat_id in enumerate(COCO80_TO_COCO91): + COCO91_TO_COCO80[_cat_id] = _idx +del _idx, _cat_id + class CocoDetectionMAP(Metric): """Wraps :class:`pycocotools.cocoeval.COCOeval` for a single ``iouType``.""" diff --git a/model_converter/tests/unit/test_adapters.py b/model_converter/tests/unit/test_adapters.py index 35a45970..2f6f43a5 100644 --- a/model_converter/tests/unit/test_adapters.py +++ b/model_converter/tests/unit/test_adapters.py @@ -97,8 +97,17 @@ def test_forward(self): assert labels.shape == (10,) assert masks.shape == (10, 28, 28) # squeezed from (10, 1, 28, 28) - # Labels should be shifted by -1 - expected_labels = mock_predictions[0]["labels"] - 1 + # Labels should be remapped from COCO 91-class IDs to 80-class indices, + # then offset by -1 to compensate for the Model API MaskRCNN wrapper's +1. + from model_converter.metrics.coco_detection import COCO91_TO_COCO80 + + expected_labels = torch.tensor( + [ + COCO91_TO_COCO80[int(cat_id)] - 1 if COCO91_TO_COCO80[int(cat_id)] is not None else -1 + for cat_id in mock_predictions[0]["labels"] + ], + dtype=torch.int64, + ) assert torch.equal(labels, expected_labels) def test_forward_with_tensor_features(self): diff --git a/model_converter/tests/unit/test_cli.py b/model_converter/tests/unit/test_cli.py index 02d9555d..19647e82 100644 --- a/model_converter/tests/unit/test_cli.py +++ b/model_converter/tests/unit/test_cli.py @@ -152,6 +152,21 @@ def test_coco_v1(self, converter): assert result == "person bicycle car" + def test_coco_80(self, converter): + """get_labels returns the 80 contiguous COCO classes for COCO_80.""" + from model_converter.metrics.coco_detection import COCO80_TO_COCO91 + + categories = [f"cat {i}" for i in range(91)] + mock_weights = MagicMock() + mock_weights.COCO_V1.meta = {"categories": categories} + + with patch("torchvision.models.detection.MaskRCNN_ResNet50_FPN_Weights", mock_weights): + result = converter.get_labels("COCO_80") + + expected = " ".join(f"cat_{cat_id}" for cat_id in COCO80_TO_COCO91) + assert result == expected + assert len(result.split()) == 80 + def test_unknown_label_set(self, converter): """get_labels returns None for unknown label sets.""" assert converter.get_labels("NONEXISTENT_LABELS") is None @@ -483,7 +498,7 @@ def test_successful_export(self, converter, sample_model_config): patch("openvino.convert_model", return_value=mock_ov_model), patch("openvino.save_model") as mock_save, patch.object(Path, "exists", return_value=True), - patch("model_converter.converters.pytorch.shutil.copy2"), + patch("model_converter.converters.base.shutil.copy2"), patch.object(converter, "copy_readme"), ): fp16_path, fp32_path = converter.export_to_openvino( diff --git a/model_converter/tests/unit/test_registry.py b/model_converter/tests/unit/test_registry.py index a5e477b4..478075df 100644 --- a/model_converter/tests/unit/test_registry.py +++ b/model_converter/tests/unit/test_registry.py @@ -456,7 +456,7 @@ def test_export_to_openvino_converts_and_saves_models( patch.object(converter, "_prepare_model_for_export", return_value=mock_torch_model) as mock_prepare, patch.object(converter, "_create_example_input", return_value=dummy_input) as mock_example_input, patch.object(converter, "_postprocess_openvino_model", return_value=mock_ov_model) as mock_postprocess, - patch("model_converter.converters.pytorch.shutil.copy2") as mock_copy, + patch("model_converter.converters.base.shutil.copy2") as mock_copy, patch.object(converter, "copy_readme") as mock_copy_readme, ): fp16_path, fp32_path = converter.export_to_openvino(