feat: improved model_exporter

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,15 +417,15 @@
       "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,
       "agnostic_nms": false,
       "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"
     },
     {

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

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."""