diff --git a/model_converter/README.md b/model_converter/README.md index 0aa67cefd..99faf96c3 100644 --- a/model_converter/README.md +++ b/model_converter/README.md @@ -14,6 +14,7 @@ This tool reads a JSON configuration file containing model specifications, downl - **Input/Output Naming**: Configurable input and output tensor names - **Batch Processing**: Process multiple models from a single configuration file - **Selective Conversion**: Convert specific models using the `--model` flag +- **Summary Report**: Generate a console + Markdown conversion report with the `--report` flag ## Installation @@ -30,7 +31,7 @@ uv sync ### Basic Usage ```bash -uv run model-converter examples/config.json -o ./output_models +uv run model-converter presets/config.json -o ./output_models --datasets-config datasets.json ``` ### Command-Line Options @@ -45,7 +46,13 @@ options: Output directory for converted models (default: ./converted_models) -c CACHE, --cache CACHE Cache directory for downloaded weights (default: ~/.cache/torch/hub/checkpoints) + --datasets-config DATASETS_CONFIG + Path to datasets configuration JSON file (default: presets/datasets.json) --model MODEL Process only the specified model (by model_short_name) + --library LIBRARY Comma-separated list of model libraries to process (e.g., getitune,timm) + --report [PATH] Generate a conversion summary report. Pass the flag alone to write to + /conversion_report.md, or provide a PATH to override the location. + The report is printed to the console and saved as Markdown. --list List all models in the configuration file and exit -v, --verbose Enable verbose logging ``` @@ -82,6 +89,65 @@ uv run model-converter examples/config.json -o ./output -c ./my_cache uv run model-converter examples/config.json -o ./output -v ``` +**Generate a conversion summary report:** + +```bash +# Write the report to /conversion_report.md +uv run model-converter examples/config.json -o ./output --report + +# Write the report to a custom path +uv run model-converter examples/config.json -o ./output --report ./reports/summary.md +``` + +## Conversion Summary Report + +Passing `--report` produces a summary of every processed model. The report is +printed to the console **and** saved as a Markdown file (default +`/conversion_report.md`, or the path given to `--report`). + +The report contains one row per model with the following columns: + +| Column | Description | +| ----------------- | ------------------------------------------------------------------------------------------------- | +| Model Full Name | Human-readable model name (`model_full_name`). | +| Model Type | Task/architecture type (e.g. `Classification`, `SSD`). | +| Model Library | Source library (`torchvision`, `timm`, `yolo`, `getitune`). | +| Original URL | Exact download URL — `weights_url` or the Hugging Face repository URL; `N/A` when not applicable. | +| Original Accuracy | Top-1 accuracy of the original PyTorch model (before OpenVINO conversion), or `N/A`. | +| FP32 Accuracy | Top-1 accuracy of the FP32 model, or `N/A` when not measured. | +| FP16 Accuracy | Top-1 accuracy of the FP16 model, or `N/A` when not measured. | +| INT8 Accuracy | Top-1 accuracy of the quantized INT8 model, or `N/A` when not measured. | +| Status | Outcome of the conversion (see below). | + +Accuracy is measured only for classification models that define `labels`, over the +calibration subset. The original accuracy is computed by running the source PyTorch +model on the same preprocessed validation images, so it is directly comparable to the +FP32/FP16/INT8 numbers. Models without measurable accuracy report `N/A` and a status of +`OK (no accuracy data)`. + +### Status values + +| Status | Meaning | +| ----------------------- | ---------------------------------------------------------------------------------- | +| `OK` | Converted, accuracy measured, all drops within 5%. | +| `OK (no accuracy data)` | Converted (FP16 + INT8) but no accuracy could be measured. | +| `ACCURACY DROP >5%` | Original→FP32, FP16, or INT8 top-1 accuracy dropped more than 5 percentage points. | +| `FAILED: conversion` | Model load or export failed (no FP16 produced). | +| `FAILED: quantization` | FP16 produced but the INT8 model was not produced. | +| `SKIPPED` | FP16 and INT8 models already existed, so processing was skipped. | + +### Sample report + +```markdown +# Conversion Summary Report + +| Model Full Name | Model Type | Model Library | Original URL | Original Accuracy | FP32 Accuracy | FP16 Accuracy | INT8 Accuracy | Status | +| --------------- | -------------- | ------------- | ------------------------------------------------ | ----------------- | ------------- | ------------- | ------------- | --------------------- | +| ResNet-50 | Classification | torchvision | https://download.pytorch.org/models/resnet50.pth | 80.20% | 80.12% | 80.10% | 79.85% | OK | +| EfficientNet-B0 | Classification | timm | https://huggingface.co/timm/efficientnet_b0 | 77.72% | 77.70% | 77.65% | 70.10% | ACCURACY DROP >5% | +| YOLO11n | YOLO11 | yolo | N/A | N/A | N/A | N/A | N/A | OK (no accuracy data) | +``` + ## Configuration File Format The configuration file is a JSON file with the following structure: @@ -101,7 +167,8 @@ The configuration file is a JSON file with the following structure: "input_names": ["images"], "output_names": ["output"], "model_params": null, - "model_type": "Classification" + "model_type": "Classification", + "dataset_type": "imagenet-1k" } ] } @@ -116,6 +183,36 @@ Common `model_type` values: - `"YOLOX"` - YOLOX detection models - `"SegmentationModel"` - Segmentation models +### Dataset Configuration + +The converter uses a separate `datasets.json` file to map dataset types to local filesystem paths. This allows model configurations to remain portable across different environments. + +**datasets.json format:** + +```json +{ + "datasets": { + "imagenet-1k": "/path/to/imagenet/validation", + "imagenet-21k": "/path/to/imagenet21k/validation", + "coco-detection": "/path/to/coco2017/val2017", + "coco-segmentation": "/path/to/coco2017/val2017" + } +} +``` + +Models that require calibration for INT8 quantization should specify a `dataset_type` field that matches one of the keys in `datasets.json`: + +```json +{ + "model_short_name": "efficientnet_b0", + "model_type": "Classification", + "dataset_type": "imagenet-1k", + ... +} +``` + +If a model does not specify a `dataset_type`, INT8 quantization will be skipped for that model. + ### Configuration Fields #### Required Fields @@ -141,3 +238,4 @@ For Hugging Face-backed models, use these required fields instead of `model_clas - **`output_names`** (array of strings): Names for output tensors (default: auto-generated) - **`model_params`** (object): Parameters to pass to model constructor (default: `null`) - **`model_type`** (string): Model type for model_api auto-detection (e.g., `"Classification"`, `"DetectionModel"`, `"YOLOX"`, etc.) +- **`dataset_type`** (string): Dataset type identifier that maps to a path in `datasets.json` (e.g., `"imagenet-1k"`, `"coco-detection"`). Required for INT8 quantization. If omitted, quantization is skipped for this model. diff --git a/model_converter/examples/config.json b/model_converter/examples/config.json deleted file mode 100644 index e63d6f8d2..000000000 --- a/model_converter/examples/config.json +++ /dev/null @@ -1,652 +0,0 @@ -{ - "models": [ - { - "model_short_name": "efficientnet_b0", - "model_class_name": "torchvision.models.efficientnet.efficientnet_b0", - "model_library": "torchvision", - "model_full_name": "EfficientNet-B0", - "description": "EfficientNet-B0 - Efficient convolutional neural network with compound scaling", - "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b0.html#torchvision.models.efficientnet_b0", - "weights_url": "https://download.pytorch.org/models/efficientnet_b0_rwightman-3dd342df.pth", - "input_shape": [1, 3, 224, 224], - "input_names": ["image"], - "output_names": ["logits"], - "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", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "efficientnet_b1", - "model_class_name": "torchvision.models.efficientnet.efficientnet_b1", - "model_library": "torchvision", - "model_full_name": "EfficientNet-B1", - "description": "EfficientNet-B1 - Efficient convolutional neural network with compound scaling", - "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b1.html", - "weights_url": "https://download.pytorch.org/models/efficientnet_b1_rwightman-bac287d4.pth", - "input_shape": [1, 3, 240, 240], - "input_names": ["image"], - "output_names": ["logits"], - "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", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "efficientnet_b2", - "model_class_name": "torchvision.models.efficientnet.efficientnet_b2", - "model_library": "torchvision", - "model_full_name": "EfficientNet-B2", - "description": "EfficientNet-B2 - Efficient convolutional neural network with compound scaling", - "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b2.html", - "weights_url": "https://download.pytorch.org/models/efficientnet_b2_rwightman-c35c1473.pth", - "input_shape": [1, 3, 260, 260], - "input_names": ["image"], - "output_names": ["logits"], - "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", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "efficientnet_b3", - "model_class_name": "torchvision.models.efficientnet.efficientnet_b3", - "model_library": "torchvision", - "model_full_name": "EfficientNet-B3", - "description": "EfficientNet-B3 - Efficient convolutional neural network with compound scaling", - "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b3.html", - "weights_url": "https://download.pytorch.org/models/efficientnet_b3_rwightman-b3899882.pth", - "input_shape": [1, 3, 300, 300], - "input_names": ["image"], - "output_names": ["logits"], - "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", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "efficientnet_b4", - "model_class_name": "torchvision.models.efficientnet.efficientnet_b4", - "model_library": "torchvision", - "model_full_name": "EfficientNet-B4", - "description": "EfficientNet-B4 - Efficient convolutional neural network with compound scaling", - "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b4.html", - "weights_url": "https://download.pytorch.org/models/efficientnet_b4_rwightman-23ab8bcd.pth", - "input_shape": [1, 3, 380, 380], - "input_names": ["image"], - "output_names": ["logits"], - "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", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "efficientnet_b5", - "model_class_name": "torchvision.models.efficientnet.efficientnet_b5", - "model_library": "torchvision", - "model_full_name": "EfficientNet-B5", - "description": "EfficientNet-B5 - Efficient convolutional neural network with compound scaling", - "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b5.html", - "weights_url": "https://download.pytorch.org/models/efficientnet_b5_lukemelas-1a07897c.pth", - "input_shape": [1, 3, 456, 456], - "input_names": ["image"], - "output_names": ["logits"], - "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", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "efficientnet_b6", - "model_class_name": "torchvision.models.efficientnet.efficientnet_b6", - "model_library": "torchvision", - "model_full_name": "EfficientNet-B6", - "description": "EfficientNet-B6 - Efficient convolutional neural network with compound scaling", - "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b6.html", - "weights_url": "https://download.pytorch.org/models/efficientnet_b6_lukemelas-24a108a5.pth", - "input_shape": [1, 3, 528, 528], - "input_names": ["image"], - "output_names": ["logits"], - "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", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "efficientnet_b7", - "model_class_name": "torchvision.models.efficientnet.efficientnet_b7", - "model_library": "torchvision", - "model_full_name": "EfficientNet-B7", - "description": "EfficientNet-B7 - Efficient convolutional neural network with compound scaling", - "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b7.html", - "weights_url": "https://download.pytorch.org/models/efficientnet_b7_lukemelas-c5b4e57e.pth", - "input_shape": [1, 3, 600, 600], - "input_names": ["image"], - "output_names": ["logits"], - "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", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "efficientnet_v2_s", - "model_class_name": "torchvision.models.efficientnet.efficientnet_v2_s", - "model_library": "torchvision", - "model_full_name": "EfficientNetV2-S", - "description": "EfficientNetV2-Small - Improved EfficientNet with faster training and better parameter efficiency", - "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_v2_s.html", - "weights_url": "https://download.pytorch.org/models/efficientnet_v2_s-dd5fe13b.pth", - "input_shape": [1, 3, 384, 384], - "input_names": ["image"], - "output_names": ["logits"], - "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", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "efficientnet_v2_m", - "model_class_name": "torchvision.models.efficientnet.efficientnet_v2_m", - "model_library": "torchvision", - "model_full_name": "EfficientNetV2-M", - "description": "EfficientNetV2-Medium - Improved EfficientNet with faster training and better parameter efficiency", - "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_v2_m.html", - "weights_url": "https://download.pytorch.org/models/efficientnet_v2_m-dc08266a.pth", - "input_shape": [1, 3, 480, 480], - "input_names": ["image"], - "output_names": ["logits"], - "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", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "efficientnet_v2_l", - "model_class_name": "torchvision.models.efficientnet.efficientnet_v2_l", - "model_library": "torchvision", - "model_full_name": "EfficientNetV2-L", - "description": "EfficientNetV2-Large - Improved EfficientNet with faster training and better parameter efficiency", - "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_v2_l.html", - "weights_url": "https://download.pytorch.org/models/efficientnet_v2_l-59c71312.pth", - "input_shape": [1, 3, 480, 480], - "input_names": ["image"], - "output_names": ["logits"], - "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", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "mobilenet_v2", - "model_class_name": "torchvision.models.mobilenetv2.mobilenet_v2", - "model_library": "torchvision", - "model_full_name": "MobileNetV2", - "description": "MobileNetV2 - Efficient convolutional neural network for mobile and embedded vision applications", - "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.mobilenet_v2.html", - "weights_url": "https://download.pytorch.org/models/mobilenet_v2-b0353104.pth", - "input_shape": [1, 3, 224, 224], - "input_names": ["image"], - "output_names": ["logits"], - "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", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "mobilenet_v3_small", - "model_class_name": "torchvision.models.mobilenetv3.mobilenet_v3_small", - "model_library": "torchvision", - "model_full_name": "MobileNetV3-Small", - "description": "MobileNetV3 Small - Efficient convolutional neural network for mobile and embedded vision applications", - "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.mobilenet_v3_small.html#torchvision.models.mobilenet_v3_small", - "weights_url": "https://download.pytorch.org/models/mobilenet_v3_small-047dcff4.pth", - "input_shape": [1, 3, 224, 224], - "input_names": ["image"], - "output_names": ["output1"], - "model_params": null, - "model_type": "Classification", - "reverse_input_channels": false, - "mean_values": "123.675 116.28 103.53", - "scale_values": "58.395 57.12 57.375", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "mobilenet_v3_large", - "model_class_name": "torchvision.models.mobilenetv3.mobilenet_v3_large", - "model_library": "torchvision", - "model_full_name": "MobileNetV3-Large", - "description": "MobileNetV3 Large - Efficient convolutional neural network for mobile and embedded vision applications", - "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.mobilenet_v3_large.html", - "weights_url": "https://download.pytorch.org/models/mobilenet_v3_large-8738ca79.pth", - "input_shape": [1, 3, 224, 224], - "input_names": ["image"], - "output_names": ["logits"], - "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", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "resnet18", - "model_class_name": "torchvision.models.resnet.resnet18", - "model_library": "torchvision", - "model_full_name": "ResNet-18", - "description": "ResNet-18 - 18-layer residual learning network for image classification", - "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.resnet18.html", - "weights_url": "https://download.pytorch.org/models/resnet18-f37072fd.pth", - "input_shape": [1, 3, 224, 224], - "input_names": ["image"], - "output_names": ["output"], - "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", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "resnet34", - "model_class_name": "torchvision.models.resnet.resnet34", - "model_library": "torchvision", - "model_full_name": "ResNet-34", - "description": "ResNet-34 - 34-layer residual learning network for image classification", - "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.resnet34.html", - "weights_url": "https://download.pytorch.org/models/resnet34-b627a593.pth", - "input_shape": [1, 3, 224, 224], - "input_names": ["image"], - "output_names": ["output"], - "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", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "resnet50", - "model_class_name": "torchvision.models.resnet.resnet50", - "model_library": "torchvision", - "model_full_name": "ResNet-50", - "description": "ResNet-50 - 50-layer residual learning network for image classification", - "weights_url": "https://download.pytorch.org/models/resnet50-0676ba61.pth", - "input_shape": [1, 3, 224, 224], - "input_names": ["image"], - "output_names": ["output"], - "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", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "resnet101", - "model_class_name": "torchvision.models.resnet.resnet101", - "model_library": "torchvision", - "model_full_name": "ResNet-101", - "description": "ResNet-101 - 101-layer residual learning network for image classification", - "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.resnet101.html", - "weights_url": "https://download.pytorch.org/models/resnet101-63fe2227.pth", - "input_shape": [1, 3, 224, 224], - "input_names": ["image"], - "output_names": ["output"], - "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", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "resnet152", - "model_class_name": "torchvision.models.resnet.resnet152", - "model_library": "torchvision", - "model_full_name": "ResNet-152", - "description": "ResNet-152 - 152-layer residual learning network for image classification", - "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.resnet152.html", - "weights_url": "https://download.pytorch.org/models/resnet152-394f9c45.pth", - "input_shape": [1, 3, 224, 224], - "input_names": ["image"], - "output_names": ["output"], - "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", - "license": "bsd-3-clause", - "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", - "labels": "IMAGENET1K_V1" - }, - - { - "model_short_name": "mobilenetv2_050_lamb_in1k", - "huggingface_repo": "timm/mobilenetv2_050.lamb_in1k", - "huggingface_revision": "8990230b131c231ad38b217b23cf9d2e7e70e8d0", - "model_library": "timm", - "model_full_name": "MobileNetV2-0.5 LAMB ImageNet-1k", - "description": "MobileNetV2 with 0.5x width multiplier trained with LAMB optimizer on ImageNet-1k", - "docs": "https://huggingface.co/timm/mobilenetv2_050.lamb_in1k", - "input_shape": [1, 3, 224, 224], - "input_names": ["image"], - "output_names": ["logits"], - "model_params": null, - "model_type": "Classification", - "reverse_input_channels": false, - "mean_values": "123.675 116.28 103.53", - "scale_values": "58.395 57.12 57.375", - "license": "apache-2.0", - "license_link": "https://spdx.org/licenses/Apache-2.0.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "mobilenetv2_100_ra_in1k", - "huggingface_repo": "timm/mobilenetv2_100.ra_in1k", - "huggingface_revision": "5afa12513b048c79b9147a5d210e9bdc50035481", - "model_library": "timm", - "model_full_name": "MobileNetV2-1.0 RandAugment ImageNet-1k", - "description": "MobileNetV2 with 1.0x width multiplier trained with RandAugment on ImageNet-1k", - "docs": "https://huggingface.co/timm/mobilenetv2_100.ra_in1k", - "input_shape": [1, 3, 224, 224], - "input_names": ["image"], - "output_names": ["logits"], - "model_params": null, - "model_type": "Classification", - "reverse_input_channels": false, - "mean_values": "123.675 116.28 103.53", - "scale_values": "58.395 57.12 57.375", - "license": "apache-2.0", - "license_link": "https://spdx.org/licenses/Apache-2.0.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "mobilenetv2_110d_ra_in1k", - "huggingface_repo": "timm/mobilenetv2_110d.ra_in1k", - "huggingface_revision": "143925e3675a13f37d417cde6e73a39a04d81dcd", - "model_library": "timm", - "model_full_name": "MobileNetV2-1.1 RandAugment ImageNet-1k", - "description": "MobileNetV2 with 1.1x width multiplier trained with RandAugment on ImageNet-1k", - "docs": "https://huggingface.co/timm/mobilenetv2_110d.ra_in1k", - "input_shape": [1, 3, 224, 224], - "input_names": ["image"], - "output_names": ["logits"], - "model_params": null, - "model_type": "Classification", - "reverse_input_channels": false, - "mean_values": "123.675 116.28 103.53", - "scale_values": "58.395 57.12 57.375", - "license": "apache-2.0", - "license_link": "https://spdx.org/licenses/Apache-2.0.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "mobilenetv2_120d_ra_in1k", - "huggingface_repo": "timm/mobilenetv2_120d.ra_in1k", - "huggingface_revision": "3e70fb071e6567f136c48aabf1b45eb071b5d471", - "model_library": "timm", - "model_full_name": "MobileNetV2-1.2 RandAugment ImageNet-1k", - "description": "MobileNetV2 with 1.2x width multiplier trained with RandAugment on ImageNet-1k", - "docs": "https://huggingface.co/timm/mobilenetv2_120d.ra_in1k", - "input_shape": [1, 3, 224, 224], - "input_names": ["image"], - "output_names": ["logits"], - "model_params": null, - "model_type": "Classification", - "reverse_input_channels": false, - "mean_values": "123.675 116.28 103.53", - "scale_values": "58.395 57.12 57.375", - "license": "apache-2.0", - "license_link": "https://spdx.org/licenses/Apache-2.0.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "mobilenetv2_140_ra_in1k", - "huggingface_repo": "timm/mobilenetv2_140.ra_in1k", - "huggingface_revision": "008fa2eb59c0080738c65898a7666059e11c0a75", - "model_library": "timm", - "model_full_name": "MobileNetV2-1.4 RandAugment ImageNet-1k", - "description": "MobileNetV2 with 1.4x width multiplier trained with RandAugment on ImageNet-1k", - "docs": "https://huggingface.co/timm/mobilenetv2_140.ra_in1k", - "input_shape": [1, 3, 224, 224], - "input_names": ["image"], - "output_names": ["logits"], - "model_params": null, - "model_type": "Classification", - "reverse_input_channels": false, - "mean_values": "123.675 116.28 103.53", - "scale_values": "58.395 57.12 57.375", - "license": "apache-2.0", - "license_link": "https://spdx.org/licenses/Apache-2.0.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "efficientnet_b0_ra_in1k", - "huggingface_repo": "timm/efficientnet_b0.ra_in1k", - "huggingface_revision": "1b5383e5f79cc0f7fc067e372f8f26a5fa73f26a", - "model_library": "timm", - "model_full_name": "EfficientNet-B0 RandAugment ImageNet-1k", - "description": "EfficientNet-B0 trained with RandAugment recipe on ImageNet-1k", - "docs": "https://huggingface.co/timm/efficientnet_b0.ra_in1k", - "input_shape": [1, 3, 224, 224], - "input_names": ["image"], - "output_names": ["logits"], - "model_params": null, - "model_type": "Classification", - "reverse_input_channels": false, - "mean_values": "123.675 116.28 103.53", - "scale_values": "58.395 57.12 57.375", - "license": "apache-2.0", - "license_link": "https://spdx.org/licenses/Apache-2.0.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "efficientnet_b0_ra4_e3600_r224_in1k", - "huggingface_repo": "timm/efficientnet_b0.ra4_e3600_r224_in1k", - "huggingface_revision": "1efd3beedccfd00f9de56a5875041a709f3058e2", - "model_library": "timm", - "model_full_name": "EfficientNet-B0 RA4 3600 Epochs ImageNet-1k", - "description": "EfficientNet-B0 trained with RA4 recipe for 3600 epochs on ImageNet-1k", - "docs": "https://huggingface.co/timm/efficientnet_b0.ra4_e3600_r224_in1k", - "input_shape": [1, 3, 224, 224], - "input_names": ["image"], - "output_names": ["logits"], - "model_params": null, - "model_type": "Classification", - "reverse_input_channels": false, - "mean_values": "123.675 116.28 103.53", - "scale_values": "58.395 57.12 57.375", - "license": "apache-2.0", - "license_link": "https://spdx.org/licenses/Apache-2.0.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "efficientnet_b1_ft_in1k", - "huggingface_repo": "timm/efficientnet_b1.ft_in1k", - "huggingface_revision": "1d6ddfd0ad535646fdb05bc3834913d93816152e", - "model_library": "timm", - "model_full_name": "EfficientNet-B1 Fine-tuned ImageNet-1k", - "description": "EfficientNet-B1 fine-tuned on ImageNet-1k", - "docs": "https://huggingface.co/timm/efficientnet_b1.ft_in1k", - "input_shape": [1, 3, 240, 240], - "input_names": ["image"], - "output_names": ["logits"], - "model_params": null, - "model_type": "Classification", - "reverse_input_channels": false, - "mean_values": "123.675 116.28 103.53", - "scale_values": "58.395 57.12 57.375", - "license": "apache-2.0", - "license_link": "https://spdx.org/licenses/Apache-2.0.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "efficientnet_b1_pruned_in1k", - "huggingface_repo": "timm/efficientnet_b1_pruned.in1k", - "huggingface_revision": "fe6043b912b98d991a070f3070b924ee799694f5", - "model_library": "timm", - "model_full_name": "EfficientNet-B1 Pruned ImageNet-1k", - "description": "EfficientNet-B1 with pruned weights on ImageNet-1k", - "docs": "https://huggingface.co/timm/efficientnet_b1_pruned.in1k", - "input_shape": [1, 3, 240, 240], - "input_names": ["image"], - "output_names": ["logits"], - "model_params": null, - "model_type": "Classification", - "reverse_input_channels": false, - "mean_values": "123.675 116.28 103.53", - "scale_values": "58.395 57.12 57.375", - "license": "apache-2.0", - "license_link": "https://spdx.org/licenses/Apache-2.0.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "efficientnet_b1_ra4_e3600_r240_in1k", - "huggingface_repo": "timm/efficientnet_b1.ra4_e3600_r240_in1k", - "huggingface_revision": "4e4a9c08eee4bdb118d69d27f87ace8fdf3fc743", - "model_library": "timm", - "model_full_name": "EfficientNet-B1 RA4 3600 Epochs ImageNet-1k", - "description": "EfficientNet-B1 trained with RA4 recipe for 3600 epochs at 240x240 resolution on ImageNet-1k", - "docs": "https://huggingface.co/timm/efficientnet_b1.ra4_e3600_r240_in1k", - "input_shape": [1, 3, 240, 240], - "input_names": ["image"], - "output_names": ["logits"], - "model_params": null, - "model_type": "Classification", - "reverse_input_channels": false, - "mean_values": "123.675 116.28 103.53", - "scale_values": "58.395 57.12 57.375", - "license": "apache-2.0", - "license_link": "https://spdx.org/licenses/Apache-2.0.html", - "labels": "IMAGENET1K_V1" - }, - { - "model_short_name": "vit_tiny_patch16_224_augreg_in21k", - "huggingface_repo": "timm/vit_tiny_patch16_224.augreg_in21k", - "huggingface_revision": "3d5f75e2fe58abe541d5651356278a1df3fd3ab3", - "model_library": "timm", - "model_full_name": "ViT-Tiny Patch16 224 AugReg ImageNet-21k", - "description": "Vision Transformer Tiny with 16x16 patches trained on ImageNet-21k with augmentation and regularization", - "docs": "https://huggingface.co/timm/vit_tiny_patch16_224.augreg_in21k", - "input_shape": [1, 3, 224, 224], - "input_names": ["image"], - "output_names": ["logits"], - "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", - "license": "apache-2.0", - "license_link": "https://spdx.org/licenses/Apache-2.0.html", - "labels": "IMAGENET21K" - }, - { - "model_short_name": "vit_small_patch14_dinov2.lvd142m", - "huggingface_repo": "timm/vit_small_patch14_dinov2.lvd142m", - "huggingface_revision": "4610ca143709d58a633b6397a74412c2c3842454", - "model_library": "timm", - "model_full_name": "DINOv2-Small Patch14 518 LVD-142M", - "description": "DINOv2 Small ViT backbone for image feature extraction with 384-dimensional features", - "docs": "https://huggingface.co/timm/vit_small_patch14_dinov2.lvd142m", - "input_shape": [1, 3, 518, 518], - "input_names": ["image"], - "output_names": ["output"], - "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", - "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" - } - ] -} diff --git a/model_converter/presets/config.json b/model_converter/presets/config.json new file mode 100644 index 000000000..624ef7235 --- /dev/null +++ b/model_converter/presets/config.json @@ -0,0 +1,1348 @@ +{ + "models": [ + { + "model_short_name": "efficientnet_b0", + "model_class_name": "torchvision.models.efficientnet.efficientnet_b0", + "model_library": "torchvision", + "model_full_name": "EfficientNet-B0", + "description": "EfficientNet-B0 - Efficient convolutional neural network with compound scaling", + "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b0.html#torchvision.models.efficientnet_b0", + "weights_url": "https://download.pytorch.org/models/efficientnet_b0_rwightman-3dd342df.pth", + "input_shape": [1, 3, 224, 224], + "input_names": ["image"], + "output_names": ["logits"], + "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", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_b1", + "model_class_name": "torchvision.models.efficientnet.efficientnet_b1", + "model_library": "torchvision", + "model_full_name": "EfficientNet-B1", + "description": "EfficientNet-B1 - Efficient convolutional neural network with compound scaling", + "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b1.html", + "weights_url": "https://download.pytorch.org/models/efficientnet_b1_rwightman-bac287d4.pth", + "input_shape": [1, 3, 240, 240], + "input_names": ["image"], + "output_names": ["logits"], + "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", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_b2", + "model_class_name": "torchvision.models.efficientnet.efficientnet_b2", + "model_library": "torchvision", + "model_full_name": "EfficientNet-B2", + "description": "EfficientNet-B2 - Efficient convolutional neural network with compound scaling", + "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b2.html", + "weights_url": "https://download.pytorch.org/models/efficientnet_b2_rwightman-c35c1473.pth", + "input_shape": [1, 3, 260, 260], + "input_names": ["image"], + "output_names": ["logits"], + "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", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_b3", + "model_class_name": "torchvision.models.efficientnet.efficientnet_b3", + "model_library": "torchvision", + "model_full_name": "EfficientNet-B3", + "description": "EfficientNet-B3 - Efficient convolutional neural network with compound scaling", + "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b3.html", + "weights_url": "https://download.pytorch.org/models/efficientnet_b3_rwightman-b3899882.pth", + "input_shape": [1, 3, 300, 300], + "input_names": ["image"], + "output_names": ["logits"], + "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", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_b4", + "model_class_name": "torchvision.models.efficientnet.efficientnet_b4", + "model_library": "torchvision", + "model_full_name": "EfficientNet-B4", + "description": "EfficientNet-B4 - Efficient convolutional neural network with compound scaling", + "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b4.html", + "weights_url": "https://download.pytorch.org/models/efficientnet_b4_rwightman-23ab8bcd.pth", + "input_shape": [1, 3, 380, 380], + "input_names": ["image"], + "output_names": ["logits"], + "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", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_b5", + "model_class_name": "torchvision.models.efficientnet.efficientnet_b5", + "model_library": "torchvision", + "model_full_name": "EfficientNet-B5", + "description": "EfficientNet-B5 - Efficient convolutional neural network with compound scaling", + "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b5.html", + "weights_url": "https://download.pytorch.org/models/efficientnet_b5_lukemelas-1a07897c.pth", + "input_shape": [1, 3, 456, 456], + "input_names": ["image"], + "output_names": ["logits"], + "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", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_b6", + "model_class_name": "torchvision.models.efficientnet.efficientnet_b6", + "model_library": "torchvision", + "model_full_name": "EfficientNet-B6", + "description": "EfficientNet-B6 - Efficient convolutional neural network with compound scaling", + "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b6.html", + "weights_url": "https://download.pytorch.org/models/efficientnet_b6_lukemelas-24a108a5.pth", + "input_shape": [1, 3, 528, 528], + "input_names": ["image"], + "output_names": ["logits"], + "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", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_b7", + "model_class_name": "torchvision.models.efficientnet.efficientnet_b7", + "model_library": "torchvision", + "model_full_name": "EfficientNet-B7", + "description": "EfficientNet-B7 - Efficient convolutional neural network with compound scaling", + "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b7.html", + "weights_url": "https://download.pytorch.org/models/efficientnet_b7_lukemelas-c5b4e57e.pth", + "input_shape": [1, 3, 600, 600], + "input_names": ["image"], + "output_names": ["logits"], + "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", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_v2_s", + "model_class_name": "torchvision.models.efficientnet.efficientnet_v2_s", + "model_library": "torchvision", + "model_full_name": "EfficientNetV2-S", + "description": "EfficientNetV2-Small - Improved EfficientNet with faster training and better parameter efficiency", + "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_v2_s.html", + "weights_url": "https://download.pytorch.org/models/efficientnet_v2_s-dd5fe13b.pth", + "input_shape": [1, 3, 384, 384], + "input_names": ["image"], + "output_names": ["logits"], + "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", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_v2_m", + "model_class_name": "torchvision.models.efficientnet.efficientnet_v2_m", + "model_library": "torchvision", + "model_full_name": "EfficientNetV2-M", + "description": "EfficientNetV2-Medium - Improved EfficientNet with faster training and better parameter efficiency", + "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_v2_m.html", + "weights_url": "https://download.pytorch.org/models/efficientnet_v2_m-dc08266a.pth", + "input_shape": [1, 3, 480, 480], + "input_names": ["image"], + "output_names": ["logits"], + "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", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_v2_l", + "model_class_name": "torchvision.models.efficientnet.efficientnet_v2_l", + "model_library": "torchvision", + "model_full_name": "EfficientNetV2-L", + "description": "EfficientNetV2-Large - Improved EfficientNet with faster training and better parameter efficiency", + "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_v2_l.html", + "weights_url": "https://download.pytorch.org/models/efficientnet_v2_l-59c71312.pth", + "input_shape": [1, 3, 480, 480], + "input_names": ["image"], + "output_names": ["logits"], + "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", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "mobilenet_v2", + "model_class_name": "torchvision.models.mobilenetv2.mobilenet_v2", + "model_library": "torchvision", + "model_full_name": "MobileNetV2", + "description": "MobileNetV2 - Efficient convolutional neural network for mobile and embedded vision applications", + "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.mobilenet_v2.html", + "weights_url": "https://download.pytorch.org/models/mobilenet_v2-b0353104.pth", + "input_shape": [1, 3, 224, 224], + "input_names": ["image"], + "output_names": ["logits"], + "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", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "mobilenet_v3_small", + "model_class_name": "torchvision.models.mobilenetv3.mobilenet_v3_small", + "model_library": "torchvision", + "model_full_name": "MobileNetV3-Small", + "description": "MobileNetV3 Small - Efficient convolutional neural network for mobile and embedded vision applications", + "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.mobilenet_v3_small.html#torchvision.models.mobilenet_v3_small", + "weights_url": "https://download.pytorch.org/models/mobilenet_v3_small-047dcff4.pth", + "input_shape": [1, 3, 224, 224], + "input_names": ["image"], + "output_names": ["output1"], + "model_params": null, + "model_type": "Classification", + "reverse_input_channels": false, + "mean_values": "123.675 116.28 103.53", + "scale_values": "58.395 57.12 57.375", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "mobilenet_v3_large", + "model_class_name": "torchvision.models.mobilenetv3.mobilenet_v3_large", + "model_library": "torchvision", + "model_full_name": "MobileNetV3-Large", + "description": "MobileNetV3 Large - Efficient convolutional neural network for mobile and embedded vision applications", + "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.mobilenet_v3_large.html", + "weights_url": "https://download.pytorch.org/models/mobilenet_v3_large-8738ca79.pth", + "input_shape": [1, 3, 224, 224], + "input_names": ["image"], + "output_names": ["logits"], + "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", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "resnet18", + "model_class_name": "torchvision.models.resnet.resnet18", + "model_library": "torchvision", + "model_full_name": "ResNet-18", + "description": "ResNet-18 - 18-layer residual learning network for image classification", + "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.resnet18.html", + "weights_url": "https://download.pytorch.org/models/resnet18-f37072fd.pth", + "input_shape": [1, 3, 224, 224], + "input_names": ["image"], + "output_names": ["output"], + "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", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "resnet34", + "model_class_name": "torchvision.models.resnet.resnet34", + "model_library": "torchvision", + "model_full_name": "ResNet-34", + "description": "ResNet-34 - 34-layer residual learning network for image classification", + "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.resnet34.html", + "weights_url": "https://download.pytorch.org/models/resnet34-b627a593.pth", + "input_shape": [1, 3, 224, 224], + "input_names": ["image"], + "output_names": ["output"], + "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", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "resnet50", + "model_class_name": "torchvision.models.resnet.resnet50", + "model_library": "torchvision", + "model_full_name": "ResNet-50", + "description": "ResNet-50 - 50-layer residual learning network for image classification", + "weights_url": "https://download.pytorch.org/models/resnet50-0676ba61.pth", + "input_shape": [1, 3, 224, 224], + "input_names": ["image"], + "output_names": ["output"], + "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", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "resnet101", + "model_class_name": "torchvision.models.resnet.resnet101", + "model_library": "torchvision", + "model_full_name": "ResNet-101", + "description": "ResNet-101 - 101-layer residual learning network for image classification", + "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.resnet101.html", + "weights_url": "https://download.pytorch.org/models/resnet101-63fe2227.pth", + "input_shape": [1, 3, 224, 224], + "input_names": ["image"], + "output_names": ["output"], + "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", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "resnet152", + "model_class_name": "torchvision.models.resnet.resnet152", + "model_library": "torchvision", + "model_full_name": "ResNet-152", + "description": "ResNet-152 - 152-layer residual learning network for image classification", + "docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.resnet152.html", + "weights_url": "https://download.pytorch.org/models/resnet152-394f9c45.pth", + "input_shape": [1, 3, 224, 224], + "input_names": ["image"], + "output_names": ["output"], + "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", + "license": "bsd-3-clause", + "license_link": "https://spdx.org/licenses/BSD-3-Clause.html", + "labels": "IMAGENET1K_V1", + "dataset_type": "imagenet-1k" + }, + { + "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", + "dataset_type": "coco-detection" + }, + { + "model_short_name": "mobilenetv2_050_lamb_in1k", + "huggingface_repo": "timm/mobilenetv2_050.lamb_in1k", + "huggingface_revision": "8990230b131c231ad38b217b23cf9d2e7e70e8d0", + "model_library": "timm", + "model_full_name": "MobileNetV2-0.5 LAMB ImageNet-1k", + "description": "MobileNetV2 with 0.5x width multiplier trained with LAMB optimizer on ImageNet-1k", + "docs": "https://huggingface.co/timm/mobilenetv2_050.lamb_in1k", + "input_names": ["image"], + "output_names": ["logits"], + "model_params": null, + "model_type": "Classification", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "labels": "IMAGENET1K_V1", + "resize_type": "crop", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "mobilenetv2_100_ra_in1k", + "huggingface_repo": "timm/mobilenetv2_100.ra_in1k", + "huggingface_revision": "5afa12513b048c79b9147a5d210e9bdc50035481", + "model_library": "timm", + "model_full_name": "MobileNetV2-1.0 RandAugment ImageNet-1k", + "description": "MobileNetV2 with 1.0x width multiplier trained with RandAugment on ImageNet-1k", + "docs": "https://huggingface.co/timm/mobilenetv2_100.ra_in1k", + "input_names": ["image"], + "output_names": ["logits"], + "model_params": null, + "model_type": "Classification", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "labels": "IMAGENET1K_V1", + "resize_type": "crop", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "mobilenetv2_110d_ra_in1k", + "huggingface_repo": "timm/mobilenetv2_110d.ra_in1k", + "huggingface_revision": "143925e3675a13f37d417cde6e73a39a04d81dcd", + "model_library": "timm", + "model_full_name": "MobileNetV2-1.1 RandAugment ImageNet-1k", + "description": "MobileNetV2 with 1.1x width multiplier trained with RandAugment on ImageNet-1k", + "docs": "https://huggingface.co/timm/mobilenetv2_110d.ra_in1k", + "input_names": ["image"], + "output_names": ["logits"], + "model_params": null, + "model_type": "Classification", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "labels": "IMAGENET1K_V1", + "resize_type": "crop", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "mobilenetv2_120d_ra_in1k", + "huggingface_repo": "timm/mobilenetv2_120d.ra_in1k", + "huggingface_revision": "3e70fb071e6567f136c48aabf1b45eb071b5d471", + "model_library": "timm", + "model_full_name": "MobileNetV2-1.2 RandAugment ImageNet-1k", + "description": "MobileNetV2 with 1.2x width multiplier trained with RandAugment on ImageNet-1k", + "docs": "https://huggingface.co/timm/mobilenetv2_120d.ra_in1k", + "input_names": ["image"], + "output_names": ["logits"], + "model_params": null, + "model_type": "Classification", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "labels": "IMAGENET1K_V1", + "resize_type": "crop", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "mobilenetv2_140_ra_in1k", + "huggingface_repo": "timm/mobilenetv2_140.ra_in1k", + "huggingface_revision": "008fa2eb59c0080738c65898a7666059e11c0a75", + "model_library": "timm", + "model_full_name": "MobileNetV2-1.4 RandAugment ImageNet-1k", + "description": "MobileNetV2 with 1.4x width multiplier trained with RandAugment on ImageNet-1k", + "docs": "https://huggingface.co/timm/mobilenetv2_140.ra_in1k", + "input_names": ["image"], + "output_names": ["logits"], + "model_params": null, + "model_type": "Classification", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "labels": "IMAGENET1K_V1", + "resize_type": "crop", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_b0_ra_in1k", + "huggingface_repo": "timm/efficientnet_b0.ra_in1k", + "huggingface_revision": "1b5383e5f79cc0f7fc067e372f8f26a5fa73f26a", + "model_library": "timm", + "model_full_name": "EfficientNet-B0 RandAugment ImageNet-1k", + "description": "EfficientNet-B0 trained with RandAugment recipe on ImageNet-1k", + "docs": "https://huggingface.co/timm/efficientnet_b0.ra_in1k", + "input_names": ["image"], + "output_names": ["logits"], + "model_params": null, + "model_type": "Classification", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "labels": "IMAGENET1K_V1", + "resize_type": "crop", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_b0_ra4_e3600_r224_in1k", + "huggingface_repo": "timm/efficientnet_b0.ra4_e3600_r224_in1k", + "huggingface_revision": "1efd3beedccfd00f9de56a5875041a709f3058e2", + "model_library": "timm", + "model_full_name": "EfficientNet-B0 RA4 3600 Epochs ImageNet-1k", + "description": "EfficientNet-B0 trained with RA4 recipe for 3600 epochs on ImageNet-1k", + "docs": "https://huggingface.co/timm/efficientnet_b0.ra4_e3600_r224_in1k", + "input_names": ["image"], + "output_names": ["logits"], + "model_params": null, + "model_type": "Classification", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "labels": "IMAGENET1K_V1", + "resize_type": "crop", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_b1_ft_in1k", + "huggingface_repo": "timm/efficientnet_b1.ft_in1k", + "huggingface_revision": "1d6ddfd0ad535646fdb05bc3834913d93816152e", + "model_library": "timm", + "model_full_name": "EfficientNet-B1 Fine-tuned ImageNet-1k", + "description": "EfficientNet-B1 fine-tuned on ImageNet-1k", + "docs": "https://huggingface.co/timm/efficientnet_b1.ft_in1k", + "input_names": ["image"], + "output_names": ["logits"], + "model_params": null, + "model_type": "Classification", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "labels": "IMAGENET1K_V1", + "resize_type": "crop", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_b1_pruned_in1k", + "huggingface_repo": "timm/efficientnet_b1_pruned.in1k", + "huggingface_revision": "fe6043b912b98d991a070f3070b924ee799694f5", + "model_library": "timm", + "model_full_name": "EfficientNet-B1 Pruned ImageNet-1k", + "description": "EfficientNet-B1 with pruned weights on ImageNet-1k", + "docs": "https://huggingface.co/timm/efficientnet_b1_pruned.in1k", + "input_names": ["image"], + "output_names": ["logits"], + "model_params": null, + "model_type": "Classification", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "labels": "IMAGENET1K_V1", + "resize_type": "crop", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_b1_ra4_e3600_r240_in1k", + "huggingface_repo": "timm/efficientnet_b1.ra4_e3600_r240_in1k", + "huggingface_revision": "4e4a9c08eee4bdb118d69d27f87ace8fdf3fc743", + "model_library": "timm", + "model_full_name": "EfficientNet-B1 RA4 3600 Epochs ImageNet-1k", + "description": "EfficientNet-B1 trained with RA4 recipe for 3600 epochs at 240x240 resolution on ImageNet-1k", + "docs": "https://huggingface.co/timm/efficientnet_b1.ra4_e3600_r240_in1k", + "input_names": ["image"], + "output_names": ["logits"], + "model_params": null, + "model_type": "Classification", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "labels": "IMAGENET1K_V1", + "resize_type": "crop", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "vit_tiny_patch16_224_augreg_in21k", + "huggingface_repo": "timm/vit_tiny_patch16_224.augreg_in21k", + "huggingface_revision": "3d5f75e2fe58abe541d5651356278a1df3fd3ab3", + "model_library": "timm", + "model_full_name": "ViT-Tiny Patch16 224 AugReg ImageNet-21k", + "description": "Vision Transformer Tiny with 16x16 patches trained on ImageNet-21k with augmentation and regularization", + "docs": "https://huggingface.co/timm/vit_tiny_patch16_224.augreg_in21k", + "input_names": ["image"], + "output_names": ["logits"], + "model_params": null, + "model_type": "Classification", + "quantization_model_type": "transformer", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "labels": "IMAGENET21K", + "resize_type": "crop", + "dataset_type": "imagenet-21k" + }, + { + "model_short_name": "vit_small_patch14_dinov2.lvd142m", + "huggingface_repo": "timm/vit_small_patch14_dinov2.lvd142m", + "huggingface_revision": "4610ca143709d58a633b6397a74412c2c3842454", + "model_library": "timm", + "model_full_name": "DINOv2-Small Patch14 518 LVD-142M", + "description": "DINOv2 Small ViT backbone for image feature extraction with 384-dimensional features", + "docs": "https://huggingface.co/timm/vit_small_patch14_dinov2.lvd142m", + "input_names": ["image"], + "output_names": ["output"], + "model_params": null, + "model_type": "Classification", + "quantization_model_type": "transformer", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "resize_type": "crop", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "dino_v2", + "model_library": "getitune", + "labels": "IMAGENET1K_V1", + "model_full_name": "DINOv2 Classification", + "description": "DINOv2 - Self-supervised Vision Transformer for image classification", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_task": "MULTI_CLASS_CLS", + "getitune_recipe": "dino_v2", + "model_type": "Classification", + "tags": ["image-classification", "vision"], + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_b0_cls", + "model_library": "getitune", + "labels": "IMAGENET1K_V1", + "model_full_name": "EfficientNet-B0 Classification", + "description": "EfficientNet-B0 trained with getitune for multi-class classification", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_task": "MULTI_CLASS_CLS", + "getitune_recipe": "efficientnet_b0", + "model_type": "Classification", + "tags": ["image-classification", "vision"], + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_v2_cls", + "model_library": "getitune", + "labels": "IMAGENET21K", + "model_full_name": "EfficientNet-V2 Classification", + "description": "EfficientNet-V2 trained with getitune for multi-class classification", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_task": "MULTI_CLASS_CLS", + "getitune_recipe": "efficientnet_v2", + "model_type": "Classification", + "tags": ["image-classification", "vision"], + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "dataset_type": "imagenet-21k" + }, + { + "model_short_name": "mobilenet_v3_large_cls", + "model_library": "getitune", + "labels": "IMAGENET1K_V1", + "model_full_name": "MobileNet-V3 Large Classification", + "description": "MobileNet-V3 Large trained with getitune for multi-class classification", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_task": "MULTI_CLASS_CLS", + "getitune_recipe": "mobilenet_v3_large", + "model_type": "Classification", + "tags": ["image-classification", "vision"], + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_b3_cls", + "model_library": "getitune", + "labels": "IMAGENET1K_V1", + "model_full_name": "EfficientNet-B3 Classification", + "description": "EfficientNet-B3 trained with getitune for multi-class classification", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "efficientnet_b3", + "getitune_task": "MULTI_CLASS_CLS", + "model_type": "Classification", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["image-classification", "vision"], + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "vit_tiny_cls", + "model_library": "getitune", + "labels": "IMAGENET1K_V1", + "model_full_name": "ViT-Tiny Classification", + "description": "Vision Transformer Tiny trained with getitune for multi-class classification", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "vit_tiny", + "getitune_task": "MULTI_CLASS_CLS", + "model_type": "Classification", + "quantization_model_type": "transformer", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["image-classification", "vision"], + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_b0_multilabel", + "model_library": "getitune", + "model_full_name": "EfficientNet-B0 Multi-Label Classification", + "description": "EfficientNet-B0 trained with getitune for multi-label classification", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_task": "MULTI_LABEL_CLS", + "getitune_recipe": "efficientnet_b0", + "model_type": "Classification", + "tags": ["image-classification", "multi-label-classification", "vision"], + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_v2_multilabel", + "model_library": "getitune", + "model_full_name": "EfficientNet-V2 Multi-Label Classification", + "description": "EfficientNet-V2 trained with getitune for multi-label classification", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_task": "MULTI_LABEL_CLS", + "getitune_recipe": "efficientnet_v2", + "model_type": "Classification", + "tags": ["image-classification", "multi-label-classification", "vision"], + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "dino_v2_multilabel", + "model_library": "getitune", + "model_full_name": "DINOv2 Multi-Label Classification", + "description": "DINOv2 Vision Transformer trained with getitune for multi-label classification", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "dino_v2", + "getitune_task": "MULTI_LABEL_CLS", + "model_type": "Classification", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["image-classification", "multi-label", "vision"], + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "efficientnet_b3_multilabel", + "model_library": "getitune", + "model_full_name": "EfficientNet-B3 Multi-Label Classification", + "description": "EfficientNet-B3 trained with getitune for multi-label classification", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "efficientnet_b3", + "getitune_task": "MULTI_LABEL_CLS", + "model_type": "Classification", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["image-classification", "multi-label", "vision"], + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "mobilenet_v3_large_multilabel", + "model_library": "getitune", + "model_full_name": "MobileNet-V3-Large Multi-Label Classification", + "description": "MobileNet-V3-Large trained with getitune for multi-label classification", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "mobilenet_v3_large", + "getitune_task": "MULTI_LABEL_CLS", + "model_type": "Classification", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["image-classification", "multi-label", "vision"], + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "vit_tiny_multilabel", + "model_library": "getitune", + "model_full_name": "ViT-Tiny Multi-Label Classification", + "description": "Vision Transformer Tiny trained with getitune for multi-label classification", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "vit_tiny", + "getitune_task": "MULTI_LABEL_CLS", + "model_type": "Classification", + "quantization_model_type": "transformer", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["image-classification", "multi-label", "vision"], + "dataset_type": "imagenet-1k" + }, + { + "model_short_name": "yolox_s", + "model_library": "getitune", + "model_full_name": "YOLOX-S", + "description": "YOLOX-S single-stage object detection model trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_task": "DETECTION", + "getitune_recipe": "yolox_s", + "model_type": "YOLOX", + "tags": ["object-detection", "vision"], + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "dataset_type": "coco-detection" + }, + { + "model_short_name": "yolox_l", + "model_library": "getitune", + "model_full_name": "YOLOX-L", + "description": "YOLOX-L single-stage object detection model trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_task": "DETECTION", + "getitune_recipe": "yolox_l", + "model_type": "YOLOX", + "tags": ["object-detection", "vision"], + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "dataset_type": "coco-detection" + }, + { + "model_short_name": "yolox_x", + "model_library": "getitune", + "model_full_name": "YOLOX-X", + "description": "YOLOX-X single-stage object detection model trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_task": "DETECTION", + "getitune_recipe": "yolox_x", + "model_type": "YOLOX", + "tags": ["object-detection", "vision"], + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "dataset_type": "coco-detection" + }, + { + "model_short_name": "ssd_mobilenet_v2", + "model_library": "getitune", + "model_full_name": "SSD MobileNet-V2", + "description": "SSD with MobileNet-V2 backbone for object detection trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_task": "DETECTION", + "getitune_recipe": "ssd_mobilenetv2", + "model_type": "SSD", + "tags": ["object-detection", "vision"], + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "dataset_type": "coco-detection" + }, + { + "model_short_name": "atss_mobilenet_v2", + "model_library": "getitune", + "model_full_name": "ATSS MobileNet-V2", + "description": "ATSS with MobileNet-V2 backbone for object detection trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_task": "DETECTION", + "getitune_recipe": "atss_mobilenetv2", + "model_type": "SSD", + "tags": ["object-detection", "vision"], + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "dataset_type": "coco-detection" + }, + { + "model_short_name": "atss_r50_fpn", + "model_library": "getitune", + "model_full_name": "ATSS ResNet-50 FPN", + "description": "ATSS with ResNet-50 FPN backbone for object detection trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_task": "DETECTION", + "getitune_recipe": "atss_r50_fpn", + "model_type": "SSD", + "tags": ["object-detection", "vision"], + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "dataset_type": "coco-detection" + }, + { + "model_short_name": "yolox_tiny", + "model_library": "getitune", + "model_full_name": "YOLOX-Tiny", + "description": "YOLOX-Tiny single-stage object detection model trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "yolox_tiny", + "getitune_task": "DETECTION", + "model_type": "YOLOX", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["object-detection", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "deim_dfine_l", + "model_library": "getitune", + "model_full_name": "DEIM D-FINE-L", + "description": "DEIM D-FINE Large transformer-based object detection model trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "deim_dfine_l", + "getitune_task": "DETECTION", + "model_type": "SSD", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["object-detection", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "deim_dfine_m", + "model_library": "getitune", + "model_full_name": "DEIM D-FINE-M", + "description": "DEIM D-FINE Medium transformer-based object detection model trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "deim_dfine_m", + "getitune_task": "DETECTION", + "model_type": "SSD", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["object-detection", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "deim_dfine_x", + "model_library": "getitune", + "model_full_name": "DEIM D-FINE-X", + "description": "DEIM D-FINE Extra-Large transformer-based object detection model trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "deim_dfine_x", + "getitune_task": "DETECTION", + "model_type": "SSD", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["object-detection", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "deimv2_l", + "model_library": "getitune", + "model_full_name": "DEIMv2-L", + "description": "DEIMv2 Large transformer-based object detection model trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "deimv2_l", + "getitune_task": "DETECTION", + "model_type": "SSD", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["object-detection", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "deimv2_m", + "model_library": "getitune", + "model_full_name": "DEIMv2-M", + "description": "DEIMv2 Medium transformer-based object detection model trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "deimv2_m", + "getitune_task": "DETECTION", + "model_type": "SSD", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["object-detection", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "deimv2_s", + "model_library": "getitune", + "model_full_name": "DEIMv2-S", + "description": "DEIMv2 Small transformer-based object detection model trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "deimv2_s", + "getitune_task": "DETECTION", + "model_type": "SSD", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["object-detection", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "dfine_x", + "model_library": "getitune", + "model_full_name": "D-FINE-X", + "description": "D-FINE Extra-Large transformer-based object detection model trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "dfine_x", + "getitune_task": "DETECTION", + "model_type": "SSD", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["object-detection", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "rfdetr_large", + "model_library": "getitune", + "model_full_name": "RF-DETR Large", + "description": "RF-DETR Large transformer-based object detection model trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "rfdetr_large", + "getitune_task": "DETECTION", + "model_type": "SSD", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["object-detection", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "rfdetr_medium", + "model_library": "getitune", + "model_full_name": "RF-DETR Medium", + "description": "RF-DETR Medium transformer-based object detection model trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "rfdetr_medium", + "getitune_task": "DETECTION", + "model_type": "SSD", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["object-detection", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "rfdetr_small", + "model_library": "getitune", + "model_full_name": "RF-DETR Small", + "description": "RF-DETR Small transformer-based object detection model trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "rfdetr_small", + "getitune_task": "DETECTION", + "model_type": "SSD", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["object-detection", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "rtdetr_50", + "model_library": "getitune", + "model_full_name": "RT-DETR ResNet-50", + "description": "RT-DETR with ResNet-50 backbone for real-time object detection trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "rtdetr_50", + "getitune_task": "DETECTION", + "model_type": "SSD", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["object-detection", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "maskrcnn_r50", + "model_library": "getitune", + "model_full_name": "Mask R-CNN ResNet-50", + "description": "Mask R-CNN with ResNet-50 for instance segmentation trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_task": "INSTANCE_SEGMENTATION", + "getitune_recipe": "maskrcnn_r50", + "model_type": "MaskRCNN", + "tags": ["image-segmentation", "instance-segmentation", "vision"], + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "dataset_type": "coco-detection" + }, + { + "model_short_name": "maskrcnn_efficientnet_b2b", + "model_library": "getitune", + "model_full_name": "Mask R-CNN EfficientNet-B2B", + "description": "Mask R-CNN with EfficientNet-B2B backbone for instance segmentation trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_task": "INSTANCE_SEGMENTATION", + "getitune_recipe": "maskrcnn_efficientnetb2b", + "model_type": "MaskRCNN", + "tags": ["image-segmentation", "instance-segmentation", "vision"], + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "dataset_type": "coco-detection" + }, + { + "model_short_name": "maskrcnn_swint", + "model_library": "getitune", + "model_full_name": "Mask R-CNN Swin-Tiny", + "description": "Mask R-CNN with Swin-Tiny backbone for instance segmentation trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "maskrcnn_swint", + "getitune_task": "INSTANCE_SEGMENTATION", + "model_type": "MaskRCNN", + "quantization_model_type": "transformer", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["image-segmentation", "instance-segmentation", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "rfdetr_seg_large", + "model_library": "getitune", + "model_full_name": "RF-DETR Segmentation Large", + "description": "RF-DETR Large for instance segmentation trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "rfdetr_seg_large", + "getitune_task": "INSTANCE_SEGMENTATION", + "model_type": "MaskRCNN", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["image-segmentation", "instance-segmentation", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "rfdetr_seg_medium", + "model_library": "getitune", + "model_full_name": "RF-DETR Segmentation Medium", + "description": "RF-DETR Medium for instance segmentation trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "rfdetr_seg_medium", + "getitune_task": "INSTANCE_SEGMENTATION", + "model_type": "MaskRCNN", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["image-segmentation", "instance-segmentation", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "rfdetr_seg_small", + "model_library": "getitune", + "model_full_name": "RF-DETR Segmentation Small", + "description": "RF-DETR Small for instance segmentation trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "rfdetr_seg_small", + "getitune_task": "INSTANCE_SEGMENTATION", + "model_type": "MaskRCNN", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["image-segmentation", "instance-segmentation", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "rfdetr_seg_xlarge", + "model_library": "getitune", + "model_full_name": "RF-DETR Segmentation X-Large", + "description": "RF-DETR Extra-Large for instance segmentation trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "rfdetr_seg_xlarge", + "getitune_task": "INSTANCE_SEGMENTATION", + "model_type": "MaskRCNN", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["image-segmentation", "instance-segmentation", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "rtmdet_inst_tiny", + "model_library": "getitune", + "model_full_name": "RTMDet Instance Tiny", + "description": "RTMDet Tiny for instance segmentation trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "rtmdet_inst_tiny", + "getitune_task": "INSTANCE_SEGMENTATION", + "model_type": "MaskRCNN", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["image-segmentation", "instance-segmentation", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "litehrnet_18_semseg", + "model_library": "getitune", + "model_full_name": "Lite-HRNet-18 Semantic Segmentation", + "description": "Lite-HRNet-18 for semantic segmentation trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_task": "SEMANTIC_SEGMENTATION", + "getitune_recipe": "litehrnet_18", + "model_type": "Segmentation", + "tags": ["image-segmentation", "semantic-segmentation", "vision"], + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "dataset_type": "ade20k" + }, + { + "model_short_name": "segnext_b_semseg", + "model_library": "getitune", + "model_full_name": "SegNeXt-B Semantic Segmentation", + "description": "SegNeXt-B for semantic segmentation trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_task": "SEMANTIC_SEGMENTATION", + "getitune_recipe": "segnext_b", + "model_type": "Segmentation", + "tags": ["image-segmentation", "semantic-segmentation", "vision"], + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "dataset_type": "ade20k" + }, + { + "model_short_name": "segnext_s_semseg", + "model_library": "getitune", + "model_full_name": "SegNeXt-S Semantic Segmentation", + "description": "SegNeXt-S for semantic segmentation trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_task": "SEMANTIC_SEGMENTATION", + "getitune_recipe": "segnext_s", + "model_type": "Segmentation", + "tags": ["image-segmentation", "semantic-segmentation", "vision"], + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "dataset_type": "ade20k" + }, + { + "model_short_name": "dino_v2_semseg", + "model_library": "getitune", + "model_full_name": "DINOv2 Semantic Segmentation", + "description": "DINOv2 Vision Transformer for semantic segmentation trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "dino_v2", + "getitune_task": "SEMANTIC_SEGMENTATION", + "model_type": "Segmentation", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["image-segmentation", "semantic-segmentation", "vision"], + "dataset_type": "ade20k" + }, + { + "model_short_name": "segnext_t_semseg", + "model_library": "getitune", + "model_full_name": "SegNeXt-T Semantic Segmentation", + "description": "SegNeXt-T for semantic segmentation trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "segnext_t", + "getitune_task": "SEMANTIC_SEGMENTATION", + "model_type": "Segmentation", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["image-segmentation", "semantic-segmentation", "vision"], + "dataset_type": "ade20k" + }, + { + "model_short_name": "maskrcnn_r50_rotated", + "model_library": "getitune", + "model_full_name": "Mask R-CNN ResNet-50 Rotated Detection", + "description": "Mask R-CNN with ResNet-50 for rotated object detection trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_task": "ROTATED_DETECTION", + "getitune_recipe": "maskrcnn_r50", + "model_type": "MaskRCNN", + "tags": ["object-detection", "rotated-detection", "vision"], + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "dataset_type": "coco-detection" + }, + { + "model_short_name": "maskrcnn_efficientnetb2b_rotated", + "model_library": "getitune", + "model_full_name": "Mask R-CNN EfficientNet-B2B Rotated Detection", + "description": "Mask R-CNN with EfficientNet-B2B backbone for rotated object detection trained with getitune", + "docs": "https://github.com/open-edge-platform/training_extensions", + "getitune_recipe": "maskrcnn_efficientnetb2b", + "getitune_task": "ROTATED_DETECTION", + "model_type": "MaskRCNN", + "license": "apache-2.0", + "license_link": "https://spdx.org/licenses/Apache-2.0.html", + "tags": ["object-detection", "rotated-detection", "vision"], + "dataset_type": "coco-detection" + }, + { + "model_short_name": "YOLO11n", + "model_library": "yolo", + "model_full_name": "YOLO11n", + "description": "YOLO11 Nano - ultrafast object detection model by Ultralytics", + "docs": "https://docs.ultralytics.com/models/yolo11/", + "yolo_version": "yolo11n", + "model_type": "YOLO11", + "license": "agpl-3.0", + "license_link": "https://spdx.org/licenses/AGPL-3.0-only.html", + "dataset_type": "coco-detection" + }, + { + "model_short_name": "YOLO11s", + "model_library": "yolo", + "model_full_name": "YOLO11s", + "description": "YOLO11 Small - fast object detection model by Ultralytics", + "docs": "https://docs.ultralytics.com/models/yolo11/", + "yolo_version": "yolo11s", + "model_type": "YOLO11", + "license": "agpl-3.0", + "license_link": "https://spdx.org/licenses/AGPL-3.0-only.html", + "dataset_type": "coco-detection" + }, + { + "model_short_name": "YOLO11m", + "model_library": "yolo", + "model_full_name": "YOLO11m", + "description": "YOLO11 Medium - balanced object detection model by Ultralytics", + "docs": "https://docs.ultralytics.com/models/yolo11/", + "yolo_version": "yolo11m", + "model_type": "YOLO11", + "license": "agpl-3.0", + "license_link": "https://spdx.org/licenses/AGPL-3.0-only.html", + "dataset_type": "coco-detection" + }, + { + "model_short_name": "YOLO11l", + "model_library": "yolo", + "model_full_name": "YOLO11l", + "description": "YOLO11 Large - accurate object detection model by Ultralytics", + "docs": "https://docs.ultralytics.com/models/yolo11/", + "yolo_version": "yolo11l", + "model_type": "YOLO11", + "license": "agpl-3.0", + "license_link": "https://spdx.org/licenses/AGPL-3.0-only.html", + "dataset_type": "coco-detection" + }, + { + "model_short_name": "YOLO11x", + "model_library": "yolo", + "model_full_name": "YOLO11x", + "description": "YOLO11 Extra-large - highest accuracy object detection model by Ultralytics", + "docs": "https://docs.ultralytics.com/models/yolo11/", + "yolo_version": "yolo11x", + "model_type": "YOLO11", + "license": "agpl-3.0", + "license_link": "https://spdx.org/licenses/AGPL-3.0-only.html", + "dataset_type": "coco-detection" + } + ] +} diff --git a/model_converter/presets/datasets.json b/model_converter/presets/datasets.json new file mode 100644 index 000000000..19cff36aa --- /dev/null +++ b/model_converter/presets/datasets.json @@ -0,0 +1,9 @@ +{ + "datasets": { + "imagenet-1k": "/path/to/imagenet1k", + "imagenet-21k": "/path/to/imagenet21k", + "coco-detection": "/path/to/coco-detection", + "ade20k": "/path/to/ade20k", + "coco-segmentation": "/path/to/coco-segmentation" + } +} diff --git a/model_converter/pyproject.toml b/model_converter/pyproject.toml index ffebb6d1d..e429c86ee 100644 --- a/model_converter/pyproject.toml +++ b/model_converter/pyproject.toml @@ -34,8 +34,11 @@ dependencies = [ "onnx", "opencv-python-headless<4.14", "openvino>=2025.3", + "openvino-model-api", + "pycocotools", "timm", "torch", + "torchmetrics", "torchvision", "transformers", "ultralytics", @@ -43,7 +46,6 @@ dependencies = [ [project.scripts] model-converter = "model_converter.model_converter:main" -model-converter-yolo = "model_converter.yolo.yolo:main" [project.urls] Homepage = "https://github.com/open-edge-platform/model_api" @@ -52,7 +54,6 @@ Repository = "https://github.com/open-edge-platform/model_api.git" [dependency-groups] tests = [ - "openvino-model-api", "pytest", "pytest-cov", "pytest-mock", diff --git a/model_converter/src/model_converter/__init__.py b/model_converter/src/model_converter/__init__.py index 05acc4b9a..8bd3ce86b 100644 --- a/model_converter/src/model_converter/__init__.py +++ b/model_converter/src/model_converter/__init__.py @@ -6,5 +6,25 @@ """Tools for converting models to OpenVINO IR.""" from .cli import ModelConverter, list_models, main +from .converters import ( + CONVERTER_REGISTRY, + BaseConverter, + GetituneConverter, + PyTorchConverter, + TimmConverter, + TorchvisionConverter, + YoloConverter, +) -__all__ = ["ModelConverter", "list_models", "main"] +__all__ = [ + "CONVERTER_REGISTRY", + "BaseConverter", + "GetituneConverter", + "ModelConverter", + "PyTorchConverter", + "TimmConverter", + "TorchvisionConverter", + "YoloConverter", + "list_models", + "main", +] diff --git a/model_converter/src/model_converter/cli.py b/model_converter/src/model_converter/cli.py index 41d18de19..1952707b9 100644 --- a/model_converter/src/model_converter/cli.py +++ b/model_converter/src/model_converter/cli.py @@ -13,57 +13,59 @@ """ import argparse -import importlib import json import logging -import shutil import sys from pathlib import Path from typing import Any -import cv2 -import numpy as np -import torch -import torch.nn as nn - -from model_converter.adapters import get_adapter -from model_converter.downloaders import URLDownloader - -_MODEL_API_METADATA_FIELDS = ( - "resize_type", - "pad_value", - "input_dtype", - "confidence_threshold", - "postprocess_semantic_masks", - "nms_execute", - "iou_threshold", - "agnostic_nms", - "nms_max_predictions", +from model_converter.converters import CONVERTER_REGISTRY, BaseConverter +from model_converter.converters.getitune import GetituneConverter +from model_converter.dataset_registry import DatasetRegistry +from model_converter.reporting import ( + ConversionResult, + format_console_table, ) class ModelConverter: - """Handles conversion of PyTorch models to OpenVINO format.""" + """Facade for model conversion that dispatches to specialized converters. + + Routes each model configuration to the appropriate converter based on + the ``model_library`` field. Maintains backward compatibility with + existing usage patterns. + """ def __init__( self, output_dir: Path, cache_dir: Path, verbose: bool = False, - dataset_path: Path | None = None, + dataset_registry: DatasetRegistry | None = None, + training_extensions_dir: Path | None = None, + report_path: Path | None = None, + measure_accuracy: bool = True, ): - """ - Initialize the ModelConverter. + """Initialize the ModelConverter. Args: output_dir: Directory to save converted models cache_dir: Directory to cache downloaded weights verbose: Enable verbose logging - dataset_path: Path to calibration dataset for quantization + dataset_registry: Dataset registry for resolving dataset paths + training_extensions_dir: Path to training_extensions repo (for getitune models) + report_path: Path to the Markdown report file. When set, each + non-skipped conversion result is written to the file immediately + after export. + measure_accuracy: When ``False``, skip per-model accuracy + measurement during quantization. """ self.output_dir = Path(output_dir) self.cache_dir = Path(cache_dir) - self.dataset_path = Path(dataset_path) if dataset_path else None + self.dataset_registry = dataset_registry + self.training_extensions_dir = Path(training_extensions_dir) if training_extensions_dir else None + self.report_path = Path(report_path) if report_path else None + self.measure_accuracy = measure_accuracy self.output_dir.mkdir(parents=True, exist_ok=True) self.cache_dir.mkdir(parents=True, exist_ok=True) @@ -75,791 +77,46 @@ def __init__( ) self.logger = logging.getLogger(__name__) - # Initialize downloaders - self._url_downloader = URLDownloader(cache_dir=self.cache_dir) + # Initialize converters + self._converters: dict[str, BaseConverter] = {} + self._verbose = verbose - def get_labels(self, label_set: str) -> str | None: - """ - Get label list for a given label set. + def _get_converter(self, model_library: str) -> BaseConverter: + """Get or create a converter for the given model library. Args: - label_set: Name of the label set (e.g., "IMAGENET1K_V1") + model_library: Library name (torchvision, timm, yolo, getitune) Returns: - Space-separated string of labels, or None if not found + Appropriate converter instance """ - if label_set == "IMAGENET1K_V1": - from torchvision.models._meta import _IMAGENET_CATEGORIES - - categories = _IMAGENET_CATEGORIES - categories = [label.replace(" ", "_") for label in categories] - return " ".join(categories) - - if label_set == "IMAGENET21K": - from timm.data import ImageNetInfo - - info = ImageNetInfo("imagenet21k") - categories = info.label_descriptions() - categories = [desc.split(",")[0].strip().replace(" ", "_") for desc in categories] - return " ".join(categories) - - if label_set == "COCO_V1": - from torchvision.models.detection import MaskRCNN_ResNet50_FPN_Weights - - categories = MaskRCNN_ResNet50_FPN_Weights.COCO_V1.meta["categories"] - categories = [label.replace(" ", "_") for label in categories] - return " ".join(categories) - - return None - - def load_model_class( - self, - class_path: str, - ) -> type: - """ - Dynamically load a model class from a Python path. - - Args: - class_path: Full Python path to the class (e.g., 'torchvision.models.resnet.resnet50') - - Returns: - The model class - """ - try: - module_path, class_name = class_path.rsplit(".", 1) - self.logger.debug(f"Importing module: {module_path}") - # nosemgrep: python.lang.security.audit.non-literal-import.non-literal-import - module = importlib.import_module( - module_path, - ) - model_class = getattr(module, class_name) - self.logger.debug(f"Loaded class: {class_name}") - return model_class - except Exception as e: - self.logger.error(f"Failed to import module {module_path}: {e}") - raise - - def load_checkpoint( - self, - checkpoint_path: Path, - ) -> dict[str, Any]: - """ - Load PyTorch checkpoint file. - - Args: - checkpoint_path: Path to checkpoint file - - Returns: - Checkpoint dictionary - """ - try: - checkpoint = torch.load( # nosemgrep: trailofbits.python.pickles-in-pytorch.pickles-in-pytorch - checkpoint_path, - map_location="cpu", - weights_only=True, - ) - self.logger.debug(f"Loaded checkpoint from: {checkpoint_path}") - return checkpoint - except Exception as e: - self.logger.error(f"Failed to load checkpoint: {e}") - raise - - def load_huggingface_model( - self, - repo_id: str, - revision: str, - model_library: str = "timm", - model_params: dict[str, Any] | None = None, - ) -> nn.Module: - """ - Load a model from Hugging Face Hub. - - Args: - repo_id: Hugging Face repository ID - revision: Immutable revision/commit SHA for the Hugging Face repository - model_library: Library to use ('timm', 'transformers', etc.) - model_params: Optional parameters for model loading - - Returns: - Loaded model instance - """ - try: - if model_library == "timm": - import timm - - repo_ref = f"hf-hub:{repo_id}@{revision}" - self.logger.info(f"Loading timm model: {repo_ref}") - model = timm.create_model( - repo_ref, - pretrained=True, - cache_dir=self.cache_dir, - **(model_params or {}), + if model_library not in self._converters: + converter_cls = CONVERTER_REGISTRY.get(model_library) + if converter_cls is None: + error_msg = ( + f"Unsupported model_library: '{model_library}'. " + f"Supported libraries: {list(CONVERTER_REGISTRY.keys())}" ) - elif model_library == "transformers": - from transformers import AutoModel - - self.logger.info(f"Loading transformers model: {repo_id}@{revision}") - model = AutoModel.from_pretrained( - repo_id, - revision=revision, - cache_dir=self.cache_dir, - **(model_params or {}), - ) - else: - error_msg = f"Unsupported model library: {model_library}" - raise ValueError(error_msg) - - model.eval() - self.logger.info("✓ Hugging Face model loaded successfully") - return model - - except Exception as e: - self.logger.error(f"Failed to load Hugging Face model: {e}") - raise - - def create_model( - self, - model_class: type, - checkpoint: dict[str, Any], - model_params: dict[str, Any] | None = None, - ) -> nn.Module: - """ - Create and initialize model instance. - - Args: - model_class: Model class to instantiate - checkpoint: Checkpoint containing model weights - model_params: Optional parameters for model initialization - - Returns: - Initialized model instance - """ - try: - # Handle torch.nn.Module case (checkpoint contains full model) - if model_class == torch.nn.Module: - if "model" in checkpoint: - model = checkpoint["model"] - elif "state_dict" in checkpoint: - # Cannot reconstruct architecture from state_dict alone - error_msg = ( - "Checkpoint contains only state_dict. Please specify the model class instead of torch.nn.Module" - ) - raise ValueError(error_msg) - else: - # Assume checkpoint is the model itself - model = checkpoint - - if not isinstance(model, nn.Module): - error_msg = "Checkpoint does not contain a valid model" - raise ValueError(error_msg) - else: - # Instantiate model class - model = model_class(**model_params) if model_params else model_class() - - # Load weights - if "state_dict" in checkpoint: - state_dict = checkpoint["state_dict"] - elif "model" in checkpoint: - if isinstance(checkpoint["model"], nn.Module): - return checkpoint["model"] - state_dict = checkpoint["model"] - else: - state_dict = checkpoint - - model.load_state_dict(state_dict, strict=False) - - model.eval() - self.logger.info("✓ Model created and loaded successfully") - return model - - except Exception as e: - self.logger.error(f"Failed to create model: {e}") - raise - - def copy_readme( - self, - model_config: dict[str, Any], - output_folder: Path, - variant: str = "fp16", - ) -> None: - """ - Copy README template to model folder and replace placeholders. - - Args: - model_config: Model configuration used to fill template placeholders - output_folder: Folder where the model is saved - variant: Model variant ('fp16' or 'int8') - """ - try: - model_short_name = str(model_config.get("model_short_name", "")).strip() - model_library = str(model_config.get("model_library", "timm")).strip() - model_license = str(model_config.get("license", "")).strip() - model_license_link = str(model_config.get("license_link", "")).strip() - docs = str(model_config.get("docs", "")).strip() - - def template_placeholder(name: str) -> str: - return f"<<{name}>>" - - if not model_short_name: - error_msg = "Model config must define a non-empty model_short_name" - raise ValueError(error_msg) - - if not model_license_link: - error_msg = f"Model '{model_short_name}' must define a non-empty license_link" raise ValueError(error_msg) - if not model_license: - error_msg = f"Model '{model_short_name}' must define a non-empty license" - raise ValueError(error_msg) - - if not docs: - self.logger.warning( - f"Model '{model_short_name}' does not define 'docs' field. Placeholder will be empty.", - ) - - # Determine which README template to use based on model library - template_name = f"README-{model_library}-{variant}.md" - template_path = Path(__file__).parent.parent / "templates" / template_name - - if not template_path.exists(): - self.logger.warning(f"README template not found: {template_path}") - return - - # Read template - readme_content = template_path.read_text() - - placeholders = { - template_placeholder("license"): model_license, - template_placeholder("license_link"): model_license_link, - template_placeholder("model_name"): model_short_name, - template_placeholder("model_short_name"): model_short_name, - template_placeholder("variant"): variant, - template_placeholder("docs"): docs, - } - - for key, value in model_config.items(): - if value is None: - continue - if isinstance(value, (str, int, float, bool)): - placeholders[template_placeholder(key)] = str(value) - - for placeholder, value in placeholders.items(): - readme_content = readme_content.replace(placeholder, value) - - # Write to model folder - output_readme = output_folder / "README.md" - output_readme.write_text(readme_content) - self.logger.debug(f"Copied README to: {output_readme}") - - except (OSError, UnicodeError, ValueError) as e: - self.logger.warning(f"Failed to copy README: {e}") - - def _collect_dataset_entries(self, image_dir: Path) -> list[tuple[Path, int]]: - """Collect dataset image paths with their class labels.""" - image_entries: list[tuple[Path, int]] = [] - for class_dir in sorted(image_dir.iterdir()): - if class_dir.is_dir(): - class_label = int(class_dir.name) - for pattern in ["*.JPEG", "*.jpg", "*.png"]: - for img_path in class_dir.glob(pattern): - image_entries.append((img_path, class_label)) - return image_entries - - def _preprocess_calibration_image( - self, - img_path: Path, - width: int, - height: int, - mean: np.ndarray, - scale: np.ndarray, - reverse_input_channels: bool, - ) -> np.ndarray | None: - """Load and preprocess a single calibration image.""" - img = cv2.imread(str(img_path)) - if img is None: - return None - - img = cv2.resize(img, (width, height)) - img = img.astype(np.float32) - - if reverse_input_channels: - img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) - - img = (img - mean) / scale - img = img.transpose(2, 0, 1) - return np.expand_dims(img, axis=0) - - def create_calibration_dataset( - self, - input_shape: list[int], - mean_values: str | None = None, - scale_values: str | None = None, - reverse_input_channels: bool = True, - subset_size: int = 5000, - return_labels: bool = False, - ) -> tuple[list[np.ndarray], list[int]] | list[np.ndarray]: - """ - Create calibration dataset from sample validation images. - - Args: - input_shape: Target input shape [batch, channels, height, width] - mean_values: Space-separated mean values for normalization - scale_values: Space-separated scale values for normalization - reverse_input_channels: Whether to reverse RGB to BGR - subset_size: Number of images to use for calibration - return_labels: Whether to return labels along with images - - Returns: - List of preprocessed image arrays, or tuple of (images, labels) - """ - if not self.dataset_path or not self.dataset_path.exists(): - self.logger.warning("Dataset path not provided or doesn't exist. Skipping quantization.") - return [] - - # Parse mean and scale values - mean = np.array([float(x) for x in mean_values.split()]) if mean_values else np.array([0, 0, 0]) - scale = np.array([float(x) for x in scale_values.split()]) if scale_values else np.array([1, 1, 1]) - - _, _, height, width = input_shape - calibration_data: list[np.ndarray] = [] - - # Find all images in the dataset - image_dir = self.dataset_path - if not image_dir.exists(): - self.logger.error(f"Image directory not found: {image_dir}") - return ([], []) - - image_entries = self._collect_dataset_entries(image_dir) - if not image_entries: - self.logger.error("No images found in dataset") - return ([], []) - - self.logger.info(f"Found {len(image_entries)} images in dataset") - self.logger.info(f"Using {min(subset_size, len(image_entries))} images for calibration") - - if return_labels: - labels: list[int] = [] - for i, (img_path, class_label) in enumerate(image_entries[:subset_size]): - try: - img = self._preprocess_calibration_image( - img_path=img_path, - width=width, - height=height, - mean=mean, - scale=scale, - reverse_input_channels=reverse_input_channels, - ) - if img is None: - continue - - calibration_data.append(img) - labels.append(class_label) - - if (i + 1) % 50 == 0: - self.logger.debug(f"Processed {i + 1}/{subset_size} images") - - except (cv2.error, OSError, TypeError, ValueError) as e: - self.logger.warning(f"Failed to process {img_path}: {e}") - continue - - self.logger.info(f"✓ Created calibration dataset with {len(calibration_data)} images") - return calibration_data, labels - - for i, (img_path, _) in enumerate(image_entries[:subset_size]): - try: - img = self._preprocess_calibration_image( - img_path=img_path, - width=width, - height=height, - mean=mean, - scale=scale, - reverse_input_channels=reverse_input_channels, - ) - if img is None: - continue - - calibration_data.append(img) - - if (i + 1) % 50 == 0: - self.logger.debug(f"Processed {i + 1}/{subset_size} images") - - except (cv2.error, OSError, TypeError, ValueError) as e: - self.logger.warning(f"Failed to process {img_path}: {e}") - continue - - self.logger.info(f"✓ Created calibration dataset with {len(calibration_data)} images") - return calibration_data, [] - - def validate_model( - self, - model_path: Path, - validation_data: list[np.ndarray], - labels: list[int], - ) -> float: - """ - Validate OpenVINO model and compute top-1 accuracy. - - Args: - model_path: Path to the OpenVINO model (.xml) - validation_data: List of validation images - labels: List of ground truth labels - - Returns: - Top-1 accuracy (0.0 to 1.0) - """ - try: - import openvino as ov - - core = ov.Core() - model = core.read_model(model_path) - compiled_model = core.compile_model(model, device_name="CPU") - output_layer = compiled_model.outputs[0] - - predictions: list[int] = [] - for img in validation_data: - result = compiled_model(img)[output_layer] - pred_class = np.argmax(result, axis=1)[0] - predictions.append(pred_class) - - # Compute accuracy - correct = sum(predicted == label for predicted, label in zip(predictions, labels)) - return correct / len(labels) - - except (ImportError, OSError, RuntimeError, TypeError, ValueError) as e: - self.logger.error(f"Failed to validate model: {e}") - return 0.0 - - def quantize_model( - self, - model_path: Path, - calibration_data: list[np.ndarray], - model_config: dict[str, Any], - preset: str = "accuracy", - validation_data: list[np.ndarray] | None = None, - validation_labels: list[int] | None = None, - ) -> Path: - """ - Quantize OpenVINO model to INT8 using NNCF. - - Args: - model_path: Path to the FP32 OpenVINO model (.xml) - calibration_data: List of calibration images - model_config: Model configuration used for README rendering - preset: Quantization preset ('accuracy', 'performance', 'mixed') - validation_data: Optional validation images for accuracy measurement - validation_labels: Optional validation labels for accuracy measurement - - Returns: - Path to the quantized model - """ - if not calibration_data: - self.logger.warning("No calibration data provided. Skipping quantization.") - return model_path - - try: - import nncf - import openvino as ov - - self.logger.info(f"Quantizing model with {len(calibration_data)} calibration samples") - self.logger.info(f"Using preset: {preset}") - - # Load the model - core = ov.Core() - model = core.read_model(model_path) - - # Create calibration dataset generator - def calibration_dataset(): - for data in calibration_data: - yield data - - # Map preset string to NNCF enum - preset_map = { - "performance": nncf.QuantizationPreset.PERFORMANCE, - "mixed": nncf.QuantizationPreset.MIXED, + kwargs: dict[str, Any] = { + "output_dir": self.output_dir, + "cache_dir": self.cache_dir, + "verbose": self._verbose, + "dataset_registry": self.dataset_registry, + "report_path": self.report_path, + "measure_accuracy": self.measure_accuracy, } - nncf_preset = preset_map.get(preset.lower(), nncf.QuantizationPreset.MIXED) - - # Quantize the model - quantized_model = nncf.quantize( - model, - calibration_dataset=nncf.Dataset(calibration_dataset()), - preset=nncf_preset, - subset_size=len(calibration_data), - ) - - # Extract model name from the FP32 model path - # The FP32 path is like: output_dir/model_name-fp16-ov/model_name_fp32.xml - model_name = model_path.stem # Gets model_name_fp32 from model_name_fp32.xml - # Remove _fp32 suffix if present - if model_name.endswith("_fp32"): - model_name = model_name[:-5] - - # Create output folder with -int8-ov suffix - output_folder = model_path.parent.parent / f"{model_name}-int8-ov" - output_folder.mkdir(parents=True, exist_ok=True) - - # Save quantized model with model name inside the folder - output_path = output_folder / f"{model_name}.xml" - ov.save_model(quantized_model, output_path, compress_to_fp16=True) - self.logger.info(f"✓ Quantized model saved: {output_path}") - - # 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) - - # Validate accuracy if validation data provided - if validation_data and validation_labels: - self.logger.info("Validating FP32 model accuracy...") - fp32_accuracy = self.validate_model(model_path, validation_data, validation_labels) - self.logger.info(f"FP32 Top-1 Accuracy: {fp32_accuracy * 100:.2f}%") - - self.logger.info("Validating INT8 model accuracy...") - int8_accuracy = self.validate_model(output_path, validation_data, validation_labels) - self.logger.info(f"INT8 Top-1 Accuracy: {int8_accuracy * 100:.2f}%") - - accuracy_drop = (fp32_accuracy - int8_accuracy) * 100 - self.logger.info(f"Accuracy Drop: {accuracy_drop:.2f}%") - - # 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}") - - # Copy README for INT8 model - self.copy_readme( - model_config, - output_folder, - variant="int8", - ) - - return output_path - - except ImportError: - self.logger.error("NNCF not installed. Install with: pip install nncf") - return model_path - except (OSError, RuntimeError, TypeError, ValueError) as e: - self.logger.error(f"Failed to quantize model: {e}") - import traceback - - self.logger.debug(traceback.format_exc()) - return model_path - - def export_to_openvino( - self, - model: nn.Module, - input_shape: list[int], - output_path: Path, - model_config: dict[str, Any], - input_names: list[str] | None = None, - output_names: list[str] | None = None, - metadata: dict[tuple[str, str], str] | None = None, - ) -> tuple[Path, Path]: - """ - Export PyTorch model to OpenVINO format. - - Args: - model: PyTorch model to export - input_shape: Input tensor shape [batch, channels, height, width] - output_path: Path to save the model (without extension) - model_config: Model configuration used for README rendering - input_names: Names for input tensors - output_names: Names for output tensors - metadata: Metadata to embed in the model - - Returns: - Tuple of (fp16_model_path, fp32_model_path) - FP16 for final use, FP32 for quantization - """ - import openvino as ov - - try: - model = self._prepare_model_for_export(model, model_config) - model.eval() - dummy_input = self._create_example_input(input_shape, model_config) - self.logger.info("Direct PyTorch to OpenVINO conversion") - ov_model = ov.convert_model(model, example_input=dummy_input) - self.logger.info("✓ PyTorch to OpenVINO conversion complete") - - # Reshape model to fixed input shape (remove dynamic dimensions) - first_input = ov_model.input(0) - input_name_for_reshape = next(iter(first_input.get_names())) if first_input.get_names() else 0 - - self.logger.debug(f"Setting fixed input shape: {input_shape}") - ov_model.reshape({input_name_for_reshape: input_shape}) - - # Post-process the model - ov_model = self._postprocess_openvino_model( - ov_model, - input_names=input_names, - output_names=output_names, - metadata=metadata, - ) - - # Create output folder with -fp16-ov suffix - model_name = output_path.name - output_folder = output_path.parent / f"{model_name}-fp16-ov" - output_folder.mkdir(parents=True, exist_ok=True) - - # Save FP32 model for quantization (temporary) - fp32_xml_path = output_folder / f"{model_name}_fp32.xml" - ov.save_model(ov_model, fp32_xml_path, compress_to_fp16=False) - self.logger.debug(f"Saved FP32 model for quantization: {fp32_xml_path}") - - # Save the FP16 model (final) - xml_path = output_folder / f"{model_name}.xml" - ov.save_model(ov_model, xml_path, compress_to_fp16=True) - 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) - - # 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}") - - # Copy README for FP16 model - self.copy_readme( - model_config, - output_folder, - variant="fp16", - ) - - return xml_path, fp32_xml_path - - except Exception as e: - self.logger.error(f"Failed to export model: {e}") - raise - - def _prepare_model_for_export(self, model: nn.Module, model_config: dict[str, Any]) -> nn.Module: - """Prepare model for OpenVINO conversion.""" - model_type = str(model_config.get("model_type", "")) - adapted = get_adapter(model_type, model) - if adapted is not model: - self.logger.info(f"Applied export adapter for model type: {model_type}") - return adapted - - def _create_example_input(self, input_shape: list[int], model_config: dict[str, Any]) -> torch.Tensor: - """Create example input suitable for the configured model type.""" - if str(model_config.get("model_type", "")).lower() == "maskrcnn": - return torch.rand(*input_shape) - return torch.randn(*input_shape) - - def _postprocess_openvino_model( - self, - model: Any, - input_names: list[str] | None = None, - output_names: list[str] | None = None, - metadata: dict[tuple[str, str], str] | None = None, - ) -> Any: - """ - Post-process OpenVINO model (set names, add metadata). - - Args: - model: OpenVINO model - input_names: Names for input tensors - output_names: Names for output tensors - metadata: Metadata to embed - Returns: - Post-processed model - """ - # Set input names - if input_names: - for i, name in enumerate(input_names): - if i < len(model.inputs): - model.input(i).set_names({name}) - self.logger.debug(f"Set input {i} name to: {name}") - - # Set output names - if output_names: - for i, name in enumerate(output_names): - if i < len(model.outputs): - model.output(i).set_names({name}) - self.logger.debug(f"Set output {i} name to: {name}") - - # Add metadata - if metadata: - for key, value in metadata.items(): - model.set_rt_info(value, list(key)) - self.logger.debug(f"Set metadata {key}: {value}") - - return model - - def _load_model_from_config(self, config: dict[str, Any]) -> Any: - """Load a PyTorch model based on configuration (HuggingFace or traditional weights).""" - huggingface_repo = config.get("huggingface_repo") - if huggingface_repo: - huggingface_revision = config.get("huggingface_revision") - if not huggingface_revision: - error_msg = "Hugging Face models must define 'huggingface_revision' with an immutable commit SHA" - raise ValueError(error_msg) + if converter_cls == GetituneConverter: + kwargs["training_extensions_dir"] = self.training_extensions_dir - model_library = config.get("model_library", "timm") - model_params = config.get("model_params") - return self.load_huggingface_model( - repo_id=huggingface_repo, - revision=huggingface_revision, - model_library=model_library, - model_params=model_params, - ) - - # Traditional PyTorch model workflow - weights_url = config["weights_url"] - weights_path = self._url_downloader.download(url=weights_url) - - model_class_name = config.get("model_class_name", "torch.nn.Module") - model_class = self.load_model_class(model_class_name) - - checkpoint = self.load_checkpoint(weights_path) - - model_params = config.get("model_params") - return self.create_model(model_class, checkpoint, model_params) - - def _quantize_and_cleanup(self, config: dict[str, Any], fp32_model_path: Path, **kwargs: Any) -> None: - """Run INT8 quantization and clean up temporary FP32 model files.""" - model_type = kwargs["model_type"] - self.logger.info("Creating calibration dataset for INT8 quantization") - return_validation_labels = model_type == "Classification" and bool(config.get("labels")) - - if return_validation_labels: - self.logger.info("Creating validation dataset for accuracy measurement") - validation_data, validation_labels = self.create_calibration_dataset( - input_shape=kwargs["input_shape"], - mean_values=kwargs["mean_values"], - scale_values=kwargs["scale_values"], - reverse_input_channels=kwargs["reverse_input_channels"], - subset_size=300, - return_labels=return_validation_labels, - ) + self._converters[model_library] = converter_cls(**kwargs) - if validation_data: - self.quantize_model( - model_path=fp32_model_path, - calibration_data=validation_data, - model_config=config, - preset="mixed", - validation_data=validation_data if validation_labels else None, - validation_labels=validation_labels or None, - ) - - # 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}") + return self._converters[model_library] def process_model_config(self, config: dict[str, Any]) -> bool: - """ - Process a single model configuration. + """Process a single model configuration by dispatching to the appropriate converter. Args: config: Model configuration dictionary @@ -867,119 +124,26 @@ def process_model_config(self, config: dict[str, Any]) -> bool: Returns: True if successful, False otherwise """ - model_short_name = config.get("model_short_name", "unknown") - model_license = config.get("license") - model_license_link = config.get("license_link") - - # 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") - return True - + model_library = config.get("model_library", "torchvision") try: - 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) - - model = self._load_model_from_config(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"]) - - # Prepare metadata from config (with defaults for normalization) - 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", "") - - metadata = { - ("model_info", "model_type"): model_type, - ("model_info", "model_short_name"): model_short_name, - ("model_info", "reverse_input_channels"): self._metadata_value(reverse_input_channels), - ("model_info", "mean_values"): mean_values, - ("model_info", "scale_values"): scale_values, - } - - for metadata_field in _MODEL_API_METADATA_FIELDS: - if metadata_field in config and config[metadata_field] is not None: - metadata["model_info", metadata_field] = self._metadata_value(config[metadata_field]) - - # Add labels if specified in config - labels_config = config.get("labels") - if labels_config: - labels = self.get_labels(labels_config) - if labels: - metadata["model_info", "labels"] = labels - self.logger.info(f"Added {labels_config} labels to metadata") - else: - self.logger.warning(f"Could not load labels for: {labels_config}") - - output_path = self.output_dir / model_short_name - fp16_model_path, fp32_model_path = self.export_to_openvino( - model=model, - input_shape=input_shape, - output_path=output_path, - model_config=config, - input_names=input_names, - output_names=output_names, - metadata=metadata, - ) - - # Quantize the model if dataset is available - if self.dataset_path and self.dataset_path.exists(): - 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, - ) - - 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}") - import traceback - - self.logger.debug(traceback.format_exc()) + converter = self._get_converter(model_library) + return converter.process_model_config(config) + except ValueError as e: + self.logger.error(f"✗ {e}") return False - @staticmethod - def _metadata_value(value: Any) -> str: - """Convert config values to Model API rt_info string values.""" - if isinstance(value, (list, tuple)): - return " ".join(str(item) for item in value) - return str(value) - def process_config_file( self, config_path: Path, model_filter: str | None = None, + library_filter: list[str] | None = None, ) -> tuple[int, int]: - """ - Process models from a configuration file. + """Process models from a configuration file. Args: config_path: Path to JSON configuration file model_filter: Optional model short name to process (process only this model) + library_filter: Optional list of library names to process (process only these libraries) Returns: Tuple of (successful_count, failed_count) @@ -999,6 +163,17 @@ def process_config_file( self.logger.info(f"Configuration validated: {len(models)} models found") + # Filter by library if requested + if library_filter: + unknown_libs = set(library_filter) - set(CONVERTER_REGISTRY.keys()) + for lib in unknown_libs: + self.logger.warning(f"Unknown library '{lib}' (available: {list(CONVERTER_REGISTRY.keys())})") + models = [m for m in models if m.get("model_library") in library_filter] + if not models: + self.logger.error(f"No models found for libraries: {library_filter}") + return 0, 0 + self.logger.info(f"Filtering by libraries: {library_filter} ({len(models)} models)") + # Filter models if requested if model_filter: models = [m for m in models if m.get("model_short_name") == model_filter] @@ -1018,9 +193,25 @@ def process_config_file( return successful, failed + def collect_results(self) -> list[ConversionResult]: + """Aggregate conversion results from every instantiated converter. + + Returns: + The combined list of per-model conversion results. + """ + results: list[ConversionResult] = [] + for converter in self._converters.values(): + results.extend(converter.results) + return results + + +def list_models(config_path: Path, library_filter: list[str] | None = None) -> None: + """List all models in a configuration file. -def list_models(config_path: Path) -> None: - """List all models in a configuration file.""" + Args: + config_path: Path to JSON configuration file + library_filter: Optional list of library names to filter by + """ try: with config_path.open() as f: config = json.load(f) @@ -1030,19 +221,26 @@ def list_models(config_path: Path) -> None: models = config.get("models", []) + if library_filter: + unknown_libs = set(library_filter) - set(CONVERTER_REGISTRY.keys()) + for lib in unknown_libs: + print(f"Warning: Unknown library '{lib}' (available: {list(CONVERTER_REGISTRY.keys())})", file=sys.stderr) + models = [m for m in models if m.get("model_library") in library_filter] + if not models: print("No models found in configuration") return print(f"\nFound {len(models)} models:\n") - print(f"{'Short Name':<30} {'Full Name':<40} {'Type':<20}") - print("-" * 90) + print(f"{'Short Name':<30} {'Full Name':<40} {'Library':<15} {'Type':<20}") + print("-" * 105) for model in models: short_name = model.get("model_short_name", "N/A") full_name = model.get("model_full_name", "N/A") model_type = model.get("model_type", "N/A") - print(f"{short_name:<30} {full_name:<40} {model_type:<20}") + library = model.get("model_library", "N/A") + print(f"{short_name:<30} {full_name:<40} {library:<15} {model_type:<20}") print() @@ -1091,11 +289,10 @@ def main(): ) parser.add_argument( - "-d", - "--dataset", + "--datasets-config", type=Path, - default=Path.home() / "model_api" / "validation_dataset", - help=("Path to calibration dataset for INT8 quantization (default: ~/model_api/validation_dataset)"), + default=Path("presets/datasets.json"), + help="Path to datasets configuration JSON file (default: presets/datasets.json)", ) parser.add_argument( @@ -1104,12 +301,59 @@ def main(): help="Process only the specified model (by model_short_name)", ) + parser.add_argument( + "--library", + type=str, + default=None, + help=( + "Comma-separated list of model libraries to process (e.g., getitune,timm). " + "If not provided, all libraries are exported." + ), + ) + parser.add_argument( "--list", action="store_true", help="List all models in the configuration file and exit", ) + parser.add_argument( + "--report", + nargs="?", + const="", + default="", + metavar="PATH", + help=( + "Generate a conversion summary report (enabled by default). Pass the flag alone to write to " + "/conversion_report.md, or provide a PATH to override the location. " + "The report is printed to the console and saved as Markdown. " + "Use --no-report to disable." + ), + ) + + parser.add_argument( + "--no-report", + dest="report", + action="store_const", + const=None, + help="Disable the conversion summary report.", + ) + + parser.add_argument( + "--measure-accuracy", + dest="measure_accuracy", + action="store_true", + default=True, + help="Measure per-model accuracy during quantization (default).", + ) + + parser.add_argument( + "--no-measure-accuracy", + dest="measure_accuracy", + action="store_false", + help="Skip per-model accuracy measurement during quantization.", + ) + parser.add_argument( "-v", "--verbose", @@ -1117,8 +361,20 @@ def main(): help="Enable verbose logging", ) + parser.add_argument( + "--training-extensions-dir", + type=Path, + default=None, + help="Path to cloned training_extensions repo (required for getitune models)", + ) + args = parser.parse_args() + # Parse library filter + library_filter: list[str] | None = None + if args.library: + library_filter = [lib.strip() for lib in args.library.split(",") if lib.strip()] + # Check if config file exists if not args.config.exists(): print(f"Error: Configuration file not found: {args.config}", file=sys.stderr) @@ -1126,7 +382,7 @@ def main(): # List models and exit if args.list: - list_models(args.config) + list_models(args.config, library_filter=library_filter) return 0 # Setup logging @@ -1140,23 +396,42 @@ def main(): logger.info(f"Loading configuration from: {args.config}") try: + report_path: Path | None = None + if args.report is not None: + report_path = Path(args.report) if args.report else args.output / "conversion_report.md" + + # Load dataset registry + dataset_registry: DatasetRegistry | None = None + if args.datasets_config.exists(): + try: + dataset_registry = DatasetRegistry(args.datasets_config) + logger.info(f"Loaded dataset registry from: {args.datasets_config}") + except (FileNotFoundError, ValueError) as e: + logger.warning(f"Failed to load dataset registry: {e}") + logger.warning("Quantization will be skipped for all models") + else: + logger.warning(f"Dataset configuration not found: {args.datasets_config}") + logger.warning("Quantization will be skipped for all models") + # Create converter converter = ModelConverter( output_dir=args.output, cache_dir=args.cache, verbose=args.verbose, - dataset_path=args.dataset, + dataset_registry=dataset_registry, + training_extensions_dir=args.training_extensions_dir, + report_path=report_path, + measure_accuracy=args.measure_accuracy, ) logger.info(f"Output directory: {args.output}") logger.info(f"Cache directory: {args.cache}") - if args.dataset: - logger.info(f"Calibration dataset: {args.dataset}") # Process models successful, failed = converter.process_config_file( config_path=args.config, model_filter=args.model, + library_filter=library_filter, ) # Print summary @@ -1167,6 +442,11 @@ def main(): logger.info(f" Total: {successful + failed}") logger.info("=" * 80) + if report_path is not None: + results = converter.collect_results() + print(format_console_table(results)) + logger.info(f"Conversion report written to: {report_path}") + return 0 if failed == 0 else 1 except (ValueError, RuntimeError, ImportError, FileNotFoundError) as e: logger.error(f"Failed to process model: {e}") diff --git a/model_converter/src/model_converter/converters/__init__.py b/model_converter/src/model_converter/converters/__init__.py new file mode 100644 index 000000000..7e74c619b --- /dev/null +++ b/model_converter/src/model_converter/converters/__init__.py @@ -0,0 +1,24 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Converters package for model_converter.""" + +from model_converter.converters.base import BaseConverter +from model_converter.converters.getitune import GetituneConverter +from model_converter.converters.pytorch import PyTorchConverter +from model_converter.converters.registry import CONVERTER_REGISTRY +from model_converter.converters.timm import TimmConverter +from model_converter.converters.torchvision import TorchvisionConverter +from model_converter.converters.yolo import YoloConverter + +__all__ = [ + "CONVERTER_REGISTRY", + "BaseConverter", + "GetituneConverter", + "PyTorchConverter", + "TimmConverter", + "TorchvisionConverter", + "YoloConverter", +] diff --git a/model_converter/src/model_converter/converters/base.py b/model_converter/src/model_converter/converters/base.py new file mode 100644 index 000000000..7232e1b3a --- /dev/null +++ b/model_converter/src/model_converter/converters/base.py @@ -0,0 +1,850 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Base converter class for model conversion pipelines.""" + +import json +import logging +import shutil +from abc import ABC, abstractmethod +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import cv2 +import numpy as np + +from model_converter.datasets import CalibrationSample, reader_for +from model_converter.metrics.coco_detection import COCO80_TO_COCO91 +from model_converter.reporting import ( + AccuracyResults, + ConversionResult, + determine_status, + original_url_for_config, + upsert_result, +) + +if TYPE_CHECKING: + from model_converter.dataset_registry import DatasetRegistry + from model_converter.metrics import CocoDetectionMAP, Metric + + +class BaseConverter(ABC): + """Abstract base class for model converters. + + Provides shared functionality for calibration dataset creation, + quantization, README rendering, and model validation. + """ + + def __init__( + self, + output_dir: Path, + cache_dir: Path, + verbose: bool = False, + dataset_registry: "DatasetRegistry | None" = None, + report_path: Path | None = None, + measure_accuracy: bool = True, + ): + """Initialize the BaseConverter. + + Args: + output_dir: Directory to save converted models + cache_dir: Directory to cache downloaded weights + verbose: Enable verbose logging + dataset_registry: Dataset registry for resolving dataset paths + report_path: Path to the Markdown report file. When set, each + non-skipped result is upserted into the file immediately after + conversion. + measure_accuracy: When ``False``, skip per-model accuracy + measurement even if a metric strategy is available. + """ + self.output_dir = Path(output_dir) + self.cache_dir = Path(cache_dir) + self.dataset_registry = dataset_registry + self.report_path = Path(report_path) if report_path else None + self.measure_accuracy = measure_accuracy + self.output_dir.mkdir(parents=True, exist_ok=True) + self.cache_dir.mkdir(parents=True, exist_ok=True) + + self.results: list[ConversionResult] = [] + + log_level = logging.DEBUG if verbose else logging.INFO + logging.basicConfig( + level=log_level, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + ) + self.logger = logging.getLogger(__name__) + + @abstractmethod + def process_model_config(self, config: dict[str, Any]) -> bool: + """Process a single model configuration. + + Args: + config: Model configuration dictionary + + Returns: + True if successful, False otherwise + """ + + def _resolve_dataset_path(self, config: dict[str, Any]) -> Path | None: + """Resolve dataset path from model configuration. + + Extracts the 'dataset_type' field from the model config and resolves + it to a filesystem path using the dataset registry. Returns None if + no dataset_type is specified or no registry is available. + + Args: + config: Model configuration dictionary + + Returns: + Path to the dataset directory, or None if not applicable + + Raises: + ValueError: If dataset_type is specified but not in registry + """ + if self.dataset_registry is None: + return None + + return self.dataset_registry.resolve_from_config(config) + + def _build_result(self, config: dict[str, Any]) -> ConversionResult: + """Create a ConversionResult seeded from a model configuration.""" + model_short_name = str(config.get("model_short_name", "unknown")) + return ConversionResult( + model_short_name=model_short_name, + model_full_name=str(config.get("model_full_name", model_short_name)), + model_type=str(config.get("model_type", "")), + model_library=str(config.get("model_library", "")), + original_url=original_url_for_config(config), + ) + + def _metric_for_config( + self, + config: dict[str, Any], + dataset_path: Path | None, + ) -> "Metric | None": + """Build the task-appropriate :class:`Metric` for a model config. + + Resolves the COCO annotation file when the dataset_type calls for it + and forwards ``getitune_task`` so that ``Classification`` model_type + with ``MULTI_LABEL_CLS`` is routed to multilabel mAP rather than + top-1. + """ + from model_converter.datasets.factory import _COCO_ANNOTATION_FILES + from model_converter.metrics import metric_for + + dataset_type = config.get("dataset_type") + model_type = config.get("model_type") + task = config.get("getitune_task") + annotation_file: Path | None = None + if dataset_path is not None and dataset_type in _COCO_ANNOTATION_FILES: + annotation_file = dataset_path / "annotations" / _COCO_ANNOTATION_FILES[dataset_type] + return metric_for(dataset_type, model_type, annotation_file=annotation_file, task=task) + + def _collect_validation_samples( + self, + dataset_path: Path | None, + dataset_type: str | None, + subset_size: int = 500, + ) -> "list[CalibrationSample]": + """Collect up to ``subset_size`` raw samples from the dataset reader. + + Used by metric paths that need raw image paths (and per-task GT + pointers like COCO ``image_id`` or ADE20K ``mask_path``) rather than + preprocessed calibration tensors. Returns an empty list when the + dataset is unavailable, unreadable, or unknown. + """ + if dataset_path is None or not dataset_path.exists(): + return [] + try: + reader = reader_for(dataset_type, dataset_path) + except ValueError: + return [] + samples: list[CalibrationSample] = [] + try: + for sample in reader: + samples.append(sample) + if len(samples) >= subset_size: + break + except (FileNotFoundError, OSError, ValueError) as e: + self.logger.warning(f"Failed to enumerate validation samples: {e}") + return [] + return samples + + def _record_result( + self, + result: ConversionResult, + *, + converted: bool, + quantized: bool, + skipped: bool = False, + accuracy: AccuracyResults | None = None, + ) -> ConversionResult: + """Finalize a ConversionResult: copy accuracies, set status, and store it.""" + if accuracy is not None and accuracy.measured: + result.original_accuracy = accuracy.original_accuracy + result.fp32_accuracy = accuracy.fp32_accuracy + result.fp16_accuracy = accuracy.fp16_accuracy + result.int8_accuracy = accuracy.int8_accuracy + result.status, result.status_detail = determine_status( + result, + converted=converted, + quantized=quantized, + skipped=skipped, + ) + self.results.append(result) + if self.report_path is not None and not skipped: + upsert_result(result, self.report_path) + return result + + def copy_readme( + self, + model_config: dict[str, Any], + output_folder: Path, + variant: str = "fp16", + ) -> None: + """Copy README template to model folder and replace placeholders. + + Args: + model_config: Model configuration used to fill template placeholders + output_folder: Folder where the model is saved + variant: Model variant ('fp16' or 'int8') + """ + try: + model_short_name = str(model_config.get("model_short_name", "")).strip() + model_library = str(model_config.get("model_library", "timm")).strip() + model_license = str(model_config.get("license", "")).strip() + model_license_link = str(model_config.get("license_link", "")).strip() + docs = str(model_config.get("docs", "")).strip() + + def template_placeholder(name: str) -> str: + return f"<<{name}>>" + + if not model_short_name: + error_msg = "Model config must define a non-empty model_short_name" + raise ValueError(error_msg) + + if not model_license_link: + error_msg = f"Model '{model_short_name}' must define a non-empty license_link" + raise ValueError(error_msg) + + if not model_license: + error_msg = f"Model '{model_short_name}' must define a non-empty license" + raise ValueError(error_msg) + + if not docs: + self.logger.warning( + f"Model '{model_short_name}' does not define 'docs' field. Placeholder will be empty.", + ) + + # Determine which README template to use based on model library + template_name = f"README-{model_library}-{variant}.md" + template_path = Path(__file__).parent.parent / "templates" / template_name + + if not template_path.exists(): + self.logger.warning(f"README template not found: {template_path}") + return + + # Read template + readme_content = template_path.read_text() + + placeholders = { + template_placeholder("license"): model_license, + template_placeholder("license_link"): model_license_link, + template_placeholder("model_name"): model_short_name, + template_placeholder("model_short_name"): model_short_name, + template_placeholder("variant"): variant, + template_placeholder("docs"): docs, + } + + # Handle tags list → YAML formatting + tags = model_config.get("tags") + if tags and isinstance(tags, list): + tags_yaml = "\n".join(f" - {tag}" for tag in tags) + placeholders[template_placeholder("tags_yaml")] = tags_yaml + else: + placeholders[template_placeholder("tags_yaml")] = "" + + for key, value in model_config.items(): + if value is None: + continue + if isinstance(value, (str, int, float, bool)): + placeholders[template_placeholder(key)] = str(value) + + for placeholder, value in placeholders.items(): + readme_content = readme_content.replace(placeholder, value) + + # Write to model folder + output_readme = output_folder / "README.md" + output_readme.write_text(readme_content) + self.logger.debug(f"Copied README to: {output_readme}") + + except (OSError, UnicodeError, ValueError) as e: + self.logger.warning(f"Failed to copy README: {e}") + + def get_labels(self, label_set: str) -> str | None: + """Get label list for a given label set. + + Args: + label_set: Name of the label set (e.g., "IMAGENET1K_V1") + + Returns: + Space-separated string of labels, or None if not found + """ + if label_set == "IMAGENET1K_V1": + from torchvision.models._meta import _IMAGENET_CATEGORIES + + categories = _IMAGENET_CATEGORIES + categories = [label.replace(" ", "_") for label in categories] + return " ".join(categories) + + if label_set == "IMAGENET21K": + from timm.data import ImageNetInfo + + info = ImageNetInfo("imagenet21k") + categories = info.label_descriptions() + categories = [desc.split(",")[0].strip().replace(" ", "_") for desc in categories] + return " ".join(categories) + + if label_set == "COCO_V1": + from torchvision.models.detection import MaskRCNN_ResNet50_FPN_Weights + + categories = MaskRCNN_ResNet50_FPN_Weights.COCO_V1.meta["categories"] + categories = [label.replace(" ", "_") for label in categories] + return " ".join(categories) + + return None + + def _collect_dataset_entries( + self, + image_dir: Path, + dataset_type: str | None = None, + ) -> list[tuple[Path, int]]: + """Collect dataset image paths with their class labels. + + Dispatches by ``dataset_type`` to the appropriate + :class:`~model_converter.datasets.DatasetReader`. ``None`` (the + default) preserves the legacy class-folder behaviour. + """ + reader = reader_for(dataset_type, image_dir) + return [(sample.image_path, sample.label) for sample in reader] + + @staticmethod + def _crop_resize(img: np.ndarray, width: int, height: int) -> np.ndarray: + """Center-crop to target aspect ratio, then resize. + + Matches the standard ImageNet evaluation pipeline used by timm and + torchvision (resize shorter side, then center-crop), giving better + accuracy than a plain stretch-to-fit resize. + + Args: + img: Input image in HWC layout. + width: Target width in pixels. + height: Target height in pixels. + + Returns: + Cropped and resized image. + """ + h, w = img.shape[:2] + desired_ar = width / height + if desired_ar == 1: + side = min(h, w) + y0 = (h - side) // 2 + x0 = (w - side) // 2 + img = img[y0 : y0 + side, x0 : x0 + side] + elif w / h > desired_ar: # image is too wide — crop width + new_w = int(h * desired_ar) + x0 = (w - new_w) // 2 + img = img[:, x0 : x0 + new_w] + else: # image is too tall — crop height + new_h = int(w / desired_ar) + y0 = (h - new_h) // 2 + img = img[y0 : y0 + new_h, :] + return cv2.resize(img, (width, height)) + + def _preprocess_calibration_image( + self, + img_path: Path, + width: int, + height: int, + mean: np.ndarray, + scale: np.ndarray, + reverse_input_channels: bool, + resize_type: str = "standard", + ) -> np.ndarray | None: + """Load and preprocess a single calibration image.""" + img = cv2.imread(str(img_path)) + if img is None: + return None + + img = self._crop_resize(img, width, height) if resize_type == "crop" else cv2.resize(img, (width, height)) + img = img.astype(np.float32) + + if reverse_input_channels: + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + + img = (img - mean) / scale + img = img.transpose(2, 0, 1) + return np.expand_dims(img, axis=0) + + def create_calibration_dataset( + self, + input_shape: list[int], + mean_values: str | None = None, + scale_values: str | None = None, + reverse_input_channels: bool = True, + subset_size: int = 500, + return_labels: bool = False, + resize_type: str = "standard", + dataset_path: Path | None = None, + dataset_type: str | None = None, + ) -> tuple[list[np.ndarray], list[int]]: + """Create calibration dataset from sample validation images. + + Args: + input_shape: Target input shape [batch, channels, height, width] + mean_values: Space-separated mean values for normalization + scale_values: Space-separated scale values for normalization + reverse_input_channels: Whether to reverse RGB to BGR + subset_size: Number of images to use for calibration + return_labels: Whether to return labels along with images + resize_type: Resize strategy — ``"crop"`` center-crops to the + target aspect ratio before resizing (matches standard ImageNet + evaluation), ``"standard"`` stretches directly to the target + size. + dataset_path: Path to dataset directory (overrides registry resolution) + dataset_type: Optional dataset_type identifier (e.g. ``"coco-detection"``) + used to dispatch path enumeration to the matching reader. Defaults + to ``None`` for backward-compatible class-folder behaviour. + + Returns: + Tuple of (images, labels); both empty lists when dataset is unavailable. + """ + if not dataset_path or not dataset_path.exists(): + self.logger.warning("Dataset path not provided or doesn't exist. Skipping quantization.") + return [], [] + + # Parse mean and scale values + mean = np.array([float(x) for x in mean_values.split()]) if mean_values else np.array([0, 0, 0]) + scale = np.array([float(x) for x in scale_values.split()]) if scale_values else np.array([1, 1, 1]) + + _, _, height, width = input_shape + calibration_data: list[np.ndarray] = [] + + # Find all images in the dataset + image_dir = dataset_path + if not image_dir.exists(): + self.logger.error(f"Image directory not found: {image_dir}") + return ([], []) + + image_entries = self._collect_dataset_entries(image_dir, dataset_type=dataset_type) + if not image_entries: + self.logger.error("No images found in dataset") + return ([], []) + + self.logger.info(f"Found {len(image_entries)} images in dataset") + self.logger.info(f"Using {min(subset_size, len(image_entries))} images for calibration") + + if return_labels: + labels: list[int] = [] + for i, (img_path, class_label) in enumerate(image_entries[:subset_size]): + try: + img = self._preprocess_calibration_image( + img_path=img_path, + width=width, + height=height, + mean=mean, + scale=scale, + reverse_input_channels=reverse_input_channels, + resize_type=resize_type, + ) + if img is None: + continue + + calibration_data.append(img) + labels.append(class_label) + + if (i + 1) % 50 == 0: + self.logger.debug(f"Processed {i + 1}/{subset_size} images") + + except (cv2.error, OSError, TypeError, ValueError) as e: + self.logger.warning(f"Failed to process {img_path}: {e}") + continue + + self.logger.info(f"✓ Created calibration dataset with {len(calibration_data)} images") + return calibration_data, labels + + for i, (img_path, _) in enumerate(image_entries[:subset_size]): + try: + img = self._preprocess_calibration_image( + img_path=img_path, + width=width, + height=height, + mean=mean, + scale=scale, + reverse_input_channels=reverse_input_channels, + resize_type=resize_type, + ) + if img is None: + continue + + calibration_data.append(img) + + if (i + 1) % 50 == 0: + self.logger.debug(f"Processed {i + 1}/{subset_size} images") + + except (cv2.error, OSError, TypeError, ValueError) as e: + self.logger.warning(f"Failed to process {img_path}: {e}") + continue + + self.logger.info(f"✓ Created calibration dataset with {len(calibration_data)} images") + return calibration_data, [] + + def validate_model( + self, + model_path: Path, + validation_data: list[np.ndarray], + labels: list[int], + ) -> float: + """Validate OpenVINO model and compute top-1 accuracy. + + Args: + model_path: Path to the OpenVINO model (.xml) + validation_data: List of validation images + labels: List of ground truth labels + + Returns: + Top-1 accuracy (0.0 to 1.0) + """ + try: + import openvino as ov + + core = ov.Core() + model = core.read_model(model_path) + compiled_model = core.compile_model(model, device_name="CPU") + output_layer = compiled_model.outputs[0] + + predictions: list[int] = [] + for img in validation_data: + result = compiled_model(img)[output_layer] + pred_class = np.argmax(result, axis=1)[0] + predictions.append(pred_class) + + # Compute accuracy + correct = sum(predicted == label for predicted, label in zip(predictions, labels)) + return correct / len(labels) + + except (ImportError, OSError, RuntimeError, TypeError, ValueError) as e: + self.logger.error(f"Failed to validate model: {e}") + return 0.0 + + def _measure_metric( + self, + model_path: Path, + samples: "list[CalibrationSample]", + metric: "Metric", + ) -> float | None: + """Measure a task-specific :class:`Metric` over an OpenVINO model via Model API. + + Loads the model through ``model_api.models.Model.create_model`` so that + task-correct preprocessing and postprocessing (including resize-info + reversal for detection and per-class argmax for segmentation) are + applied. Each sample's raw image is decoded with OpenCV and passed to + the wrapper, whose result is dispatched into ``metric.update`` by + metric type. + + Args: + model_path: Path to the OpenVINO ``.xml`` model whose rt_info + identifies the model_type used by ``Model.create_model``. + samples: Per-image ground-truth pointers (image path, optional + COCO ``image_id``, optional ADE20K ``mask_path``). + metric: A reset-or-fresh metric instance. The metric is reset + before iterating and ``compute()`` is returned at the end. + + Returns: + The scalar metric value, or ``None`` if the Model API import + fails or model loading raises. + """ + try: + from model_api.models import Model + except ImportError as e: + self.logger.error(f"Failed to import model_api: {e}") + return None + try: + metric.reset() + wrapper = Model.create_model(str(model_path)) + except (OSError, RuntimeError, TypeError, ValueError) as e: + self.logger.error(f"Failed to load model {model_path}: {e}") + return None + for sample in samples: + img = cv2.imread(str(sample.image_path)) + if img is None: + continue + try: + result = wrapper(img) + except (RuntimeError, TypeError, ValueError) as e: + self.logger.debug(f"Inference failed for {sample.image_path}: {e}") + continue + try: + self._update_metric_with_result(metric, result, sample) + except (RuntimeError, TypeError, ValueError, IndexError) as e: + self.logger.debug(f"Metric update failed for {sample.image_path}: {e}") + continue + return float(metric.compute()) + + def _update_metric_with_result( + self, + metric: "Metric", + result: Any, + sample: "CalibrationSample", + ) -> None: + """Translate a Model API result into the right ``metric.update`` call.""" + from model_converter.metrics import ( + CocoDetectionMAP, + MultilabelMAP, + SemSegMIoU, + ) + + if isinstance(metric, MultilabelMAP): + scores = getattr(result, "raw_scores", None) + if scores is None: + return + gt = np.zeros(metric.num_labels, dtype=np.int64) + if 0 <= sample.label < metric.num_labels: + gt[sample.label] = 1 + metric.update(prediction=np.asarray(scores, dtype=np.float32), ground_truth=gt) + return + + if isinstance(metric, CocoDetectionMAP): + if sample.image_id is None: + return + if metric.iou_type == "bbox": + self._feed_bbox_predictions(metric, result, sample) + return + + if isinstance(metric, SemSegMIoU): + if sample.mask_path is None: + return + gt_raw = cv2.imread(str(sample.mask_path), cv2.IMREAD_GRAYSCALE) + if gt_raw is None: + return + # ADE20K convention: 0=unlabeled (ignore), 1..N=classes - shift down by 1. + gt_mask = gt_raw.astype(np.int32) - 1 + gt_mask[gt_mask < 0] = metric.ignore_index + pred_mask = getattr(result, "resultImage", result) + pred_mask = np.asarray(pred_mask) + if pred_mask.shape[:2] != gt_mask.shape[:2]: + pred_mask = cv2.resize( + pred_mask, + (gt_mask.shape[1], gt_mask.shape[0]), + interpolation=cv2.INTER_NEAREST, + ) + metric.update(pred_mask, gt_mask) + + @staticmethod + def _feed_bbox_predictions( + metric: "CocoDetectionMAP", + result: Any, + sample: "CalibrationSample", + ) -> None: + bboxes = getattr(result, "bboxes", None) + labels = getattr(result, "labels", None) + 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), + }, + ) + metric.update(predictions=preds) + + def quantize_model( + self, + model_path: Path, + calibration_data: list[np.ndarray], + model_config: dict[str, Any], + preset: str = "accuracy", + validation_data: list[np.ndarray] | None = None, + validation_labels: list[int] | None = None, + accuracy_results: AccuracyResults | None = None, + validation_samples: "list[CalibrationSample] | None" = None, + metric: "Metric | None" = None, + ) -> Path: + """Quantize OpenVINO model to INT8 using NNCF. + + Args: + model_path: Path to the FP32 OpenVINO model (.xml) + calibration_data: List of calibration images + model_config: Model configuration used for README rendering + preset: Quantization preset ('accuracy', 'performance', 'mixed') + validation_data: Optional validation images for accuracy measurement + validation_labels: Optional validation labels for accuracy measurement + accuracy_results: Optional collector populated with measured accuracies + and the INT8 success flag. + validation_samples: Optional per-image GT pointers consumed by + Model API-based metrics (multilabel mAP, COCO mAP, mIoU). + metric: Optional task-specific :class:`Metric`. When supplied and + not a :class:`TopOneAccuracy`, accuracy is measured via the + Model API path using ``validation_samples`` instead of the + preprocessed-tensor classification path. + + Returns: + Path to the quantized model + """ + if not calibration_data: + self.logger.warning("No calibration data provided. Skipping quantization.") + return model_path + + try: + import nncf + import openvino as ov + + self.logger.info(f"Quantizing model with {len(calibration_data)} calibration samples") + self.logger.info(f"Using preset: {preset}") + + # Load the model + core = ov.Core() + model = core.read_model(model_path) + + # Map preset string to NNCF enum + preset_map = { + "performance": nncf.QuantizationPreset.PERFORMANCE, + "mixed": nncf.QuantizationPreset.MIXED, + } + nncf_preset = preset_map.get(preset.lower(), nncf.QuantizationPreset.MIXED) + + quantize_kwargs: dict[str, Any] = {} + model_type = model_config.get("quantization_model_type") + if model_type and model_type.lower() == "transformer": + quantize_kwargs["model_type"] = nncf.ModelType.TRANSFORMER + + quantized_model = nncf.quantize( + model, + calibration_dataset=nncf.Dataset(calibration_data), + preset=nncf_preset, + subset_size=len(calibration_data), + **quantize_kwargs, + ) + + # Extract model name from the FP32 model path + model_name = model_path.stem + if model_name.endswith("_fp32"): + model_name = model_name[:-5] + + # Create output folder with -int8-ov suffix + output_folder = model_path.parent.parent / f"{model_name}-int8-ov" + output_folder.mkdir(parents=True, exist_ok=True) + + # Save quantized model with model name inside the folder + output_path = output_folder / f"{model_name}.xml" + ov.save_model(quantized_model, output_path, compress_to_fp16=True) + self.logger.info(f"✓ Quantized model saved: {output_path}") + if accuracy_results is not None: + 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) + + # Validate accuracy if validation data provided. + # Two paths: (1) Top-1 classification uses preprocessed validation_data + labels via raw + # OpenVINO; (2) other task metrics (multilabel mAP, COCO mAP, mIoU) iterate raw image + # paths through the Model API wrapper so per-task postprocessing is applied. + from model_converter.metrics import TopOneAccuracy + + fp16_model_path = model_path.parent / f"{model_name}.xml" + metric_path_active = ( + metric is not None and not isinstance(metric, TopOneAccuracy) and bool(validation_samples) + ) + if metric_path_active: + assert metric is not None # guaranteed by metric_path_active + assert validation_samples is not None # guaranteed by metric_path_active + metric_name = getattr(metric, "name", "metric") + self.logger.info(f"Validating FP32 model {metric_name}...") + fp32_metric = self._measure_metric(model_path, validation_samples, metric) + self.logger.info(f"FP32 {metric_name}: {fp32_metric}") + + fp16_metric: float | None = None + if fp16_model_path.exists(): + self.logger.info(f"Validating FP16 model {metric_name}...") + fp16_metric = self._measure_metric(fp16_model_path, validation_samples, metric) + self.logger.info(f"FP16 {metric_name}: {fp16_metric}") + else: + self.logger.warning(f"FP16 model not found for accuracy measurement: {fp16_model_path}") + + self.logger.info(f"Validating INT8 model {metric_name}...") + int8_metric = self._measure_metric(output_path, validation_samples, metric) + self.logger.info(f"INT8 {metric_name}: {int8_metric}") + + if accuracy_results is not None: + accuracy_results.fp32_accuracy = fp32_metric + accuracy_results.fp16_accuracy = fp16_metric + accuracy_results.int8_accuracy = int8_metric + accuracy_results.measured = True + elif validation_data and validation_labels: + self.logger.info("Validating FP32 model accuracy...") + fp32_accuracy = self.validate_model(model_path, validation_data, validation_labels) + self.logger.info(f"FP32 Top-1 Accuracy: {fp32_accuracy * 100:.2f}%") + + fp16_accuracy: float | None = None + if fp16_model_path.exists(): + self.logger.info("Validating FP16 model accuracy...") + fp16_accuracy = self.validate_model(fp16_model_path, validation_data, validation_labels) + self.logger.info(f"FP16 Top-1 Accuracy: {fp16_accuracy * 100:.2f}%") + else: + self.logger.warning(f"FP16 model not found for accuracy measurement: {fp16_model_path}") + + self.logger.info("Validating INT8 model accuracy...") + int8_accuracy = self.validate_model(output_path, validation_data, validation_labels) + self.logger.info(f"INT8 Top-1 Accuracy: {int8_accuracy * 100:.2f}%") + + accuracy_drop = (fp32_accuracy - int8_accuracy) * 100 + self.logger.info(f"Accuracy Drop: {accuracy_drop:.2f}%") + + if accuracy_results is not None: + accuracy_results.fp32_accuracy = fp32_accuracy + accuracy_results.fp16_accuracy = fp16_accuracy + accuracy_results.int8_accuracy = int8_accuracy + 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}") + + # Copy README for INT8 model + self.copy_readme( + model_config, + output_folder, + variant="int8", + ) + + return output_path + + except ImportError: + self.logger.error("NNCF not installed. Install with: pip install nncf") + return model_path + except (OSError, RuntimeError, TypeError, ValueError) as e: + self.logger.error(f"Failed to quantize model: {e}") + import traceback + + self.logger.debug(traceback.format_exc()) + return model_path + + @staticmethod + def _metadata_value(value: Any) -> str: + """Convert config values to Model API rt_info string values.""" + if isinstance(value, (list, tuple)): + return " ".join(str(item) for item in value) + return str(value) diff --git a/model_converter/src/model_converter/converters/getitune.py b/model_converter/src/model_converter/converters/getitune.py new file mode 100644 index 000000000..8134949ae --- /dev/null +++ b/model_converter/src/model_converter/converters/getitune.py @@ -0,0 +1,450 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""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 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. + + Invokes the export_pretrained_models.py script as a subprocess to export + models, then repackages the output and applies INT8 quantization. + """ + + def __init__(self, training_extensions_dir: Path | None = None, **kwargs: Any): + """Initialize GetituneConverter. + + Args: + training_extensions_dir: Path to cloned training_extensions repository + **kwargs: Arguments passed to BaseConverter + """ + super().__init__(**kwargs) + self.training_extensions_dir = Path(training_extensions_dir) if training_extensions_dir else None + + def process_model_config(self, config: dict[str, Any]) -> bool: + """Process a getitune model configuration. + + Runs the export_pretrained_models.py script, repackages the output + into the standard layout, and optionally quantizes to INT8. + + Args: + config: Model configuration dictionary + + Returns: + True if successful, False otherwise + """ + 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) + return True + + try: + if not self.training_extensions_dir: + error_msg = ( + "training_extensions_dir is required for getitune models. " + "Use --training-extensions-dir to specify the path." + ) + raise ValueError(error_msg) + + if not self.training_extensions_dir.exists(): + 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.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) + + # Export model using training_extensions script + exported_model_path = self._run_export(config) + + # Repackage into standard layout + self._repackage_model(config, exported_model_path) + + # Quantize if enabled and dataset is available + accuracy: AccuracyResults | None = None + quantization_attempted = bool(config.get("quantize", True) and config.get("dataset_type")) + 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, + accuracy=accuracy, + ) + + 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 + + def _run_export(self, config: dict[str, Any]) -> Path: + """Run the export_pretrained_models.py script. + + Invokes the script within the training_extensions/library uv-managed + environment so that all ``getitune`` dependencies are available. + + Args: + config: Model configuration dictionary + + Returns: + Path to the exported model XML file + """ + assert self.training_extensions_dir is not None + + getitune_task = config.get("getitune_task") + getitune_recipe = config.get("getitune_recipe") + + if not getitune_task: + error_msg = f"Model '{config.get('model_short_name')}' must define 'getitune_task'" + raise ValueError(error_msg) + if not getitune_recipe: + error_msg = f"Model '{config.get('model_short_name')}' must define 'getitune_recipe'" + raise ValueError(error_msg) + + # Create temporary directory for export output + temp_dir = tempfile.mkdtemp(prefix="getitune_export_") + temp_output = Path(temp_dir) + + export_script = (self.training_extensions_dir / "export_pretrained_models.py").resolve() + if not export_script.exists(): + error_msg = f"Export script not found: {export_script}" + raise FileNotFoundError(error_msg) + + library_dir = (self.training_extensions_dir / "library").resolve() + library_pyproject = library_dir / "pyproject.toml" + if not library_pyproject.exists(): + error_msg = ( + f"getitune library project not found at {library_dir}. " + f"Ensure training_extensions contains a 'library/' subdirectory with a pyproject.toml." + ) + raise FileNotFoundError(error_msg) + + cmd = [ + "uv", + "run", + "--project", + str(library_dir), + "--extra", + "cpu", + "python", + str(export_script), + "--task", + getitune_task, + "--model", + getitune_recipe, + "--output-dir", + str(temp_output), + "--format", + "OPENVINO", + "--precision", + "FP16", + ] + + self.logger.info(f"Running export command: {' '.join(cmd)}") + + result = subprocess.run( # noqa: S603 # nosec B603 — cmd is built from validated config, no shell + cmd, + cwd=str(self.training_extensions_dir), + capture_output=True, + text=True, + check=False, + ) + + if result.returncode != 0: + self.logger.error(f"Export script stderr: {result.stderr}") + error_msg = f"Export script failed with return code {result.returncode}: {result.stderr}" + raise RuntimeError(error_msg) + + self.logger.debug(f"Export script stdout: {result.stdout}") + + # Locate the exported model + # Expected path: temp_output/{task_lower}/{recipe}/exported_model.xml + expected_path = temp_output / getitune_task.lower() / getitune_recipe / "exported_model.xml" + + if not expected_path.exists(): + # Try searching for any .xml file in the output + xml_files = list(temp_output.rglob("*.xml")) + if xml_files: + expected_path = xml_files[0] + else: + error_msg = f"No exported model found in {temp_output}" + raise FileNotFoundError(error_msg) + + self.logger.info(f"Found exported model: {expected_path}") + return expected_path + + def _repackage_model(self, config: dict[str, Any], exported_model_path: Path) -> None: + """Repackage exported model into standard layout. + + Args: + config: Model configuration dictionary + exported_model_path: Path to the exported model XML + """ + model_short_name = config.get("model_short_name", "unknown") + + # Create output folder + output_folder = self.output_dir / f"{model_short_name}-fp16-ov" + output_folder.mkdir(parents=True, exist_ok=True) + + # Copy model files (XML + BIN) + target_xml = output_folder / f"{model_short_name}.xml" + shutil.copy2(exported_model_path, target_xml) + + bin_path = exported_model_path.with_suffix(".bin") + if bin_path.exists(): + shutil.copy2(bin_path, output_folder / f"{model_short_name}.bin") + + # Also save an FP32 copy for quantization + fp32_xml = output_folder / f"{model_short_name}_fp32.xml" + shutil.copy2(exported_model_path, fp32_xml) + if bin_path.exists(): + shutil.copy2(bin_path, output_folder / f"{model_short_name}_fp32.bin") + + self.logger.info(f"✓ Model repackaged to: {target_xml}") + + # Overwrite the exporter's placeholder labels with real class names. + self._apply_config_labels(config, target_xml, fp32_xml) + + # Extract and save model_info as config.json + try: + import openvino as ov + + 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) + 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") + + # Copy README + self.copy_readme(config, output_folder, variant="fp16") + + # Cleanup temp directory + temp_dir = exported_model_path.parent + while temp_dir.name != "getitune_export_" and "getitune_export_" not in temp_dir.name: + temp_dir = temp_dir.parent + if temp_dir == temp_dir.parent: + break + if "getitune_export_" in temp_dir.name: + shutil.rmtree(temp_dir, ignore_errors=True) + + def _apply_config_labels(self, config: dict[str, Any], *model_paths: Path) -> None: + """Overwrite ``model_info/labels`` rt_info with configured class names. + + The getitune exporter embeds only numeric placeholder ids and no + human-readable labels. When the configuration defines a ``labels`` set + (e.g. ``IMAGENET1K_V1``), resolve it to the real class names and rewrite + the rt_info of each given OpenVINO model in place so the repackaged + ``config.json`` (and any model quantized from these files) carries the + correct labels. + + Args: + config: Model configuration dictionary. + *model_paths: OpenVINO model XML files to update in place. + """ + labels_config = config.get("labels") + if not labels_config: + return + + labels = self.get_labels(labels_config) + if not labels: + self.logger.warning(f"Could not load labels for: {labels_config}") + return + + try: + import openvino as ov + + core = ov.Core() + for model_path in model_paths: + if not model_path.exists(): + continue + model = core.read_model(model_path) + model.set_rt_info(labels, ["model_info", "labels"]) + # Save to a temporary path first: OpenVINO memory-maps the + # source weights file, so writing back to the same path would + # truncate the file while it is still being read, corrupting + # the model (and crashing the process). + tmp_xml = model_path.with_name(f"{model_path.stem}_labeled_tmp.xml") + ov.save_model(model, tmp_xml, compress_to_fp16=not model_path.stem.endswith("_fp32")) + tmp_xml.replace(model_path) + tmp_xml.with_suffix(".bin").replace(model_path.with_suffix(".bin")) + self.logger.info(f"✓ Applied {labels_config} labels to exported model(s)") + except (ImportError, RuntimeError) as e: + self.logger.warning(f"Could not apply labels to exported model: {e}") + + @staticmethod + def _read_preprocessing_from_model( + ov: Any, + model_path: Path, + ) -> tuple[list[int], str, str, bool]: + """Read preprocessing parameters from the model's rt_info metadata. + + Args: + ov: The openvino module + model_path: Path to the OpenVINO model XML file + + Returns: + Tuple of (input_shape, mean_values, scale_values, reverse_input_channels) + """ + core = ov.Core() + model = core.read_model(model_path) + + # Get input shape from the model's input layer + input_shape = list(model.input(0).shape) + + # Read preprocessing params from model_info rt_info + def _get_rt_str(key: str, default: str) -> str: + try: + return model.get_rt_info(["model_info", key]).astype(str) + except RuntimeError: + return default + + mean_values = _get_rt_str("mean_values", "0 0 0") + scale_values = _get_rt_str("scale_values", "1 1 1") + reverse_input_channels = _get_rt_str("reverse_input_channels", "True").lower() in ("true", "1", "yes") + + return input_shape, mean_values, scale_values, reverse_input_channels + + def _quantize_exported_model(self, config: dict[str, Any]) -> AccuracyResults: + """Quantize the exported FP16 model to INT8. + + Reads preprocessing parameters (input_shape, mean_values, scale_values, + reverse_input_channels) from the exported model's rt_info metadata, + which is embedded by the getitune exporter. + + Args: + config: Model configuration dictionary + + Returns: + The accuracies measured during quantization and the INT8 success flag. + """ + import openvino as ov + + model_short_name = config.get("model_short_name", "unknown") + accuracy = AccuracyResults() + + # Use FP32 model for quantization (better quality) + fp32_model_path = self.output_dir / f"{model_short_name}-fp16-ov" / f"{model_short_name}_fp32.xml" + + if not fp32_model_path.exists(): + # Fall back to the FP16 model + fp32_model_path = self.output_dir / f"{model_short_name}-fp16-ov" / f"{model_short_name}.xml" + + # Extract preprocessing parameters from model rt_info + input_shape, mean_values, scale_values, reverse_input_channels = self._read_preprocessing_from_model( + ov, + fp32_model_path, + ) + + # Pick a metric strategy for this dataset/model_type combo. Top-1 + # uses the preprocessed-tensor classification path; other metrics + # (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 + + self.logger.info("Creating calibration dataset for INT8 quantization") + if is_top1: + self.logger.info("Creating validation dataset for accuracy measurement") + + calibration_data, validation_labels = self.create_calibration_dataset( + input_shape=input_shape, + mean_values=mean_values, + scale_values=scale_values, + reverse_input_channels=reverse_input_channels, + subset_size=500, + return_labels=is_top1, + dataset_path=dataset_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 + ) + + if calibration_data: + self.quantize_model( + model_path=fp32_model_path, + calibration_data=calibration_data, + model_config=config, + preset="mixed", + validation_data=calibration_data if validation_labels else None, + validation_labels=validation_labels or None, + validation_samples=validation_samples, + metric=metric, + accuracy_results=accuracy, + ) + + # 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}") + + return accuracy diff --git a/model_converter/src/model_converter/converters/pytorch.py b/model_converter/src/model_converter/converters/pytorch.py new file mode 100644 index 000000000..03f92b2a9 --- /dev/null +++ b/model_converter/src/model_converter/converters/pytorch.py @@ -0,0 +1,425 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""PyTorch-based converter shared by torchvision and timm converters.""" + +import importlib +import json +import shutil +from pathlib import Path +from typing import TYPE_CHECKING, 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", + "input_dtype", + "confidence_threshold", + "postprocess_semantic_masks", + "nms_execute", + "iou_threshold", + "agnostic_nms", + "nms_max_predictions", +) + + +class PyTorchConverter(BaseConverter): + """Shared converter for PyTorch-based models (torchvision, timm). + + Provides common export-to-OpenVINO logic, model loading utilities, + and quantization workflow. + """ + + def load_model_class(self, class_path: str) -> type: + """Dynamically load a model class from a Python path. + + Args: + class_path: Full Python path to the class (e.g., 'torchvision.models.resnet.resnet50') + + Returns: + The model class + """ + try: + module_path, class_name = class_path.rsplit(".", 1) + self.logger.debug(f"Importing module: {module_path}") + # nosemgrep: python.lang.security.audit.non-literal-import.non-literal-import + module = importlib.import_module(module_path) + model_class = getattr(module, class_name) + self.logger.debug(f"Loaded class: {class_name}") + return model_class + except Exception as e: + self.logger.error(f"Failed to import module {module_path}: {e}") + raise + + def load_checkpoint(self, checkpoint_path: Path) -> dict[str, Any]: + """Load PyTorch checkpoint file. + + Args: + checkpoint_path: Path to checkpoint file + + Returns: + Checkpoint dictionary + """ + try: + checkpoint = torch.load( # nosemgrep: trailofbits.python.pickles-in-pytorch.pickles-in-pytorch + checkpoint_path, + map_location="cpu", + weights_only=True, + ) + self.logger.debug(f"Loaded checkpoint from: {checkpoint_path}") + return checkpoint + except Exception as e: + self.logger.error(f"Failed to load checkpoint: {e}") + raise + + def create_model( + self, + model_class: type, + checkpoint: dict[str, Any], + model_params: dict[str, Any] | None = None, + ) -> nn.Module: + """Create and initialize model instance. + + Args: + model_class: Model class to instantiate + checkpoint: Checkpoint containing model weights + model_params: Optional parameters for model initialization + + Returns: + Initialized model instance + """ + try: + if model_class == torch.nn.Module: + if "model" in checkpoint: + model = checkpoint["model"] + elif "state_dict" in checkpoint: + error_msg = ( + "Checkpoint contains only state_dict. Please specify the model class instead of torch.nn.Module" + ) + raise ValueError(error_msg) + else: + model = checkpoint + + if not isinstance(model, nn.Module): + error_msg = "Checkpoint does not contain a valid model" + raise ValueError(error_msg) + else: + model = model_class(**model_params) if model_params else model_class() + + if "state_dict" in checkpoint: + state_dict = checkpoint["state_dict"] + elif "model" in checkpoint: + if isinstance(checkpoint["model"], nn.Module): + return checkpoint["model"] + state_dict = checkpoint["model"] + else: + state_dict = checkpoint + + model.load_state_dict(state_dict, strict=False) + + model.eval() + self.logger.info("✓ Model created and loaded successfully") + return model + + except Exception as e: + self.logger.error(f"Failed to create model: {e}") + raise + + def export_to_openvino( + self, + model: nn.Module, + input_shape: list[int], + output_path: Path, + model_config: dict[str, Any], + input_names: list[str] | None = None, + output_names: list[str] | None = None, + metadata: dict[tuple[str, str], str] | None = None, + ) -> tuple[Path, Path]: + """Export PyTorch model to OpenVINO format. + + Args: + model: PyTorch model to export + input_shape: Input tensor shape [batch, channels, height, width] + output_path: Path to save the model (without extension) + model_config: Model configuration used for README rendering + input_names: Names for input tensors + output_names: Names for output tensors + metadata: Metadata to embed in the model + + Returns: + Tuple of (fp16_model_path, fp32_model_path) - FP16 for final use, FP32 for quantization + """ + import openvino as ov + + try: + model = self._prepare_model_for_export(model, model_config) + model.eval() + dummy_input = self._create_example_input(input_shape, model_config) + self.logger.info("Direct PyTorch to OpenVINO conversion") + ov_model = ov.convert_model(model, example_input=dummy_input) + self.logger.info("✓ PyTorch to OpenVINO conversion complete") + + # Reshape model to fixed input shape (remove dynamic dimensions) + first_input = ov_model.input(0) + input_name_for_reshape = next(iter(first_input.get_names())) if first_input.get_names() else 0 + + self.logger.debug(f"Setting fixed input shape: {input_shape}") + ov_model.reshape({input_name_for_reshape: input_shape}) + + # Post-process the model + ov_model = self._postprocess_openvino_model( + ov_model, + input_names=input_names, + output_names=output_names, + metadata=metadata, + ) + + # Create output folder with -fp16-ov suffix + model_name = output_path.name + output_folder = output_path.parent / f"{model_name}-fp16-ov" + output_folder.mkdir(parents=True, exist_ok=True) + + # Save FP32 model for quantization (temporary) + fp32_xml_path = output_folder / f"{model_name}_fp32.xml" + ov.save_model(ov_model, fp32_xml_path, compress_to_fp16=False) + self.logger.debug(f"Saved FP32 model for quantization: {fp32_xml_path}") + + # Save the FP16 model (final) + xml_path = output_folder / f"{model_name}.xml" + ov.save_model(ov_model, xml_path, compress_to_fp16=True) + 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) + + # 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}") + + # Copy README for FP16 model + self.copy_readme( + model_config, + output_folder, + variant="fp16", + ) + + return xml_path, fp32_xml_path + + except Exception as e: + self.logger.error(f"Failed to export model: {e}") + raise + + def _prepare_model_for_export(self, model: nn.Module, model_config: dict[str, Any]) -> nn.Module: + """Prepare model for OpenVINO conversion.""" + model_type = str(model_config.get("model_type", "")) + adapted = get_adapter(model_type, model) + if adapted is not model: + self.logger.info(f"Applied export adapter for model type: {model_type}") + return adapted + + def _create_example_input(self, input_shape: list[int], model_config: dict[str, Any]) -> torch.Tensor: + """Create example input suitable for the configured model type.""" + if str(model_config.get("model_type", "")).lower() == "maskrcnn": + return torch.rand(*input_shape) + return torch.randn(*input_shape) + + def _postprocess_openvino_model( + self, + model: Any, + input_names: list[str] | None = None, + output_names: list[str] | None = None, + metadata: dict[tuple[str, str], str] | None = None, + ) -> Any: + """Post-process OpenVINO model (set names, add metadata). + + Args: + model: OpenVINO model + input_names: Names for input tensors + output_names: Names for output tensors + metadata: Metadata to embed + + Returns: + Post-processed model + """ + if input_names: + for i, name in enumerate(input_names): + if i < len(model.inputs): + model.input(i).set_names({name}) + self.logger.debug(f"Set input {i} name to: {name}") + + if output_names: + for i, name in enumerate(output_names): + if i < len(model.outputs): + model.output(i).set_names({name}) + self.logger.debug(f"Set output {i} name to: {name}") + + if metadata: + for key, value in metadata.items(): + model.set_rt_info(value, list(key)) + self.logger.debug(f"Set metadata {key}: {value}") + + return model + + def _build_metadata(self, config: dict[str, Any]) -> dict[tuple[str, str], str]: + """Build metadata dictionary from model config.""" + model_short_name = config.get("model_short_name", "unknown") + 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", "") + + metadata = { + ("model_info", "model_type"): model_type, + ("model_info", "model_short_name"): model_short_name, + ("model_info", "reverse_input_channels"): self._metadata_value(reverse_input_channels), + ("model_info", "mean_values"): mean_values, + ("model_info", "scale_values"): scale_values, + } + + for metadata_field in _MODEL_API_METADATA_FIELDS: + if metadata_field in config and config[metadata_field] is not None: + metadata["model_info", metadata_field] = self._metadata_value(config[metadata_field]) + + # Add labels if specified in config + labels_config = config.get("labels") + if labels_config: + labels = self.get_labels(labels_config) + if labels: + metadata["model_info", "labels"] = labels + self.logger.info(f"Added {labels_config} labels to metadata") + else: + self.logger.warning(f"Could not load labels for: {labels_config}") + + return metadata + + def validate_torch_model( + self, + model: nn.Module, + validation_data: list[Any], + labels: list[int], + ) -> float | None: + """Validate the original PyTorch model and compute top-1 accuracy. + + Runs inference on the same preprocessed validation tensors used for the + OpenVINO models, so the result is comparable to the FP32/FP16/INT8 + accuracies. + + Args: + model: The original PyTorch model (before OpenVINO conversion). + validation_data: Preprocessed validation images (NCHW numpy arrays). + labels: Ground truth class labels. + + Returns: + Top-1 accuracy (0.0 to 1.0), or ``None`` if validation failed. + """ + try: + model.eval() + predictions: list[int] = [] + with torch.no_grad(): + for img in validation_data: + output = model(torch.from_numpy(img).float()) + if isinstance(output, (tuple, list)): + output = output[0] + predictions.append(int(torch.argmax(output, dim=1)[0].item())) + + correct = sum(predicted == label for predicted, label in zip(predictions, labels)) + return correct / len(labels) + + except (RuntimeError, TypeError, ValueError) as e: + self.logger.error(f"Failed to validate PyTorch model: {e}") + return None + + def _quantize_and_cleanup(self, config: dict[str, Any], fp32_model_path: Path, **kwargs: Any) -> AccuracyResults: + """Run INT8 quantization and clean up temporary FP32 model files. + + Returns: + The accuracies measured during quantization and the INT8 success flag. + """ + model_type = kwargs["model_type"] + accuracy = AccuracyResults() + self.logger.info("Creating calibration dataset for INT8 quantization") + # Pick a metric strategy. Top-1 uses the preprocessed-tensor 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 + resize_type = config.get("resize_type", "standard") + + if is_top1: + self.logger.info("Creating validation dataset for accuracy measurement") + validation_data, validation_labels = self.create_calibration_dataset( + input_shape=kwargs["input_shape"], + mean_values=kwargs["mean_values"], + scale_values=kwargs["scale_values"], + reverse_input_channels=kwargs["reverse_input_channels"], + subset_size=500, + return_labels=is_top1, + resize_type=resize_type, + dataset_path=dataset_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 + ) + + if validation_data: + torch_model = kwargs.get("torch_model") + if validation_labels and torch_model is not None: + self.logger.info("Validating original PyTorch model accuracy...") + original_accuracy = self.validate_torch_model(torch_model, validation_data, validation_labels) + if original_accuracy is not None: + self.logger.info(f"Original Top-1 Accuracy: {original_accuracy * 100:.2f}%") + accuracy.original_accuracy = original_accuracy + accuracy.measured = True + + self.quantize_model( + model_path=fp32_model_path, + calibration_data=validation_data, + model_config=config, + preset="mixed", + validation_data=validation_data if validation_labels else None, + validation_labels=validation_labels or None, + validation_samples=validation_samples, + metric=metric, + accuracy_results=accuracy, + ) + + # 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}") + + return accuracy diff --git a/model_converter/src/model_converter/converters/registry.py b/model_converter/src/model_converter/converters/registry.py new file mode 100644 index 000000000..afd981842 --- /dev/null +++ b/model_converter/src/model_converter/converters/registry.py @@ -0,0 +1,19 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Converter registry mapping model_library names to converter classes.""" + +from model_converter.converters.base import BaseConverter +from model_converter.converters.getitune import GetituneConverter +from model_converter.converters.timm import TimmConverter +from model_converter.converters.torchvision import TorchvisionConverter +from model_converter.converters.yolo import YoloConverter + +CONVERTER_REGISTRY: dict[str, type[BaseConverter]] = { + "torchvision": TorchvisionConverter, + "timm": TimmConverter, + "yolo": YoloConverter, + "getitune": GetituneConverter, +} diff --git a/model_converter/src/model_converter/converters/timm.py b/model_converter/src/model_converter/converters/timm.py new file mode 100644 index 000000000..0cd0e306e --- /dev/null +++ b/model_converter/src/model_converter/converters/timm.py @@ -0,0 +1,251 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Timm (HuggingFace) model converter.""" + +from typing import TYPE_CHECKING, Any + +import torch.nn as nn + +from model_converter.converters.pytorch import PyTorchConverter + +if TYPE_CHECKING: + from model_converter.reporting import AccuracyResults + + +class TimmConverter(PyTorchConverter): + """Converter for timm models hosted on HuggingFace Hub. + + Loads models via timm/transformers from HuggingFace, + exports to OpenVINO, and optionally quantizes to INT8. + """ + + def load_huggingface_model( + self, + repo_id: str, + revision: str, + model_library: str = "timm", + model_params: dict[str, Any] | None = None, + ) -> nn.Module: + """Load a model from Hugging Face Hub. + + Args: + repo_id: Hugging Face repository ID + revision: Immutable revision/commit SHA for the Hugging Face repository + model_library: Library to use ('timm', 'transformers', etc.) + model_params: Optional parameters for model loading + + Returns: + Loaded model instance + """ + try: + if model_library == "timm": + import timm + + repo_ref = f"hf-hub:{repo_id}@{revision}" + self.logger.info(f"Loading timm model: {repo_ref}") + model = timm.create_model( + repo_ref, + pretrained=True, + cache_dir=self.cache_dir, + **(model_params or {}), + ) + elif model_library == "transformers": + from transformers import AutoModel + + self.logger.info(f"Loading transformers model: {repo_id}@{revision}") + model = AutoModel.from_pretrained( + repo_id, + revision=revision, + cache_dir=self.cache_dir, + **(model_params or {}), + ) + else: + error_msg = f"Unsupported model library: {model_library}" + raise ValueError(error_msg) + + model.eval() + self.logger.info("✓ Hugging Face model loaded successfully") + return model + + except Exception as e: + self.logger.error(f"Failed to load Hugging Face model: {e}") + raise + + def _apply_timm_data_config(self, model: nn.Module, config: dict[str, Any]) -> None: + """Override preprocessing config with timm's canonical values. + + Reads the model's ``pretrained_cfg`` via :func:`timm.data.resolve_data_config` + and updates ``mean_values``, ``scale_values`` and ``input_shape`` in + ``config`` in place. timm stores ``mean``/``std`` as 0..1 floats, so they + are scaled to the 0..255 pixel range used by Model API metadata. + + ``reverse_input_channels`` is forced to ``True`` because timm models are + trained on RGB images while images are decoded as BGR by OpenCV, so the + channels must always be swapped. timm does not expose this through + ``resolve_data_config``, but it is an invariant for these models. + + Args: + model: The loaded timm model. + config: Model configuration dictionary, mutated in place. + """ + # timm models consume RGB; OpenCV decodes BGR, so the swap is always needed. + self._override_config_value(config, "reverse_input_channels", resolved=True) + + try: + from timm.data import resolve_data_config + except ImportError: + self.logger.warning("timm not available; keeping configured preprocessing values") + return + + try: + data_config = resolve_data_config({}, model=model) + except (RuntimeError, ValueError, KeyError, TypeError) as e: + self.logger.warning(f"Could not resolve timm data config, keeping configured values: {e}") + return + + mean = data_config.get("mean") + std = data_config.get("std") + input_size = data_config.get("input_size") + + if mean is not None: + resolved_mean = " ".join(f"{value * 255:g}" for value in mean) + self._override_config_value(config, "mean_values", resolved_mean) + if std is not None: + resolved_scale = " ".join(f"{value * 255:g}" for value in std) + self._override_config_value(config, "scale_values", resolved_scale) + if input_size is not None and len(input_size) == 3: + channels, height, width = input_size + resolved_shape = [1, int(channels), int(height), int(width)] + self._override_config_value(config, "input_shape", resolved_shape) + + def _override_config_value(self, config: dict[str, Any], key: str, resolved: Any) -> None: + """Set ``config[key]`` to ``resolved``, logging when it changes a value.""" + existing = config.get(key) + if existing is not None and existing != resolved: + self.logger.info(f"Overriding {key}: config={existing!r} -> timm={resolved!r}") + config[key] = resolved + + def process_model_config(self, config: dict[str, Any]) -> bool: + """Process a timm/HuggingFace model configuration. + + Args: + config: Model configuration dictionary + + Returns: + True if successful, False otherwise + """ + 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) + 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) + + # Load model from HuggingFace + huggingface_repo = config.get("huggingface_repo") + huggingface_revision = config.get("huggingface_revision") + + if not huggingface_repo: + error_msg = f"Timm model '{model_short_name}' must define 'huggingface_repo'" + raise ValueError(error_msg) + if not huggingface_revision: + error_msg = "Hugging Face models must define 'huggingface_revision' with an immutable commit SHA" + raise ValueError(error_msg) + + model_library = config.get("model_library", "timm") + model_params = config.get("model_params") + model = self.load_huggingface_model( + repo_id=huggingface_repo, + revision=huggingface_revision, + model_library=model_library, + model_params=model_params, + ) + + # Override preprocessing parameters with the values timm ships for + # this specific checkpoint. Hand-maintained config values are easy + # to get wrong (e.g. ImageNet vs. inception normalization), which + # silently destroys accuracy, so timm's pretrained_cfg is treated as + # the source of truth. + if model_library == "timm": + 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", "") + + metadata = self._build_metadata(config) + + output_path = self.output_dir / model_short_name + fp16_model_path, fp32_model_path = self.export_to_openvino( + model=model, + input_shape=input_shape, + output_path=output_path, + model_config=config, + input_names=input_names, + output_names=output_names, + metadata=metadata, + ) + + # Quantize the model if dataset is available + accuracy: AccuracyResults | None = None + quantization_attempted = bool(config.get("quantize", True) and config.get("dataset_type")) + if quantization_attempted: + 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, + 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, + accuracy=accuracy, + ) + + 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 diff --git a/model_converter/src/model_converter/converters/torchvision.py b/model_converter/src/model_converter/converters/torchvision.py new file mode 100644 index 000000000..d351544f6 --- /dev/null +++ b/model_converter/src/model_converter/converters/torchvision.py @@ -0,0 +1,133 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Torchvision model converter.""" + +from typing import TYPE_CHECKING, Any + +from model_converter.converters.pytorch import PyTorchConverter +from model_converter.downloaders import URLDownloader + +if TYPE_CHECKING: + from model_converter.reporting import AccuracyResults + + +class TorchvisionConverter(PyTorchConverter): + """Converter for torchvision models. + + Downloads weights from URL, loads the model class dynamically, + exports to OpenVINO, and optionally quantizes to INT8. + """ + + def __init__(self, **kwargs: Any): + """Initialize TorchvisionConverter.""" + super().__init__(**kwargs) + self._url_downloader = URLDownloader(cache_dir=self.cache_dir) + + def process_model_config(self, config: dict[str, Any]) -> bool: + """Process a torchvision model configuration. + + Args: + config: Model configuration dictionary + + Returns: + True if successful, False otherwise + """ + 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) + 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) + + # Download weights and load model + weights_url = config["weights_url"] + weights_path = self._url_downloader.download(url=weights_url) + + model_class_name = config.get("model_class_name", "torch.nn.Module") + model_class = self.load_model_class(model_class_name) + + checkpoint = self.load_checkpoint(weights_path) + + model_params = config.get("model_params") + 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", "") + + metadata = self._build_metadata(config) + + output_path = self.output_dir / model_short_name + fp16_model_path, fp32_model_path = self.export_to_openvino( + model=model, + input_shape=input_shape, + output_path=output_path, + model_config=config, + input_names=input_names, + output_names=output_names, + metadata=metadata, + ) + + # Quantize the model if dataset is available + accuracy: AccuracyResults | None = None + quantization_attempted = bool(config.get("quantize", True) and config.get("dataset_type")) + if quantization_attempted: + 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, + 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, + accuracy=accuracy, + ) + + 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 diff --git a/model_converter/src/model_converter/converters/yolo.py b/model_converter/src/model_converter/converters/yolo.py new file mode 100644 index 000000000..1b8010068 --- /dev/null +++ b/model_converter/src/model_converter/converters/yolo.py @@ -0,0 +1,298 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""YOLO model converter.""" + +import json +import shutil +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import cv2 +from defusedxml.ElementTree import ParseError, parse + +from model_converter.converters.base import BaseConverter +from model_converter.metrics.coco_detection import COCO80_TO_COCO91, CocoDetectionMAP +from model_converter.reporting import AccuracyResults + +if TYPE_CHECKING: + from model_converter.datasets import CalibrationSample + +MODEL_VERSIONS = ["yolo11n", "yolo11s", "yolo11m", "yolo11l", "yolo11x"] + + +# Re-export for backward compatibility (the canonical definition lives in metrics/coco_detection.py). +_COCO80_TO_COCO91 = COCO80_TO_COCO91 + + +class YoloConverter(BaseConverter): + """Converter for Ultralytics YOLO models. + + Uses the Ultralytics library to export YOLO models to OpenVINO format, + then repackages and adds metadata. + """ + + def process_model_config(self, config: dict[str, Any]) -> bool: + """Process a YOLO model configuration. + + Args: + config: Model configuration dictionary with 'yolo_version' field + + Returns: + True if successful, False otherwise + """ + model_short_name = config.get("model_short_name", "unknown") + yolo_version = config.get("yolo_version", model_short_name) + + # Check if both FP16 and INT8 models already exist + 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) + return True + + try: + from ultralytics import YOLO + + self.logger.info("=" * 80) + self.logger.info(f"Processing YOLO model: {model_short_name}") + self.logger.info("=" * 80) + + yolo_size = yolo_version[-1] # n, s, m, l, or x + + # Load model from cache directory to avoid polluting the working directory + model_path = self.cache_dir / f"{yolo_version}.pt" + model = YOLO(str(model_path)) + + # Export regular OpenVINO model (FP16) + self.logger.info(f"Exporting {yolo_version} to OpenVINO FP16 format...") + model.export(format="openvino", half=True) + + # Ultralytics exports next to the .pt file in the cache directory + old_name = self.cache_dir / f"{yolo_version}_openvino_model" + fp16_folder.mkdir(parents=True, exist_ok=True) + if old_name.exists(): + if fp16_folder.exists(): + shutil.rmtree(fp16_folder) + old_name.rename(fp16_folder) + + # Update model_type in XML metadata + xml_file = fp16_folder / f"{yolo_version}.xml" + if xml_file.exists(): + self._update_model_type_in_xml(xml_file, "YOLO11") + + # Copy README template for fp16 + self._copy_yolo_readme("README-yolo-fp16.md", fp16_folder, yolo_size) + + # Export INT8 quantized model + self.logger.info(f"Exporting {yolo_version} to OpenVINO INT8 format...") + model.export(format="openvino", int8=True, data="coco128.yaml") + + # Rename output folder for INT8 + old_name_int8 = self.cache_dir / f"{yolo_version}_int8_openvino_model" + if old_name_int8.exists(): + if int8_folder.exists(): + shutil.rmtree(int8_folder) + old_name_int8.rename(int8_folder) + + # Update model_type in XML metadata + xml_file = int8_folder / f"{yolo_version}.xml" + if xml_file.exists(): + self._update_model_type_in_xml(xml_file, "YOLO11") + + # Copy README template for int8 + self._copy_yolo_readme("README-yolo-int8.md", int8_folder, yolo_size) + + self.logger.info(f"✓ Successfully converted {model_short_name}") + quantized = (int8_folder / f"{yolo_version}.xml").exists() + accuracy: AccuracyResults | None = None + if self.measure_accuracy and quantized: + accuracy = self._measure_yolo_accuracy(config, yolo_version, fp16_folder, int8_folder) + self._record_result(self._build_result(config), converted=True, quantized=quantized, accuracy=accuracy) + 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 + + def _measure_yolo_accuracy( + self, + config: dict[str, Any], + yolo_version: str, + fp16_folder: Path, + int8_folder: Path, + ) -> AccuracyResults | None: + """Measure original PT model, FP16 OV, and INT8 OV mAP on the COCO validation subset. + + Uses the same 500-image COCO subset for all three measurements so the + numbers in the report are directly comparable. + + Args: + config: Model configuration dictionary (must contain ``dataset_type``). + yolo_version: Ultralytics model identifier (e.g. ``"yolo11n"``). + fp16_folder: Directory containing the exported FP16 OpenVINO model. + int8_folder: Directory containing the exported INT8 OpenVINO model. + + Returns: + Populated :class:`AccuracyResults`, or ``None`` when the dataset or + metric is unavailable. + """ + from model_converter.datasets.factory import _COCO_ANNOTATION_FILES + + dataset_path = self._resolve_dataset_path(config) + if dataset_path is None or not dataset_path.exists(): + self.logger.warning("COCO dataset not available — skipping accuracy measurement for YOLO") + return None + + dataset_type = config.get("dataset_type") + if dataset_type not in _COCO_ANNOTATION_FILES: + self.logger.warning(f"Unsupported dataset_type {dataset_type!r} — skipping accuracy measurement") + return None + + annotation_file = dataset_path / "annotations" / _COCO_ANNOTATION_FILES[dataset_type] + if not annotation_file.exists(): + self.logger.warning(f"COCO annotation file not found: {annotation_file}") + return None + + samples = self._collect_validation_samples(dataset_path, dataset_type, subset_size=500) + if not samples: + self.logger.warning("No validation samples found — skipping accuracy measurement for YOLO") + return None + + metric = self._metric_for_config(config, dataset_path) + if metric is None: + self.logger.warning("No metric available for this config — skipping accuracy measurement for YOLO") + return None + + accuracy = AccuracyResults() + accuracy.metric_name = metric.name + + # Original PT model accuracy — run Ultralytics native inference on each sample. + pt_model_path = self.cache_dir / f"{yolo_version}.pt" + original_map = self._measure_original_accuracy(pt_model_path, samples, annotation_file) + accuracy.original_accuracy = original_map + + # FP16 OV model accuracy via Model API. + fp16_model_path = fp16_folder / f"{yolo_version}.xml" + if fp16_model_path.exists(): + self.logger.info("Measuring FP16 model mAP...") + accuracy.fp16_accuracy = self._measure_metric(fp16_model_path, samples, metric) + self.logger.info(f"FP16 mAP: {accuracy.fp16_accuracy}") + + # INT8 OV model accuracy via Model API. + int8_model_path = int8_folder / f"{yolo_version}.xml" + if int8_model_path.exists(): + self.logger.info("Measuring INT8 model mAP...") + accuracy.int8_accuracy = self._measure_metric(int8_model_path, samples, metric) + self.logger.info(f"INT8 mAP: {accuracy.int8_accuracy}") + + accuracy.measured = True + return accuracy + + def _measure_original_accuracy( + self, + pt_model_path: Path, + samples: "list[CalibrationSample]", + annotation_file: Path, + ) -> float | None: + """Measure mAP of the original YOLO PT model using direct Ultralytics inference. + + Runs the ``.pt`` model on each COCO sample, converts Ultralytics 0-79 + class indices to COCO 91-class category IDs via :data:`COCO80_TO_COCO91`, + and evaluates with :class:`CocoDetectionMAP`. + + Args: + pt_model_path: Path to the Ultralytics ``.pt`` weights file. + samples: COCO validation samples with ``image_path`` and ``image_id``. + annotation_file: Path to the COCO ``instances_val2017.json`` file. + + Returns: + mAP@IoU=0.50:0.95, or ``None`` on failure. + """ + try: + from ultralytics import YOLO + + model = YOLO(str(pt_model_path)) + except (ImportError, FileNotFoundError, RuntimeError) as e: + self.logger.error(f"Failed to load PT model for original accuracy measurement: {e}") + return None + + metric = CocoDetectionMAP(annotation_file=annotation_file, iou_type="bbox") + predictions: list[dict[str, Any]] = [] + + for sample in samples: + if sample.image_id is None: + continue + img = cv2.imread(str(sample.image_path)) + if img is None: + continue + try: + results = model(img, verbose=False)[0] + except (RuntimeError, TypeError, ValueError) as e: + self.logger.debug(f"PT inference failed for {sample.image_path}: {e}") + continue + + boxes_xyxy = results.boxes.xyxy.cpu().tolist() + cls_ids = results.boxes.cls.cpu().tolist() + 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), + }, + ) + + metric.update(predictions=predictions) + original_map = float(metric.compute()) + self.logger.info(f"Original PT model mAP: {original_map}") + return original_map + + def _copy_yolo_readme(self, template_name: str, dest_dir: Path, yolo_size: str) -> None: + """Copy a YOLO README template, replacing <>.""" + template_path = Path(__file__).parent.parent / "templates" / template_name + if not template_path.exists(): + self.logger.warning(f"YOLO README template not found: {template_path}") + return + content = template_path.read_text() + content = content.replace("<>", yolo_size) + (dest_dir / "README.md").write_text(content) + self.logger.debug(f"Copied {template_name} -> {dest_dir / 'README.md'} (size={yolo_size})") + + def _update_model_type_in_xml(self, xml_path: Path, model_type: str = "YOLO11") -> None: + """Update the model_type value in the OpenVINO XML file.""" + try: + tree = parse(xml_path) + root = tree.getroot() + + for rt_info in root.findall(".//rt_info"): + for model_info in rt_info.findall(".//model_info"): + for model_type_elem in model_info.findall(".//model_type"): + model_type_elem.set("value", model_type) + + tree.write(xml_path, encoding="utf-8", xml_declaration=True) + self.logger.debug(f"Updated model_type to {model_type} in {xml_path}") + + model_info_dict = {} + for rt_info in root.findall(".//rt_info"): + for model_info in rt_info.findall(".//model_info"): + for child in model_info: + model_info_dict[child.tag] = child.attrib["value"] + with (xml_path.parent / "config.json").open("w") as f: + json.dump(model_info_dict, f, indent=4) + + except (OSError, ParseError) as error: + self.logger.warning(f"Failed to update {xml_path}: {error}") diff --git a/model_converter/src/model_converter/dataset_registry.py b/model_converter/src/model_converter/dataset_registry.py new file mode 100644 index 000000000..bbe4ccaa2 --- /dev/null +++ b/model_converter/src/model_converter/dataset_registry.py @@ -0,0 +1,165 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Dataset registry for managing calibration dataset paths.""" + +import json +import logging +from pathlib import Path +from typing import Any + + +class DatasetRegistry: + """Registry for mapping dataset types to local filesystem paths. + + Loads a JSON configuration file that maps dataset type identifiers + (e.g., "imagenet-1k", "coco-detection") to local directory paths. + Provides validation and helpful error messages when datasets are missing. + + Example datasets.json format: + { + "datasets": { + "imagenet-1k": "/path/to/imagenet/validation", + "imagenet-21k": "/path/to/imagenet21k/validation", + "coco-detection": "/path/to/coco2017/val" + } + } + """ + + def __init__(self, config_path: Path): + """Initialize the DatasetRegistry from a JSON configuration file. + + Args: + config_path: Path to the datasets configuration JSON file + + Raises: + FileNotFoundError: If the config file doesn't exist + ValueError: If the JSON is invalid or missing required structure + """ + self.config_path = Path(config_path) + self.logger = logging.getLogger(__name__) + + if not self.config_path.exists(): + error_msg = f"Dataset configuration file not found: {self.config_path}" + raise FileNotFoundError(error_msg) + + try: + with self.config_path.open() as f: + config = json.load(f) + except json.JSONDecodeError as e: + error_msg = f"Invalid JSON in dataset configuration file {self.config_path}: {e}" + raise ValueError(error_msg) from e + except (OSError, PermissionError) as e: + error_msg = f"Failed to read dataset configuration file {self.config_path}: {e}" + raise ValueError(error_msg) from e + + if not isinstance(config, dict) or "datasets" not in config: + error_msg = ( + f"Invalid dataset configuration format in {self.config_path}. " + 'Expected JSON with "datasets" key: {"datasets": {...}}' + ) + raise ValueError(error_msg) + + self._datasets: dict[str, Path] = {} + datasets = config["datasets"] + + if not isinstance(datasets, dict): + error_msg = ( + f"Invalid datasets format in {self.config_path}. Expected dictionary mapping dataset types to paths." + ) + raise ValueError(error_msg) + + # Convert string paths to Path objects + for dataset_type, path in datasets.items(): + if not isinstance(dataset_type, str) or not isinstance(path, str): + error_msg = ( + f"Invalid entry in datasets configuration: {dataset_type} -> {path}. " + "Both type and path must be strings." + ) + raise ValueError(error_msg) + self._datasets[dataset_type] = Path(path) + + self.logger.info(f"Loaded dataset registry with {len(self._datasets)} dataset types") + self.logger.debug(f"Available dataset types: {list(self._datasets.keys())}") + + def get_path(self, dataset_type: str, *, validate_exists: bool = False) -> Path: + """Get the filesystem path for a dataset type. + + Args: + dataset_type: Dataset type identifier (e.g., "imagenet-1k") + validate_exists: If True, verify the path exists on the filesystem + + Returns: + Path to the dataset directory + + Raises: + ValueError: If the dataset type is not registered + FileNotFoundError: If validate_exists=True and path doesn't exist + """ + if dataset_type not in self._datasets: + available = ", ".join(sorted(self._datasets.keys())) + error_msg = ( + f"Dataset type '{dataset_type}' not found in registry. " + f"Available types: {available if available else '(none)'}" + ) + raise ValueError(error_msg) + + path = self._datasets[dataset_type] + + if validate_exists and not path.exists(): + error_msg = ( + f"Dataset path for type '{dataset_type}' does not exist: {path}. " + f"Please verify the path in {self.config_path}" + ) + raise FileNotFoundError(error_msg) + + return path + + def has_type(self, dataset_type: str) -> bool: + """Check if a dataset type is registered. + + Args: + dataset_type: Dataset type identifier to check + + Returns: + True if the type is registered, False otherwise + """ + return dataset_type in self._datasets + + def list_types(self) -> list[str]: + """Get a list of all registered dataset types. + + Returns: + Sorted list of dataset type identifiers + """ + return sorted(self._datasets.keys()) + + def resolve_from_config(self, model_config: dict[str, Any]) -> Path | None: + """Resolve dataset path from a model configuration. + + Extracts the 'dataset_type' field from the model config and + returns the corresponding path. Returns None if no dataset_type + is specified (model doesn't require calibration dataset). + + Args: + model_config: Model configuration dictionary + + Returns: + Path to the dataset, or None if no dataset_type specified + + Raises: + ValueError: If dataset_type is specified but not in registry + """ + dataset_type = model_config.get("dataset_type") + + if dataset_type is None: + return None + + if not isinstance(dataset_type, str) or not dataset_type.strip(): + model_name = model_config.get("model_short_name", "unknown") + error_msg = f"Model '{model_name}' has invalid dataset_type: {dataset_type}" + raise ValueError(error_msg) + + return self.get_path(dataset_type.strip()) diff --git a/model_converter/src/model_converter/datasets/__init__.py b/model_converter/src/model_converter/datasets/__init__.py new file mode 100644 index 000000000..7ed45ad04 --- /dev/null +++ b/model_converter/src/model_converter/datasets/__init__.py @@ -0,0 +1,33 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Calibration/validation dataset readers. + +Provides a small abstraction over the on-disk layouts used by the +converter for different task families (classification, detection, +segmentation, etc.). + +Each reader yields :class:`CalibrationSample` objects with at least an +``image_path``. Readers for classification datasets also populate the +``label`` field with the integer class id; readers for other tasks emit a +placeholder ``0`` because per-image ground truth for those tasks is +consumed via task-specific metric modules (Phase 4), not via the +calibration loader. +""" + +from .ade20k import Ade20kReader +from .base import CalibrationSample, DatasetReader +from .class_folder import ClassFolderReader +from .coco import CocoImagesReader +from .factory import reader_for + +__all__ = [ + "Ade20kReader", + "CalibrationSample", + "ClassFolderReader", + "CocoImagesReader", + "DatasetReader", + "reader_for", +] diff --git a/model_converter/src/model_converter/datasets/ade20k.py b/model_converter/src/model_converter/datasets/ade20k.py new file mode 100644 index 000000000..18664e186 --- /dev/null +++ b/model_converter/src/model_converter/datasets/ade20k.py @@ -0,0 +1,45 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""ADE20K semantic-segmentation dataset reader.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from .base import CalibrationSample, DatasetReader + +if TYPE_CHECKING: + from collections.abc import Iterator + + +class Ade20kReader(DatasetReader): + """Enumerate ``/images/*.jpg`` paired with ``/annotations/*.png`` masks. + + Only images with a matching mask (same stem) are emitted. Mask paths + are not exposed here; downstream mIoU computation derives them from + the image stem. + """ + + def __init__(self, root) -> None: + super().__init__(root) + self.images_dir = self.root / "images" + self.annotations_dir = self.root / "annotations" + + def __iter__(self) -> Iterator[CalibrationSample]: + if not self.images_dir.exists(): + msg = f"ADE20K images directory missing: {self.images_dir}" + raise FileNotFoundError(msg) + if not self.annotations_dir.exists(): + msg = f"ADE20K annotations directory missing: {self.annotations_dir}" + raise FileNotFoundError(msg) + for img_path in sorted(self.images_dir.glob("*.jpg")): + mask_path = self.annotations_dir / f"{img_path.stem}.png" + if mask_path.exists(): + yield CalibrationSample( + image_path=img_path, + label=0, + mask_path=mask_path, + ) diff --git a/model_converter/src/model_converter/datasets/base.py b/model_converter/src/model_converter/datasets/base.py new file mode 100644 index 000000000..6cc5db5f6 --- /dev/null +++ b/model_converter/src/model_converter/datasets/base.py @@ -0,0 +1,44 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Common types for dataset readers.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Iterator + + +@dataclass(frozen=True) +class CalibrationSample: + """One image plus optional task-specific ground-truth pointers. + + ``label`` is the class id for class-folder datasets and a placeholder + ``0`` for layouts (COCO, ADE20K) where per-image classification labels + do not apply. ``image_id`` is the COCO image identifier used to match + predictions against the annotation JSON; ``mask_path`` is the ADE20K + ground-truth mask file. Both are ``None`` when not applicable. + """ + + image_path: Path + label: int + image_id: int | None = None + mask_path: Path | None = None + + +class DatasetReader(ABC): + """Abstract enumerator for calibration/validation samples on disk.""" + + def __init__(self, root: Path) -> None: + self.root = Path(root) + + @abstractmethod + def __iter__(self) -> Iterator[CalibrationSample]: + """Yield samples in deterministic order.""" diff --git a/model_converter/src/model_converter/datasets/class_folder.py b/model_converter/src/model_converter/datasets/class_folder.py new file mode 100644 index 000000000..29227227f --- /dev/null +++ b/model_converter/src/model_converter/datasets/class_folder.py @@ -0,0 +1,39 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Class-folder dataset layout (ImageNet-style).""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from .base import CalibrationSample, DatasetReader + +if TYPE_CHECKING: + from collections.abc import Iterator + from pathlib import Path + +_IMAGE_PATTERNS = ("*.JPEG", "*.jpg", "*.jpeg", "*.png") + + +class ClassFolderReader(DatasetReader): + """Enumerate ``//*.JPEG|jpg|png`` samples. + + Class-folder names must be integers. This matches the layout used by + ImageNet-1k and ImageNet-21k validation sets. + """ + + def __iter__(self) -> Iterator[CalibrationSample]: + if not self.root.exists(): + return + for class_dir in sorted(self.root.iterdir()): + if not class_dir.is_dir(): + continue + class_label = int(class_dir.name) # raises ValueError on bad layout + img_paths: set[Path] = set() + for pattern in _IMAGE_PATTERNS: + img_paths.update(class_dir.glob(pattern)) + for img_path in sorted(img_paths): + yield CalibrationSample(image_path=img_path, label=class_label) diff --git a/model_converter/src/model_converter/datasets/coco.py b/model_converter/src/model_converter/datasets/coco.py new file mode 100644 index 000000000..9f6ffdce0 --- /dev/null +++ b/model_converter/src/model_converter/datasets/coco.py @@ -0,0 +1,54 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""COCO-format dataset reader (images + annotations JSON).""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING + +from .base import CalibrationSample, DatasetReader + +if TYPE_CHECKING: + from collections.abc import Iterator + + +class CocoImagesReader(DatasetReader): + """Enumerate ``/images/*.jpg`` produced by ``download_coco.py``. + + The reader validates that the configured COCO-format annotation file is + present so downstream metric code can rely on it without re-checking. + Per-image ``label`` is a placeholder ``0`` because COCO ground truth is + consumed via the annotation JSON rather than per-image labels. The + ``image_id`` field is populated from the annotation JSON so prediction + records can reference the correct COCO image. + """ + + def __init__(self, root, annotation_filename: str) -> None: + super().__init__(root) + self.annotation_filename = annotation_filename + self.images_dir = self.root / "images" + self.annotations_path = self.root / "annotations" / annotation_filename + + def __iter__(self) -> Iterator[CalibrationSample]: + if not self.images_dir.exists(): + msg = f"COCO images directory missing: {self.images_dir}" + raise FileNotFoundError(msg) + if not self.annotations_path.exists(): + msg = f"COCO annotations file missing: {self.annotations_path}" + raise FileNotFoundError(msg) + filename_to_id = _load_filename_to_image_id(self.annotations_path) + for img_path in sorted(self.images_dir.glob("*.jpg")): + yield CalibrationSample( + image_path=img_path, + label=0, + image_id=filename_to_id.get(img_path.name), + ) + + +def _load_filename_to_image_id(annotation_path) -> dict[str, int]: + payload = json.loads(annotation_path.read_text()) + return {entry["file_name"]: int(entry["id"]) for entry in payload.get("images", [])} diff --git a/model_converter/src/model_converter/datasets/factory.py b/model_converter/src/model_converter/datasets/factory.py new file mode 100644 index 000000000..441f7e86d --- /dev/null +++ b/model_converter/src/model_converter/datasets/factory.py @@ -0,0 +1,40 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Dispatch ``dataset_type`` strings to concrete reader implementations.""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING + +from .ade20k import Ade20kReader +from .class_folder import ClassFolderReader +from .coco import CocoImagesReader + +if TYPE_CHECKING: + from .base import DatasetReader + +_CLASS_FOLDER_TYPES = frozenset({"imagenet-1k", "imagenet-21k"}) +_COCO_ANNOTATION_FILES = { + "coco-detection": "instances_val2017.json", +} + + +def reader_for(dataset_type: str | None, root: Path) -> DatasetReader: + """Return the :class:`DatasetReader` matching ``dataset_type``. + + ``None`` falls back to :class:`ClassFolderReader` so that callers without + a configured dataset type (e.g. legacy code paths) keep working. + """ + root = Path(root) + if dataset_type is None or dataset_type in _CLASS_FOLDER_TYPES: + return ClassFolderReader(root) + if dataset_type in _COCO_ANNOTATION_FILES: + return CocoImagesReader(root, annotation_filename=_COCO_ANNOTATION_FILES[dataset_type]) + if dataset_type == "ade20k": + return Ade20kReader(root) + msg = f"Unknown dataset_type: {dataset_type!r}" + raise ValueError(msg) diff --git a/model_converter/src/model_converter/metrics/__init__.py b/model_converter/src/model_converter/metrics/__init__.py new file mode 100644 index 000000000..427c2a832 --- /dev/null +++ b/model_converter/src/model_converter/metrics/__init__.py @@ -0,0 +1,28 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Task-aware accuracy metrics for the model converter pipeline. + +Each metric implements the :class:`~model_converter.metrics.base.Metric` +interface (``update``, ``compute``, ``reset``, plus a ``name`` label such as +``"top1"``, ``"mAP"``, ``"mIoU"``). The :func:`metric_for` factory dispatches +to the right metric for a given ``dataset_type`` and ``model_type``. +""" + +from model_converter.metrics.base import Metric +from model_converter.metrics.classification import TopOneAccuracy +from model_converter.metrics.coco_detection import CocoDetectionMAP +from model_converter.metrics.factory import metric_for +from model_converter.metrics.multilabel import MultilabelMAP +from model_converter.metrics.semseg_miou import SemSegMIoU + +__all__ = [ + "CocoDetectionMAP", + "Metric", + "MultilabelMAP", + "SemSegMIoU", + "TopOneAccuracy", + "metric_for", +] diff --git a/model_converter/src/model_converter/metrics/base.py b/model_converter/src/model_converter/metrics/base.py new file mode 100644 index 000000000..4ac689a69 --- /dev/null +++ b/model_converter/src/model_converter/metrics/base.py @@ -0,0 +1,35 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Abstract metric interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + + +class Metric(ABC): + """Accumulator-style metric. + + Subclasses are responsible for advertising a short label via :attr:`name` + (e.g. ``"top1"``, ``"mAP"``, ``"mIoU"``), accepting one prediction at a + time via :meth:`update`, returning a scalar via :meth:`compute`, and + clearing internal state via :meth:`reset`. + """ + + name: str + + @abstractmethod + def update(self, prediction: Any, ground_truth: Any = None) -> None: + """Accumulate a single prediction with its ground truth.""" + + @abstractmethod + def compute(self) -> float: + """Return the metric value for everything accumulated so far.""" + + @abstractmethod + def reset(self) -> None: + """Clear all accumulated state.""" diff --git a/model_converter/src/model_converter/metrics/classification.py b/model_converter/src/model_converter/metrics/classification.py new file mode 100644 index 000000000..990a0d56c --- /dev/null +++ b/model_converter/src/model_converter/metrics/classification.py @@ -0,0 +1,38 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Top-1 classification accuracy metric.""" + +from __future__ import annotations + +import numpy as np + +from model_converter.metrics.base import Metric + + +class TopOneAccuracy(Metric): + """Fraction of samples whose argmax matches the ground-truth label.""" + + name = "top1" + + def __init__(self) -> None: + self._correct = 0 + self._total = 0 + + def update(self, prediction: np.ndarray, ground_truth: int | None = None) -> None: + """Accept ``prediction`` as a logits tensor and ``ground_truth`` as a class id.""" + assert ground_truth is not None + pred_class = int(np.argmax(prediction, axis=1)[0]) + self._correct += int(pred_class == int(ground_truth)) + self._total += 1 + + def compute(self) -> float: + if self._total == 0: + return 0.0 + return self._correct / self._total + + def reset(self) -> None: + self._correct = 0 + self._total = 0 diff --git a/model_converter/src/model_converter/metrics/coco_detection.py b/model_converter/src/model_converter/metrics/coco_detection.py new file mode 100644 index 000000000..e2226fe6e --- /dev/null +++ b/model_converter/src/model_converter/metrics/coco_detection.py @@ -0,0 +1,150 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""COCO-format mean Average Precision (bbox / segm) via pycocotools.""" + +from __future__ import annotations + +import contextlib +import io +from pathlib import Path +from typing import Any + +from model_converter.metrics.base import Metric + +_VALID_IOU_TYPES = {"bbox", "segm"} + +# Canonical mapping from COCO 80-class indices (0-79) used by OpenVINO detection +# models to the original COCO 91-class category IDs (non-contiguous) in COCO JSON. +# COCO dropped 11 categories (IDs 12, 26, 29, 30, 45, 66, 68, 69, 71, 83, 91), +# so class index 11 (stop sign) maps to category ID 13, not 12, etc. +COCO80_TO_COCO91 = [ + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 20, + 21, + 22, + 23, + 24, + 25, + 27, + 28, + 31, + 32, + 33, + 34, + 35, + 36, + 37, + 38, + 39, + 40, + 41, + 42, + 43, + 44, + 46, + 47, + 48, + 49, + 50, + 51, + 52, + 53, + 54, + 55, + 56, + 57, + 58, + 59, + 60, + 61, + 62, + 63, + 64, + 65, + 67, + 70, + 72, + 73, + 74, + 75, + 76, + 77, + 78, + 79, + 80, + 81, + 82, + 84, + 85, + 86, + 87, + 88, + 89, + 90, +] + + +class CocoDetectionMAP(Metric): + """Wraps :class:`pycocotools.cocoeval.COCOeval` for a single ``iouType``.""" + + name = "mAP" + + def __init__(self, annotation_file: Path, iou_type: str = "bbox") -> None: + if iou_type not in _VALID_IOU_TYPES: + error_msg = f"Unsupported iou_type {iou_type!r}; expected one of {sorted(_VALID_IOU_TYPES)}" + raise ValueError(error_msg) + self.annotation_file = Path(annotation_file) + self.iou_type = iou_type + self._predictions: list[dict[str, Any]] = [] + + def update(self, predictions: list[dict[str, Any]] | None = None, ground_truth: Any = None) -> None: + """Accept a batch of COCO-format detection dicts. + + Each dict must contain ``image_id``, ``category_id``, ``bbox`` + (``[x, y, w, h]``), and ``score``. ``ground_truth`` is ignored — GT is + loaded from :attr:`annotation_file`. + """ + del ground_truth + if predictions is None: + return + self._predictions.extend(predictions) + + def compute(self) -> float: + from pycocotools.coco import COCO + from pycocotools.cocoeval import COCOeval + + if not self._predictions: + return 0.0 + + with contextlib.redirect_stdout(io.StringIO()): + coco_gt = COCO(str(self.annotation_file)) + coco_dt = coco_gt.loadRes(self._predictions) + evaluator = COCOeval(coco_gt, coco_dt, iouType=self.iou_type) + evaluator.evaluate() + evaluator.accumulate() + evaluator.summarize() + # stats[0] is the primary mAP @ IoU=0.50:0.95. + return float(evaluator.stats[0]) + + def reset(self) -> None: + self._predictions = [] diff --git a/model_converter/src/model_converter/metrics/factory.py b/model_converter/src/model_converter/metrics/factory.py new file mode 100644 index 000000000..af877c80e --- /dev/null +++ b/model_converter/src/model_converter/metrics/factory.py @@ -0,0 +1,90 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Dispatcher mapping ``(dataset_type, model_type)`` → :class:`Metric` instance.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from model_converter.metrics.classification import TopOneAccuracy +from model_converter.metrics.coco_detection import CocoDetectionMAP +from model_converter.metrics.multilabel import MultilabelMAP +from model_converter.metrics.semseg_miou import SemSegMIoU + +if TYPE_CHECKING: + from pathlib import Path + + from model_converter.metrics.base import Metric + +_IMAGENET_CLASSIFICATION_TYPES = {"Classification"} +_IMAGENET_MULTILABEL_TYPES = {"Classification_Multilabel"} +_COCO_DETECTION_TYPES = { + "SSD", + "YOLO", + "YOLOv4", + "YOLOv5", + "YOLOv8", + "YOLO11", + "YOLOX", + "MaskRCNN", + "RotatedDetection", + "DETR", + "Detection", +} +_SEMSEG_TYPES = {"Segmentation"} + +_MULTILABEL_TASKS = {"MULTI_LABEL_CLS"} +_MULTI_CLASS_TASKS = {"MULTI_CLASS_CLS"} +_UNSUPPORTED_TASKS = {"ROTATED_DETECTION"} + +_IMAGENET_LABEL_COUNTS = { + "imagenet-1k": 1000, + "imagenet-21k": 21843, +} + + +def metric_for( + dataset_type: str | None, + model_type: str | None, + *, + annotation_file: "Path | None" = None, + task: str | None = None, +) -> "Metric | None": + """Return a fresh metric instance for the given dataset/model combo, or ``None``. + + ``None`` means accuracy is not measured for this configuration (unknown + dataset type, headless feature extractor, rotated detection, etc.). + + The optional ``task`` argument is the OpenVINO Training Extensions task + label (``getitune_task``) and is used to disambiguate cases where + ``model_type`` alone does not carry enough information — for example, + ``Classification`` model_type with task ``MULTI_LABEL_CLS`` maps to + multilabel mAP, not top-1 accuracy. + """ + if dataset_type is None or model_type is None: + return None + + if task in _UNSUPPORTED_TASKS: + return None + + if dataset_type in _IMAGENET_LABEL_COUNTS: + if task in _MULTILABEL_TASKS or model_type in _IMAGENET_MULTILABEL_TYPES: + return MultilabelMAP(num_labels=_IMAGENET_LABEL_COUNTS[dataset_type]) + if task in _MULTI_CLASS_TASKS or model_type in _IMAGENET_CLASSIFICATION_TYPES: + return TopOneAccuracy() + return None + + if dataset_type == "coco-detection": + if annotation_file is None or model_type not in _COCO_DETECTION_TYPES: + return None + return CocoDetectionMAP(annotation_file=annotation_file, iou_type="bbox") + + if dataset_type == "ade20k": + if model_type not in _SEMSEG_TYPES: + return None + return SemSegMIoU(num_classes=150, ignore_index=255) + + return None diff --git a/model_converter/src/model_converter/metrics/multilabel.py b/model_converter/src/model_converter/metrics/multilabel.py new file mode 100644 index 000000000..d552aabbc --- /dev/null +++ b/model_converter/src/model_converter/metrics/multilabel.py @@ -0,0 +1,40 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Multilabel mean Average Precision (macro) via torchmetrics.""" + +from __future__ import annotations + +import numpy as np +import torch +from torchmetrics.classification import MultilabelAveragePrecision + +from model_converter.metrics.base import Metric + + +class MultilabelMAP(Metric): + """Threshold-free macro mAP for multi-label classification.""" + + name = "mAP" + + def __init__(self, num_labels: int) -> None: + self.num_labels = num_labels + self._impl = MultilabelAveragePrecision(num_labels=num_labels, average="macro") + + def update(self, prediction: np.ndarray, ground_truth: np.ndarray | None = None) -> None: + """Accept ``prediction`` as logits of shape ``(1, num_labels)`` and ``ground_truth`` as a 0/1 vector.""" + logits = torch.from_numpy(np.asarray(prediction, dtype=np.float32)) + target = torch.from_numpy(np.asarray(ground_truth, dtype=np.int64)) + if logits.ndim == 1: + logits = logits.unsqueeze(0) + if target.ndim == 1: + target = target.unsqueeze(0) + self._impl.update(logits, target) + + def compute(self) -> float: + return float(self._impl.compute().item()) + + def reset(self) -> None: + self._impl.reset() diff --git a/model_converter/src/model_converter/metrics/semseg_miou.py b/model_converter/src/model_converter/metrics/semseg_miou.py new file mode 100644 index 000000000..f73ad888b --- /dev/null +++ b/model_converter/src/model_converter/metrics/semseg_miou.py @@ -0,0 +1,55 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Semantic segmentation mIoU computed from a confusion matrix.""" + +from __future__ import annotations + +import numpy as np + +from model_converter.metrics.base import Metric + + +class SemSegMIoU(Metric): + """Mean Intersection-over-Union with optional ignore index. + + Accumulates an ``(num_classes, num_classes)`` confusion matrix; per-class + IoU is averaged across classes that have at least one pixel in either GT + or prediction (others are skipped to avoid biasing toward easy zeros). + """ + + name = "mIoU" + + def __init__(self, num_classes: int, ignore_index: int = 255) -> None: + self.num_classes = num_classes + self.ignore_index = ignore_index + self._confusion = np.zeros((num_classes, num_classes), dtype=np.int64) + + def update(self, prediction: np.ndarray, ground_truth: np.ndarray | None = None) -> None: + """Accept ``prediction`` and ``ground_truth`` as same-shape class-id arrays.""" + gt = np.asarray(ground_truth).ravel() + pred = np.asarray(prediction).ravel() + mask = gt != self.ignore_index + gt = gt[mask] + pred = pred[mask] + index = gt * self.num_classes + pred + binc = np.bincount(index, minlength=self.num_classes * self.num_classes) + self._confusion += binc.reshape(self.num_classes, self.num_classes) + + def compute(self) -> float: + cm = self._confusion + tp = np.diag(cm).astype(np.float64) + fp = cm.sum(axis=0).astype(np.float64) - tp + fn = cm.sum(axis=1).astype(np.float64) - tp + denom = tp + fp + fn + valid = denom > 0 + if not np.any(valid): + return 0.0 + iou = np.zeros_like(denom) + iou[valid] = tp[valid] / denom[valid] + return float(iou[valid].mean()) + + def reset(self) -> None: + self._confusion = np.zeros((self.num_classes, self.num_classes), dtype=np.int64) diff --git a/model_converter/src/model_converter/reporting.py b/model_converter/src/model_converter/reporting.py new file mode 100644 index 000000000..466207e51 --- /dev/null +++ b/model_converter/src/model_converter/reporting.py @@ -0,0 +1,287 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Conversion summary report generation. + +Collects per-model conversion outcomes (accuracies and status) in a structured +form and renders them as a console table and a Markdown report. +""" + +from __future__ import annotations + +import dataclasses +import json +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +# Status string constants +STATUS_OK = "OK" +STATUS_OK_NO_ACCURACY = "OK (no accuracy data)" +STATUS_ACCURACY_DROP = "ACCURACY DROP >5%" +STATUS_FAILED_CONVERSION = "FAILED: conversion" +STATUS_FAILED_QUANTIZATION = "FAILED: quantization" +STATUS_SKIPPED = "SKIPPED" + +#: Maximum tolerated top-1 accuracy drop (percentage points) versus FP32. +DEFAULT_ACCURACY_DROP_THRESHOLD = 5.0 + +_REPORT_COLUMNS = ( + "Model Full Name", + "Model Type", + "Model Library", + "Original URL", + "Original Accuracy", + "FP32 Accuracy", + "FP16 Accuracy", + "INT8 Accuracy", + "Status", +) + + +@dataclass +class ConversionResult: + """Structured outcome for a single model conversion. + + Accuracies are top-1 accuracy fractions in the ``[0.0, 1.0]`` range, or + ``None`` when not measured. + """ + + model_short_name: str + model_full_name: str + model_type: str + model_library: str + original_url: str | None = None + original_accuracy: float | None = None + fp32_accuracy: float | None = None + fp16_accuracy: float | None = None + int8_accuracy: float | None = None + status: str = STATUS_OK + status_detail: str = "" + + +@dataclass +class AccuracyResults: + """Container for accuracies measured during quantization.""" + + original_accuracy: float | None = None + fp32_accuracy: float | None = None + fp16_accuracy: float | None = None + int8_accuracy: float | None = None + int8_succeeded: bool = False + measured: bool = field(default=False) + metric_name: str | None = None + + +def original_url_for_config(config: dict[str, Any]) -> str | None: + """Resolve the exact download URL for a model configuration. + + Only exact download locations are returned: ``weights_url`` for PyTorch + checkpoints and the Hugging Face repository URL for ``huggingface_repo``. + Returns ``None`` when no exact download URL is available. + + Args: + config: Model configuration dictionary. + + Returns: + The download URL, or ``None`` if unavailable. + """ + weights_url = config.get("weights_url") + if weights_url: + return str(weights_url) + + hf_repo = config.get("huggingface_repo") + if hf_repo: + return f"https://huggingface.co/{hf_repo}" + + return None + + +def determine_status( + result: ConversionResult, + *, + converted: bool, + quantized: bool, + skipped: bool, + threshold: float = DEFAULT_ACCURACY_DROP_THRESHOLD, +) -> tuple[str, str]: + """Compute the status and detail for a conversion result. + + FP32 is the preferred baseline for quantization drops. When FP32 accuracy + is absent (e.g. for YOLO11 models where no FP32 OV artifact is produced), + ``original_accuracy`` is used as baseline instead. A drop is abnormal when + ``(baseline - FP16)`` or ``(baseline - INT8)`` exceeds ``threshold`` + percentage points. + + Args: + result: The conversion result holding accuracy values. + converted: Whether the FP16 model was produced. + quantized: Whether the INT8 model was produced. + skipped: Whether processing was skipped (model already existed). + threshold: Maximum tolerated accuracy drop in percentage points. + + Returns: + A tuple of ``(status, status_detail)``. + """ + if skipped: + return STATUS_SKIPPED, "FP16 and INT8 models already existed" + + if not converted: + return STATUS_FAILED_CONVERSION, "Model load or export failed" + + if not quantized: + return STATUS_FAILED_QUANTIZATION, "INT8 model was not produced" + + # Determine the baseline accuracy for drop checks. + # Prefer FP32 (standard OV export); fall back to original (e.g. YOLO PT model). + baseline = result.fp32_accuracy if result.fp32_accuracy is not None else result.original_accuracy + if baseline is None: + return STATUS_OK_NO_ACCURACY, "No accuracy measurement available" + + baseline_pct = baseline * 100 + drops: list[str] = [] + + if result.fp32_accuracy is not None and result.original_accuracy is not None: + fp32_drop = result.original_accuracy * 100 - baseline_pct + if fp32_drop > threshold: + drops.append(f"FP32 drop {fp32_drop:.2f}%") + + if result.fp16_accuracy is not None: + fp16_drop = baseline_pct - result.fp16_accuracy * 100 + if fp16_drop > threshold: + drops.append(f"FP16 drop {fp16_drop:.2f}%") + + if result.int8_accuracy is not None: + int8_drop = baseline_pct - result.int8_accuracy * 100 + if int8_drop > threshold: + drops.append(f"INT8 drop {int8_drop:.2f}%") + + if drops: + return STATUS_ACCURACY_DROP, "; ".join(drops) + + return STATUS_OK, "" + + +def _format_accuracy(value: float | None) -> str: + """Format an accuracy fraction as a percentage string or ``N/A``.""" + if value is None: + return "N/A" + return f"{value * 100:.2f}%" + + +def _row_values(result: ConversionResult) -> list[str]: + """Build the ordered cell values for a single result row.""" + return [ + result.model_full_name or result.model_short_name, + result.model_type or "N/A", + result.model_library or "N/A", + result.original_url or "N/A", + _format_accuracy(result.original_accuracy), + _format_accuracy(result.fp32_accuracy), + _format_accuracy(result.fp16_accuracy), + _format_accuracy(result.int8_accuracy), + result.status, + ] + + +def format_console_table(results: list[ConversionResult]) -> str: + """Render results as an aligned plain-text table for the console. + + Args: + results: Conversion results to render. + + Returns: + The formatted table as a string. + """ + rows = [list(_REPORT_COLUMNS)] + [_row_values(r) for r in results] + widths = [max(len(row[col]) for row in rows) for col in range(len(_REPORT_COLUMNS))] + + def render_row(row: list[str]) -> str: + return " ".join(cell.ljust(widths[i]) for i, cell in enumerate(row)) + + separator = " ".join("-" * widths[i] for i in range(len(_REPORT_COLUMNS))) + + lines = ["Conversion Summary Report", "", render_row(rows[0]), separator] + lines.extend(render_row(row) for row in rows[1:]) + return "\n".join(lines) + + +def format_markdown_report(results: list[ConversionResult]) -> str: + """Render results as a Markdown document with a summary table. + + Args: + results: Conversion results to render. + + Returns: + The Markdown report as a string. + """ + header = "| " + " | ".join(_REPORT_COLUMNS) + " |" + separator = "| " + " | ".join("---" for _ in _REPORT_COLUMNS) + " |" + + lines = ["# Conversion Summary Report", ""] + if not results: + lines.append("No models were processed.") + return "\n".join(lines) + "\n" + + lines.append(header) + lines.append(separator) + for result in results: + cells = [cell.replace("|", "\\|") for cell in _row_values(result)] + lines.append("| " + " | ".join(cells) + " |") + return "\n".join(lines) + "\n" + + +def write_markdown_report(results: list[ConversionResult], path: Path) -> None: + """Write the Markdown report to ``path``, creating parent directories. + + Args: + results: Conversion results to render. + path: Destination file path. + """ + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(format_markdown_report(results)) + + +def _load_results_from_json(path: Path) -> list[ConversionResult]: + """Load persisted ConversionResult objects from a JSON sidecar file. + + Args: + path: Path to the JSON sidecar file. + + Returns: + List of ConversionResult objects, or an empty list if the file does not exist. + """ + if not path.exists(): + return [] + try: + data = json.loads(path.read_text()) + return [ConversionResult(**entry) for entry in data] + except (json.JSONDecodeError, TypeError, KeyError): + return [] + + +def upsert_result(result: ConversionResult, path: Path) -> None: + """Upsert a single result into the report files. + + Reads the existing JSON sidecar (``path.with_suffix('.json')``), replaces + the entry whose ``model_short_name`` matches ``result``, or appends it if + absent. Then writes the updated JSON sidecar and regenerates the Markdown + report at ``path``. + + Args: + result: Conversion result to upsert. + path: Destination Markdown file path. + """ + path = Path(path) + json_path = path.with_suffix(".json") + + existing = _load_results_from_json(json_path) + updated = [r for r in existing if r.model_short_name != result.model_short_name] + updated.append(result) + + path.parent.mkdir(parents=True, exist_ok=True) + json_path.write_text(json.dumps([dataclasses.asdict(r) for r in updated], indent=2)) + path.write_text(format_markdown_report(updated)) diff --git a/model_converter/src/model_converter/templates/README-getitune-fp16.md b/model_converter/src/model_converter/templates/README-getitune-fp16.md new file mode 100644 index 000000000..11d756b3d --- /dev/null +++ b/model_converter/src/model_converter/templates/README-getitune-fp16.md @@ -0,0 +1,68 @@ +--- +license: <> +tags: +<> +--- + +# <>-fp16-ov + +- Model creator: [Geti™](https://github.com/open-edge-platform/training_extensions) +- Original model: [<>](<>) + +## Description + +This is a [Geti™](https://github.com/open-edge-platform/training_extensions) version of [<>](<>) model converted to the [OpenVINO™ IR](https://docs.openvino.ai/2026/documentation/openvino-ir-format.html) (Intermediate Representation) format with weights compressed to FP16. + +To fine-tune your model with a custom dataset, you can use Geti™ to annotate data, perform fine-tuning, and export the resulting model. + +## Compatibility + +The provided OpenVINO™ IR model is compatible with: + +- OpenVINO version 2026.1.0 and higher +- Model API 0.4.0 and higher + +## Running Model Inference with [Model API](https://github.com/open-edge-platform/model_api) + +1. Install required packages: + +```sh +pip install openvino-model-api[huggingface] +``` + + + +2. Run model inference: + +```python +import cv2 +from model_api.models import Model +from model_api.visualizer import Visualizer + +# 1. Load model +model = Model.from_pretrained("OpenVINO/<>-fp16-ov") + +# 2. Load image +image = cv2.imread("image.jpg") + +# 3. Run inference +result = model(image) + +# 4. Visualize and save results +vis = Visualizer().render(image, result) +cv2.imwrite("output.jpg", vis) +``` + +For more examples and possible optimizations, refer to the [Model API Documentation](https://open-edge-platform.github.io/model_api/latest/). + +## Limitations + +Check the original [model documentation](<>) for limitations. + +## Legal information + +The original model is distributed under the [<>](<>) license. More details can be found in [training_extensions](https://github.com/open-edge-platform/training_extensions). + +## Disclaimer + +Intel is committed to respecting human rights and avoiding causing or contributing to adverse impacts on human rights. See [Intel's Global Human Rights Principles](https://www.intel.com/content/dam/www/central-libraries/us/en/documents/policy-human-rights.pdf). Intel's products and software are intended only to be used in applications that do not cause or contribute to adverse impacts on human rights. diff --git a/model_converter/src/model_converter/templates/README-getitune-int8.md b/model_converter/src/model_converter/templates/README-getitune-int8.md new file mode 100644 index 000000000..f34918c0e --- /dev/null +++ b/model_converter/src/model_converter/templates/README-getitune-int8.md @@ -0,0 +1,77 @@ +--- +license: <> +tags: +<> +--- + +# <>-int8-ov + +- Model creator: [Geti™](https://github.com/open-edge-platform/training_extensions) +- Original model: [<>](<>) + +## Description + +This is a [Geti™](https://github.com/open-edge-platform/training_extensions) version of [<>](<>) model converted to the [OpenVINO™ IR](https://docs.openvino.ai/2026/documentation/openvino-ir-format.html) (Intermediate Representation) format with weights compressed to INT8. + +To fine-tune your model with a custom dataset, you can use Geti™ to annotate data, perform fine-tuning, and export the resulting model. + +## Quantization Parameters + +Weight compression was performed using nncf.quantize with the following parameters: + +- **Quantization method**: Post-Training Quantization (PTQ) +- **Precision**: INT8 for both weights and activations + +For more information on quantization, check the [OpenVINO model optimization guide](https://docs.openvino.ai/2026/openvino-workflow/model-optimization-guide/quantizing-models-post-training.html). + +## Compatibility + +The provided OpenVINO™ IR model is compatible with: + +- OpenVINO version 2026.1.0 and higher +- Model API 0.4.0 and higher + +## Running Model Inference with [Model API](https://github.com/open-edge-platform/model_api) + +1. Install required packages: + +```sh +pip install openvino-model-api[huggingface] +``` + + + +2. Run model inference: + +```python +import cv2 +from model_api.models import Model +from model_api.visualizer import Visualizer + +# 1. Load model +model = Model.from_pretrained("OpenVINO/<>-int8-ov") + +# 2. Load image +image = cv2.imread("image.jpg") + +# 3. Run inference +result = model(image) + +# 4. Visualize and save results +vis = Visualizer().render(image, result) +cv2.imwrite("output.jpg", vis) +``` + +For more examples and possible optimizations, refer to the [Model API Documentation](https://open-edge-platform.github.io/model_api/latest/). + +## Limitations + +Check the original [model documentation](<>) for limitations. + +## Legal information + +The original model is distributed under the [<>](<>) license. More details can be found in [training_extensions](https://github.com/open-edge-platform/training_extensions). + +## Disclaimer + +Intel is committed to respecting human rights and avoiding causing or contributing to adverse impacts on human rights. See [Intel's Global Human Rights Principles](https://www.intel.com/content/dam/www/central-libraries/us/en/documents/policy-human-rights.pdf). Intel's products and software are intended only to be used in applications that do not cause or contribute to adverse impacts on human rights. diff --git a/model_converter/src/model_converter/yolo/yolo.py b/model_converter/src/model_converter/yolo/yolo.py deleted file mode 100644 index 291a622cc..000000000 --- a/model_converter/src/model_converter/yolo/yolo.py +++ /dev/null @@ -1,117 +0,0 @@ -# -# Copyright (C) 2026 Intel Corporation -# SPDX-License-Identifier: Apache-2.0 -# - -import json -import shutil -from pathlib import Path - -from defusedxml.ElementTree import ParseError, parse -from ultralytics import YOLO - -TEMPLATES_DIR = Path(__file__).resolve().parent.parent / "templates" - - -def copy_readme_template(template_name: str, dest_dir: Path, yolo_size: str) -> None: - """Copy a README template to dest_dir/README.md, replacing <>.""" - template_path = TEMPLATES_DIR / template_name - content = template_path.read_text() - content = content.replace("<>", yolo_size) - (dest_dir / "README.md").write_text(content) - print(f"Copied {template_name} -> {dest_dir / 'README.md'} (size={yolo_size})") - - -def update_model_type_in_xml(xml_path: Path, model_type: str = "YOLO11") -> None: - """Update the model_type value in the OpenVINO XML file.""" - try: - tree = parse(xml_path) - root = tree.getroot() - - # Find and update model_type in rt_info - for rt_info in root.findall(".//rt_info"): - for model_info in rt_info.findall(".//model_info"): - for model_type_elem in model_info.findall(".//model_type"): - model_type_elem.set("value", model_type) - - tree.write(xml_path, encoding="utf-8", xml_declaration=True) - print(f"Updated model_type to {model_type} in {xml_path}") - - model_info_dict = {} - for rt_info in root.findall(".//rt_info"): - for model_info in rt_info.findall(".//model_info"): - for child in model_info: - model_info_dict[child.tag] = child.attrib["value"] - with (xml_path.parent / "config.json").open("w") as f: - json.dump(model_info_dict, f, indent=4) - - except (OSError, ParseError) as error: - print(f"Failed to update {xml_path}: {error}") - - -# YOLO11 model versions -MODEL_VERSIONS = ["yolo11n", "yolo11s", "yolo11m", "yolo11l", "yolo11x"] - - -def convert_yolo_models(model_versions: list[str] | None = None) -> None: - """Convert YOLO11 model variants to OpenVINO IR folders.""" - for version in model_versions or MODEL_VERSIONS: - print(f"Processing {version}...") - yolo_size = version[-1] # n, s, m, l, or x - - # Load model - model = YOLO(f"{version}.pt") - - # Export regular OpenVINO model - print(f"Exporting {version} to OpenVINO format...") - model.export(format="openvino", half=True) - - # Rename output folder for regular model - old_name = Path(f"{version}_openvino_model") - new_name = Path(f"YOLO{version[4:]}-fp16-ov") - if old_name.exists(): - if new_name.exists(): - shutil.rmtree(new_name) - old_name.rename(new_name) - print(f"Renamed {old_name} to {new_name}") - - # Update model_type in XML metadata - xml_file = new_name / f"{version}.xml" - if xml_file.exists(): - update_model_type_in_xml(xml_file, "YOLO11") - - # Copy README template for fp16 - copy_readme_template("README-yolo-fp16.md", new_name, yolo_size) - - # Export INT8 quantized OpenVINO model - print(f"Exporting {version} to OpenVINO INT8 format...") - model.export(format="openvino", int8=True, data="coco128.yaml") - - # Rename output folder for INT8 model - old_name_int8 = Path(f"{version}_int8_openvino_model") - new_name_int8 = Path(f"YOLO{version[4:]}-int8-ov") - if old_name_int8.exists(): - if new_name_int8.exists(): - shutil.rmtree(new_name_int8) - old_name_int8.rename(new_name_int8) - print(f"Renamed {old_name_int8} to {new_name_int8}") - - # Update model_type in XML metadata - xml_file = new_name_int8 / f"{version}.xml" - if xml_file.exists(): - update_model_type_in_xml(xml_file, "YOLO11") - - # Copy README template for int8 - copy_readme_template("README-yolo-int8.md", new_name_int8, yolo_size) - - print(f"Completed {version}\n") - - -def main() -> int: - """Run YOLO11 conversion for all configured variants.""" - convert_yolo_models() - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/model_converter/tests/unit/conftest.py b/model_converter/tests/unit/conftest.py index 47f75f5b1..73e98f2bd 100644 --- a/model_converter/tests/unit/conftest.py +++ b/model_converter/tests/unit/conftest.py @@ -24,10 +24,11 @@ def tmp_cache_dir(tmp_path): @pytest.fixture def sample_model_config(): - """Sample model configuration dictionary.""" + """Sample torchvision model configuration dictionary.""" return { "model_short_name": "test_model", "model_full_name": "Test Model", + "model_library": "torchvision", "model_class_name": "torchvision.models.resnet.resnet18", "weights_url": "https://example.com/weights.pth", "input_shape": [1, 3, 224, 224], @@ -42,12 +43,38 @@ def sample_model_config(): "scale_values": "58.395 57.12 57.375", "reverse_input_channels": True, "description": "A test model", + "dataset_type": "imagenet-1k", } @pytest.fixture -def converter(tmp_output_dir, tmp_cache_dir): - """Pre-built ModelConverter instance with temporary directories.""" +def sample_timm_config(): + """Sample timm model configuration dictionary.""" + return { + "model_short_name": "test_timm_model", + "model_full_name": "Test Timm Model", + "model_library": "timm", + "huggingface_repo": "timm/resnet50.a1_in1k", + "huggingface_revision": "abc123", + "input_shape": [1, 3, 224, 224], + "input_names": ["input"], + "output_names": ["result"], + "model_type": "Classification", + "license": "Apache-2.0", + "license_link": "https://www.apache.org/licenses/LICENSE-2.0", + "docs": "https://docs.example.com", + "labels": "IMAGENET1K_V1", + "mean_values": "123.675 116.28 103.53", + "scale_values": "58.395 57.12 57.375", + "reverse_input_channels": True, + "description": "A test timm model", + "dataset_type": "imagenet-1k", + } + + +@pytest.fixture +def facade_converter(tmp_output_dir, tmp_cache_dir): + """Pre-built facade ModelConverter instance with temporary directories.""" from model_converter.cli import ModelConverter return ModelConverter( @@ -57,19 +84,42 @@ def converter(tmp_output_dir, tmp_cache_dir): ) +@pytest.fixture +def converter(tmp_output_dir, tmp_cache_dir): + """Pre-built TorchvisionConverter instance with temporary directories.""" + from model_converter.converters.torchvision import TorchvisionConverter + + return TorchvisionConverter( + output_dir=tmp_output_dir, + cache_dir=tmp_cache_dir, + verbose=True, + ) + + +@pytest.fixture +def timm_converter(tmp_output_dir, tmp_cache_dir): + """Pre-built TimmConverter instance with temporary directories.""" + from model_converter.converters.timm import TimmConverter + + return TimmConverter( + output_dir=tmp_output_dir, + cache_dir=tmp_cache_dir, + verbose=True, + ) + + @pytest.fixture def mock_ov_model(): """Mock OpenVINO model object.""" model = MagicMock() model.inputs = [MagicMock()] model.outputs = [MagicMock()] - model.input.return_value = MagicMock() - model.output.return_value = MagicMock() - # Mock input(0).get_names() input_mock = MagicMock() input_mock.get_names.return_value = {"input"} model.input.return_value = input_mock + model.output.return_value = MagicMock() + model.get_rt_info.return_value = MagicMock(value={"model_type": "Classification"}) return model @@ -97,15 +147,13 @@ def mock_torch_model(): @pytest.fixture def dataset_dir(tmp_path): """Create a temporary calibration dataset directory structure.""" + import cv2 import numpy as np dataset_path = tmp_path / "dataset" class_dir = dataset_path / "0" class_dir.mkdir(parents=True) - # Create a dummy image file - import cv2 - img = np.zeros((224, 224, 3), dtype=np.uint8) cv2.imwrite(str(class_dir / "image_001.jpg"), img) @@ -126,3 +174,31 @@ def template_dir(tmp_path): (templates / "README-torchvision-fp16.md").write_text("# <> (<>)\nLicense: <>") (templates / ".gitattributes").write_text("*.bin filter=lfs diff=lfs merge=lfs -text\n") return templates + + +@pytest.fixture +def datasets_config(tmp_path, dataset_dir): + """Create a datasets configuration JSON file for testing.""" + import json + + config = {"datasets": {"imagenet-1k": str(dataset_dir), "coco-detection": str(tmp_path / "coco")}} + config_file = tmp_path / "datasets.json" + with config_file.open("w") as f: + json.dump(config, f) + return config_file + + +@pytest.fixture +def dataset_registry(datasets_config): + """Create a DatasetRegistry for testing.""" + from model_converter.dataset_registry import DatasetRegistry + + return DatasetRegistry(datasets_config) + + +@pytest.fixture +def mock_dataset_registry(dataset_dir): + """Create a mocked DatasetRegistry for testing.""" + mock = MagicMock() + mock.resolve_from_config.return_value = dataset_dir + return mock diff --git a/model_converter/src/model_converter/yolo/__init__.py b/model_converter/tests/unit/datasets/__init__.py similarity index 71% rename from model_converter/src/model_converter/yolo/__init__.py rename to model_converter/tests/unit/datasets/__init__.py index 67ce5b11b..d22239848 100644 --- a/model_converter/src/model_converter/yolo/__init__.py +++ b/model_converter/tests/unit/datasets/__init__.py @@ -2,5 +2,3 @@ # Copyright (C) 2026 Intel Corporation # SPDX-License-Identifier: Apache-2.0 # - -"""YOLO conversion helpers.""" diff --git a/model_converter/tests/unit/datasets/test_ade20k.py b/model_converter/tests/unit/datasets/test_ade20k.py new file mode 100644 index 000000000..155d072d9 --- /dev/null +++ b/model_converter/tests/unit/datasets/test_ade20k.py @@ -0,0 +1,75 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Unit tests for the ADE20K reader.""" + +from __future__ import annotations + +import pytest +from model_converter.datasets import Ade20kReader + + +@pytest.fixture +def ade20k_root(tmp_path): + images = tmp_path / "images" + annotations = tmp_path / "annotations" + images.mkdir() + annotations.mkdir() + for stem in ("ADE_val_00000001", "ADE_val_00000002", "ADE_val_00000003"): + (images / f"{stem}.jpg").write_bytes(b"") + (annotations / f"{stem}.png").write_bytes(b"") + return tmp_path + + +def test_enumerates_all_image_mask_pairs(ade20k_root): + reader = Ade20kReader(ade20k_root) + samples = list(reader) + assert {s.image_path.name for s in samples} == { + "ADE_val_00000001.jpg", + "ADE_val_00000002.jpg", + "ADE_val_00000003.jpg", + } + + +def test_labels_are_placeholder_zero(ade20k_root): + reader = Ade20kReader(ade20k_root) + assert all(s.label == 0 for s in reader) + + +def test_skips_images_without_matching_mask(ade20k_root): + # Add an unmatched image + (ade20k_root / "images" / "ADE_val_00000004.jpg").write_bytes(b"") + reader = Ade20kReader(ade20k_root) + names = [s.image_path.name for s in reader] + assert "ADE_val_00000004.jpg" not in names + assert len(names) == 3 + + +def test_raises_when_images_directory_missing(tmp_path): + (tmp_path / "annotations").mkdir() + reader = Ade20kReader(tmp_path) + with pytest.raises(FileNotFoundError, match="images"): + list(reader) + + +def test_raises_when_annotations_directory_missing(tmp_path): + (tmp_path / "images").mkdir() + reader = Ade20kReader(tmp_path) + with pytest.raises(FileNotFoundError, match="annotations"): + list(reader) + + +def test_results_sorted_by_image_name(ade20k_root): + reader = Ade20kReader(ade20k_root) + names = [s.image_path.name for s in reader] + assert names == sorted(names) + + +def test_mask_path_populated_for_each_sample(ade20k_root): + reader = Ade20kReader(ade20k_root) + for sample in reader: + assert sample.mask_path is not None + assert sample.mask_path == ade20k_root / "annotations" / f"{sample.image_path.stem}.png" + assert sample.mask_path.exists() diff --git a/model_converter/tests/unit/datasets/test_class_folder.py b/model_converter/tests/unit/datasets/test_class_folder.py new file mode 100644 index 000000000..841d42520 --- /dev/null +++ b/model_converter/tests/unit/datasets/test_class_folder.py @@ -0,0 +1,71 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Unit tests for ClassFolderReader.""" + +from __future__ import annotations + +import pytest +from model_converter.datasets import ClassFolderReader + + +@pytest.fixture +def class_folder_root(tmp_path): + """Build a minimal class-folder dataset (3 classes, 2 images each).""" + for class_id in (0, 1, 5): + class_dir = tmp_path / str(class_id) + class_dir.mkdir() + (class_dir / "img_a.JPEG").write_bytes(b"") + (class_dir / "img_b.png").write_bytes(b"") + return tmp_path + + +def test_enumerates_all_images_with_integer_labels(class_folder_root): + reader = ClassFolderReader(class_folder_root) + samples = list(reader) + + assert len(samples) == 6 + labels = sorted({s.label for s in samples}) + assert labels == [0, 1, 5] + assert all(s.image_path.exists() for s in samples) + + +def test_label_matches_class_folder_name(class_folder_root): + reader = ClassFolderReader(class_folder_root) + for sample in reader: + assert int(sample.image_path.parent.name) == sample.label + + +def test_ignores_non_directory_entries_at_root(class_folder_root): + (class_folder_root / "stray_file.txt").write_text("noise") + reader = ClassFolderReader(class_folder_root) + samples = list(reader) + assert len(samples) == 6 # stray file ignored + + +def test_supports_multiple_image_extensions(class_folder_root): + (class_folder_root / "0" / "img_c.jpg").write_bytes(b"") + reader = ClassFolderReader(class_folder_root) + samples = list(reader) + assert len(samples) == 7 + + +def test_raises_when_class_folder_name_not_integer(tmp_path): + bad = tmp_path / "not_an_int" + bad.mkdir() + (bad / "img.jpg").write_bytes(b"") + reader = ClassFolderReader(tmp_path) + with pytest.raises(ValueError, match="invalid literal"): + list(reader) + + +def test_empty_root_returns_no_samples(tmp_path): + reader = ClassFolderReader(tmp_path) + assert list(reader) == [] + + +def test_nonexistent_root_returns_no_samples(tmp_path): + reader = ClassFolderReader(tmp_path / "does_not_exist") + assert list(reader) == [] diff --git a/model_converter/tests/unit/datasets/test_coco.py b/model_converter/tests/unit/datasets/test_coco.py new file mode 100644 index 000000000..4288d98c3 --- /dev/null +++ b/model_converter/tests/unit/datasets/test_coco.py @@ -0,0 +1,89 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Unit tests for the COCO image reader.""" + +from __future__ import annotations + +import json + +import pytest +from model_converter.datasets import CocoImagesReader + + +@pytest.fixture +def coco_root(tmp_path): + images = tmp_path / "images" + annotations = tmp_path / "annotations" + images.mkdir() + annotations.mkdir() + for name in ("000001.jpg", "000002.jpg", "000003.jpg"): + (images / name).write_bytes(b"") + coco_payload = { + "images": [ + {"id": 1, "file_name": "000001.jpg"}, + {"id": 2, "file_name": "000002.jpg"}, + {"id": 3, "file_name": "000003.jpg"}, + ], + "annotations": [], + "categories": [], + } + (annotations / "instances_val2017.json").write_text(json.dumps(coco_payload)) + return tmp_path + + +def test_enumerates_every_image_in_images_subdir(coco_root): + reader = CocoImagesReader(coco_root, annotation_filename="instances_val2017.json") + samples = list(reader) + assert {s.image_path.name for s in samples} == {"000001.jpg", "000002.jpg", "000003.jpg"} + + +def test_labels_are_placeholder_zero(coco_root): + """COCO calibration does not carry a single classification label; reader uses 0.""" + reader = CocoImagesReader(coco_root, annotation_filename="instances_val2017.json") + assert all(s.label == 0 for s in reader) + + +def test_sorted_for_determinism(coco_root): + reader = CocoImagesReader(coco_root, annotation_filename="instances_val2017.json") + names = [s.image_path.name for s in reader] + assert names == sorted(names) + + +def test_raises_when_images_directory_missing(tmp_path): + (tmp_path / "annotations").mkdir() + reader = CocoImagesReader(tmp_path, annotation_filename="instances_val2017.json") + with pytest.raises(FileNotFoundError, match="images"): + list(reader) + + +def test_raises_when_annotation_file_missing(tmp_path): + (tmp_path / "images").mkdir() + (tmp_path / "annotations").mkdir() + reader = CocoImagesReader(tmp_path, annotation_filename="instances_val2017.json") + with pytest.raises(FileNotFoundError, match="annotations"): + list(reader) + + +def test_image_id_populated_from_annotation_json(coco_root): + reader = CocoImagesReader(coco_root, annotation_filename="instances_val2017.json") + samples = {s.image_path.name: s for s in reader} + assert samples["000001.jpg"].image_id == 1 + assert samples["000002.jpg"].image_id == 2 + assert samples["000003.jpg"].image_id == 3 + + +def test_image_id_is_none_when_filename_not_in_annotations(tmp_path): + images = tmp_path / "images" + annotations = tmp_path / "annotations" + images.mkdir() + annotations.mkdir() + (images / "missing_from_json.jpg").write_bytes(b"") + (annotations / "instances_val2017.json").write_text( + json.dumps({"images": [], "annotations": [], "categories": []}), + ) + reader = CocoImagesReader(tmp_path, annotation_filename="instances_val2017.json") + samples = list(reader) + assert samples[0].image_id is None diff --git a/model_converter/tests/unit/datasets/test_factory.py b/model_converter/tests/unit/datasets/test_factory.py new file mode 100644 index 000000000..daf47b149 --- /dev/null +++ b/model_converter/tests/unit/datasets/test_factory.py @@ -0,0 +1,43 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Unit tests for the dataset reader factory.""" + +from __future__ import annotations + +import pytest +from model_converter.datasets import ( + Ade20kReader, + ClassFolderReader, + CocoImagesReader, + reader_for, +) + + +@pytest.mark.parametrize("dataset_type", ["imagenet-1k", "imagenet-21k", None]) +def test_class_folder_used_for_classification_types(tmp_path, dataset_type): + reader = reader_for(dataset_type, tmp_path) + assert isinstance(reader, ClassFolderReader) + + +@pytest.mark.parametrize("dataset_type", ["coco-detection"]) +def test_coco_reader_used_for_coco_types(tmp_path, dataset_type): + reader = reader_for(dataset_type, tmp_path) + assert isinstance(reader, CocoImagesReader) + + +def test_ade20k_reader_used_for_ade20k_type(tmp_path): + reader = reader_for("ade20k", tmp_path) + assert isinstance(reader, Ade20kReader) + + +def test_coco_detection_picks_instances_annotation_file(tmp_path): + reader = reader_for("coco-detection", tmp_path) + assert reader.annotation_filename == "instances_val2017.json" + + +def test_unknown_dataset_type_raises_value_error(tmp_path): + with pytest.raises(ValueError, match="Unknown dataset_type"): + reader_for("does-not-exist", tmp_path) diff --git a/model_converter/tests/unit/metrics/__init__.py b/model_converter/tests/unit/metrics/__init__.py new file mode 100644 index 000000000..d22239848 --- /dev/null +++ b/model_converter/tests/unit/metrics/__init__.py @@ -0,0 +1,4 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# diff --git a/model_converter/tests/unit/metrics/test_classification.py b/model_converter/tests/unit/metrics/test_classification.py new file mode 100644 index 000000000..e56ddb9b9 --- /dev/null +++ b/model_converter/tests/unit/metrics/test_classification.py @@ -0,0 +1,38 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Tests for the top-1 classification metric.""" + +import numpy as np +import pytest +from model_converter.metrics.classification import TopOneAccuracy + + +class TestTopOneAccuracy: + def test_name_is_top1(self): + assert TopOneAccuracy().name == "top1" + + def test_perfect_predictions(self): + metric = TopOneAccuracy() + # logits with argmax matching label + metric.update(np.array([[0.1, 0.7, 0.2]]), 1) + metric.update(np.array([[0.9, 0.05, 0.05]]), 0) + assert metric.compute() == pytest.approx(1.0) + + def test_partial_correct(self): + metric = TopOneAccuracy() + metric.update(np.array([[0.1, 0.7, 0.2]]), 1) # correct + metric.update(np.array([[0.9, 0.05, 0.05]]), 1) # wrong + assert metric.compute() == pytest.approx(0.5) + + def test_reset_clears_counts(self): + metric = TopOneAccuracy() + metric.update(np.array([[0.1, 0.7, 0.2]]), 1) + metric.reset() + metric.update(np.array([[0.9, 0.05, 0.05]]), 1) + assert metric.compute() == pytest.approx(0.0) + + def test_compute_without_updates_returns_zero(self): + assert TopOneAccuracy().compute() == pytest.approx(0.0) diff --git a/model_converter/tests/unit/metrics/test_coco_detection.py b/model_converter/tests/unit/metrics/test_coco_detection.py new file mode 100644 index 000000000..c7ce23c94 --- /dev/null +++ b/model_converter/tests/unit/metrics/test_coco_detection.py @@ -0,0 +1,94 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Tests for the COCO detection mAP metric (wraps pycocotools).""" + +import json +from pathlib import Path + +import pytest +from model_converter.metrics.coco_detection import COCO80_TO_COCO91, CocoDetectionMAP + + +def _write_minimal_coco_gt(path: Path) -> None: + """Write a tiny COCO-format ground truth file with one image and one box.""" + gt = { + "images": [{"id": 1, "width": 100, "height": 100, "file_name": "img1.jpg"}], + "categories": [{"id": 1, "name": "obj", "supercategory": "none"}], + "annotations": [ + { + "id": 1, + "image_id": 1, + "category_id": 1, + "bbox": [10, 10, 20, 20], # [x, y, w, h] + "area": 400, + "iscrowd": 0, + }, + ], + } + path.write_text(json.dumps(gt)) + + +class TestCocoDetectionMAP: + def test_name_for_bbox_iou_type(self): + assert CocoDetectionMAP(annotation_file=Path("/nonexistent"), iou_type="bbox").name == "mAP" + + def test_perfect_detection_scores_one(self, tmp_path): + gt_file = tmp_path / "gt.json" + _write_minimal_coco_gt(gt_file) + metric = CocoDetectionMAP(annotation_file=gt_file, iou_type="bbox") + # Predict the exact ground-truth box with score 1.0 + metric.update( + predictions=[ + {"image_id": 1, "category_id": 1, "bbox": [10, 10, 20, 20], "score": 1.0}, + ], + ) + assert metric.compute() == pytest.approx(1.0, abs=1e-3) + + def test_no_predictions_scores_zero(self, tmp_path): + gt_file = tmp_path / "gt.json" + _write_minimal_coco_gt(gt_file) + metric = CocoDetectionMAP(annotation_file=gt_file, iou_type="bbox") + # Update with no predictions at all + metric.update(predictions=[]) + assert metric.compute() == pytest.approx(0.0) + + def test_reset_clears_predictions(self, tmp_path): + gt_file = tmp_path / "gt.json" + _write_minimal_coco_gt(gt_file) + metric = CocoDetectionMAP(annotation_file=gt_file, iou_type="bbox") + metric.update(predictions=[{"image_id": 1, "category_id": 1, "bbox": [10, 10, 20, 20], "score": 1.0}]) + metric.reset() + metric.update(predictions=[]) + assert metric.compute() == pytest.approx(0.0) + + def test_invalid_iou_type_raises(self): + with pytest.raises(ValueError, match="iou_type"): + CocoDetectionMAP(annotation_file=Path("/nonexistent"), iou_type="invalid") + + def test_update_with_none_predictions_is_a_noop(self, tmp_path): + gt_file = tmp_path / "gt.json" + _write_minimal_coco_gt(gt_file) + metric = CocoDetectionMAP(annotation_file=gt_file, iou_type="bbox") + metric.update(predictions=None) + assert metric.compute() == pytest.approx(0.0) + + +class TestCoco80ToCoco91: + def test_length_is_80(self): + assert len(COCO80_TO_COCO91) == 80 + + def test_no_duplicates(self): + assert len(set(COCO80_TO_COCO91)) == 80 + + def test_first_value_is_person(self): + assert COCO80_TO_COCO91[0] == 1 # person + + def test_last_value_is_90(self): + assert COCO80_TO_COCO91[-1] == 90 # toothbrush + + def test_stop_sign_maps_to_13_not_12(self): + """Class 11 (stop sign) must be COCO ID 13; COCO ID 12 (street sign) is absent.""" + assert COCO80_TO_COCO91[11] == 13 diff --git a/model_converter/tests/unit/metrics/test_factory.py b/model_converter/tests/unit/metrics/test_factory.py new file mode 100644 index 000000000..1a6f4cdc7 --- /dev/null +++ b/model_converter/tests/unit/metrics/test_factory.py @@ -0,0 +1,142 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Tests for the metric dispatcher.""" + +from pathlib import Path + +from model_converter.metrics import ( + CocoDetectionMAP, + MultilabelMAP, + SemSegMIoU, + TopOneAccuracy, + metric_for, +) + + +class TestMetricFor: + def test_imagenet_classification_returns_top1(self): + metric = metric_for(dataset_type="imagenet-1k", model_type="Classification") + assert isinstance(metric, TopOneAccuracy) + + def test_imagenet21k_classification_returns_top1(self): + metric = metric_for(dataset_type="imagenet-21k", model_type="Classification") + assert isinstance(metric, TopOneAccuracy) + + def test_multilabel_classification_returns_map(self): + metric = metric_for(dataset_type="imagenet-1k", model_type="Classification_Multilabel") + assert isinstance(metric, MultilabelMAP) + + def test_coco_detection_returns_bbox_map(self, tmp_path): + ann = tmp_path / "ann.json" + ann.write_text('{"images":[],"annotations":[],"categories":[]}') + metric = metric_for( + dataset_type="coco-detection", + model_type="SSD", + annotation_file=ann, + ) + assert isinstance(metric, CocoDetectionMAP) + assert metric.iou_type == "bbox" + + def test_ade20k_returns_miou(self): + metric = metric_for(dataset_type="ade20k", model_type="Segmentation") + assert isinstance(metric, SemSegMIoU) + assert metric.num_classes == 150 + assert metric.ignore_index == 255 + + def test_unknown_dataset_returns_none(self): + assert metric_for(dataset_type="not-a-dataset", model_type="Classification") is None + + def test_none_dataset_returns_none(self): + assert metric_for(dataset_type=None, model_type="Classification") is None + + def test_coco_without_annotation_file_returns_none(self): + # Without annotation_file, COCO metric cannot be constructed → None + assert metric_for(dataset_type="coco-detection", model_type="SSD") is None + + def test_dinov2_feature_extractor_returns_none(self): + # vit_small_patch14_dinov2.lvd142m has no classification head + assert metric_for(dataset_type="imagenet-1k", model_type="FeatureExtractor") is None + + def test_unknown_model_type_with_imagenet_returns_none(self): + # An unknown model_type with imagenet should not crash; returns None + assert metric_for(dataset_type="imagenet-1k", model_type=None) is None + + def test_coco_metric_uses_provided_annotation_file(self, tmp_path): + ann = tmp_path / "ann.json" + ann.write_text('{"images":[],"annotations":[],"categories":[]}') + metric = metric_for( + dataset_type="coco-detection", + model_type="SSD", + annotation_file=ann, + ) + assert isinstance(metric, CocoDetectionMAP) + # internal field — used to ensure dispatcher routes the path through + assert Path(metric.annotation_file) == ann + + def test_ade20k_with_non_segmentation_model_returns_none(self): + assert metric_for(dataset_type="ade20k", model_type="Classification") is None + + def test_multilabel_via_getitune_task(self): + # Real configs report model_type="Classification" with getitune_task="MULTI_LABEL_CLS" + metric = metric_for( + dataset_type="imagenet-1k", + model_type="Classification", + task="MULTI_LABEL_CLS", + ) + assert isinstance(metric, MultilabelMAP) + + def test_multi_class_cls_task_returns_top1(self): + metric = metric_for( + dataset_type="imagenet-1k", + model_type="Classification", + task="MULTI_CLASS_CLS", + ) + assert isinstance(metric, TopOneAccuracy) + + def test_yolo11_bbox_map(self, tmp_path): + ann = tmp_path / "ann.json" + ann.write_text('{"images":[],"annotations":[],"categories":[]}') + metric = metric_for( + dataset_type="coco-detection", + model_type="YOLO11", + annotation_file=ann, + ) + assert isinstance(metric, CocoDetectionMAP) + assert metric.iou_type == "bbox" + + def test_yolov5_bbox_map(self, tmp_path): + ann = tmp_path / "ann.json" + ann.write_text('{"images":[],"annotations":[],"categories":[]}') + metric = metric_for( + dataset_type="coco-detection", + model_type="YOLOv5", + annotation_file=ann, + ) + assert isinstance(metric, CocoDetectionMAP) + + def test_yolov8_bbox_map(self, tmp_path): + ann = tmp_path / "ann.json" + ann.write_text('{"images":[],"annotations":[],"categories":[]}') + metric = metric_for( + dataset_type="coco-detection", + model_type="YOLOv8", + annotation_file=ann, + ) + assert isinstance(metric, CocoDetectionMAP) + + def test_rotated_detection_task_returns_none(self, tmp_path): + # Rotated detection has no standard metric implemented here — skip + ann = tmp_path / "ann.json" + ann.write_text('{"images":[],"annotations":[],"categories":[]}') + assert ( + metric_for( + dataset_type="coco-detection", + model_type="MaskRCNN", + annotation_file=ann, + task="ROTATED_DETECTION", + ) + is None + ) diff --git a/model_converter/tests/unit/metrics/test_multilabel.py b/model_converter/tests/unit/metrics/test_multilabel.py new file mode 100644 index 000000000..371b9f94e --- /dev/null +++ b/model_converter/tests/unit/metrics/test_multilabel.py @@ -0,0 +1,49 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Tests for the multilabel mean Average Precision metric.""" + +import numpy as np +import pytest +from model_converter.metrics.multilabel import MultilabelMAP + + +class TestMultilabelMAP: + def test_name_is_map(self): + assert MultilabelMAP(num_labels=3).name == "mAP" + + def test_perfect_predictions_score_one(self): + metric = MultilabelMAP(num_labels=3) + # logits clearly favour the positive label per sample + metric.update(np.array([[10.0, -10.0, -10.0]]), np.array([1, 0, 0])) + metric.update(np.array([[-10.0, 10.0, -10.0]]), np.array([0, 1, 0])) + metric.update(np.array([[-10.0, -10.0, 10.0]]), np.array([0, 0, 1])) + assert metric.compute() == pytest.approx(1.0, abs=1e-3) + + def test_random_predictions_below_perfect(self): + rng = np.random.default_rng(0) + metric = MultilabelMAP(num_labels=5) + for _ in range(20): + logits = rng.normal(size=(1, 5)).astype(np.float32) + gt = rng.integers(0, 2, size=5) + metric.update(logits, gt) + score = metric.compute() + assert 0.0 <= score < 1.0 + + def test_reset_clears_state(self): + metric = MultilabelMAP(num_labels=2) + metric.update(np.array([[10.0, -10.0]]), np.array([1, 0])) + metric.reset() + metric.update(np.array([[-10.0, 10.0]]), np.array([1, 0])) + # After reset, predicting wrong label every time → score is low + assert metric.compute() < 0.6 + + def test_accepts_1d_logits_input(self): + """Per-sample logits may be passed as a 1D vector; the metric promotes to 2D.""" + metric = MultilabelMAP(num_labels=3) + metric.update(np.array([10.0, -10.0, -10.0]), np.array([1, 0, 0])) + metric.update(np.array([-10.0, 10.0, -10.0]), np.array([0, 1, 0])) + metric.update(np.array([-10.0, -10.0, 10.0]), np.array([0, 0, 1])) + assert metric.compute() == pytest.approx(1.0, abs=1e-3) diff --git a/model_converter/tests/unit/metrics/test_semseg_miou.py b/model_converter/tests/unit/metrics/test_semseg_miou.py new file mode 100644 index 000000000..c3c262ec3 --- /dev/null +++ b/model_converter/tests/unit/metrics/test_semseg_miou.py @@ -0,0 +1,59 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Tests for the semantic segmentation mIoU metric.""" + +import numpy as np +import pytest +from model_converter.metrics.semseg_miou import SemSegMIoU + + +class TestSemSegMIoU: + def test_name_is_miou(self): + assert SemSegMIoU(num_classes=3).name == "mIoU" + + def test_perfect_prediction(self): + metric = SemSegMIoU(num_classes=3) + mask = np.array([[0, 1], [2, 0]]) + metric.update(mask, mask) + assert metric.compute() == pytest.approx(1.0) + + def test_completely_wrong_prediction(self): + metric = SemSegMIoU(num_classes=2) + gt = np.array([[0, 0], [0, 0]]) + pred = np.array([[1, 1], [1, 1]]) + metric.update(pred, gt) + # No intersection for any class present → 0 + assert metric.compute() == pytest.approx(0.0) + + def test_partial_overlap(self): + # 2 classes, 4 pixels. GT = [0,0,1,1], Pred = [0,1,1,1]. + # Class 0: TP=1, FP=0, FN=1 → IoU = 1/2 = 0.5 + # Class 1: TP=2, FP=1, FN=0 → IoU = 2/3 + # mIoU = (0.5 + 2/3) / 2 ≈ 0.5833 + metric = SemSegMIoU(num_classes=2) + gt = np.array([0, 0, 1, 1]) + pred = np.array([0, 1, 1, 1]) + metric.update(pred, gt) + assert metric.compute() == pytest.approx((0.5 + 2 / 3) / 2) + + def test_ignore_index_excluded(self): + metric = SemSegMIoU(num_classes=2, ignore_index=255) + gt = np.array([0, 0, 255, 1]) + # If ignored pixel were counted, prediction "1" at index 2 would be FP for class 1. + pred = np.array([0, 0, 1, 1]) + metric.update(pred, gt) + # With ignore: class 0 IoU = 2/2 = 1.0, class 1 IoU = 1/1 = 1.0 → mIoU = 1.0 + assert metric.compute() == pytest.approx(1.0) + + def test_reset_clears_confusion_matrix(self): + metric = SemSegMIoU(num_classes=2) + metric.update(np.array([1, 1]), np.array([0, 0])) + metric.reset() + metric.update(np.array([0, 0]), np.array([0, 0])) + assert metric.compute() == pytest.approx(1.0) + + def test_compute_without_updates_returns_zero(self): + assert SemSegMIoU(num_classes=2).compute() == pytest.approx(0.0) diff --git a/model_converter/tests/unit/test_cli.py b/model_converter/tests/unit/test_cli.py index 7cc5b8ae0..02d9555d7 100644 --- a/model_converter/tests/unit/test_cli.py +++ b/model_converter/tests/unit/test_cli.py @@ -3,23 +3,26 @@ # SPDX-License-Identifier: Apache-2.0 # -"""Tests for model_converter.cli module.""" +"""Tests for the model_converter CLI facade and converter hierarchy.""" import json import sys from pathlib import Path from unittest.mock import MagicMock, patch -import cv2 import numpy as np import pytest import torch import torch.nn as nn from model_converter.cli import ModelConverter, list_models, main +from model_converter.converters.getitune import GetituneConverter +from model_converter.converters.timm import TimmConverter +from model_converter.converters.torchvision import TorchvisionConverter +from model_converter.reporting import AccuracyResults class TestModelConverterInit: - """Tests for ModelConverter.__init__.""" + """Tests for facade ModelConverter.__init__.""" def test_creates_directories(self, tmp_path): """ModelConverter creates output and cache directories.""" @@ -41,55 +44,100 @@ def test_verbose_logging(self, tmp_path): assert converter.logger is not None def test_dataset_path(self, tmp_path): - """ModelConverter stores dataset path.""" + """ModelConverter stores dataset registry.""" dataset = tmp_path / "dataset" converter = ModelConverter( output_dir=tmp_path / "out", cache_dir=tmp_path / "cache", - dataset_path=dataset, + dataset_registry=dataset, ) - assert converter.dataset_path == dataset + assert converter.dataset_registry == dataset def test_dataset_path_none(self, tmp_path): - """ModelConverter handles None dataset path.""" + """ModelConverter handles None dataset registry.""" converter = ModelConverter( output_dir=tmp_path / "out", cache_dir=tmp_path / "cache", - dataset_path=None, + dataset_registry=None, ) - assert converter.dataset_path is None + assert converter.dataset_registry is None + + +class TestModelConverterDispatch: + """Tests for facade dispatching logic.""" + + def test_get_converter_returns_torchvision_and_caches(self, facade_converter): + """_get_converter caches converter instances per library.""" + first = facade_converter._get_converter("torchvision") + second = facade_converter._get_converter("torchvision") + + assert isinstance(first, TorchvisionConverter) + assert first is second + + def test_get_converter_passes_training_extensions_dir(self, tmp_path): + """_get_converter wires training_extensions_dir for getitune.""" + training_extensions_dir = tmp_path / "training_extensions" + converter = ModelConverter( + output_dir=tmp_path / "out", + cache_dir=tmp_path / "cache", + training_extensions_dir=training_extensions_dir, + ) + + getitune_converter = converter._get_converter("getitune") + + assert isinstance(getitune_converter, GetituneConverter) + assert getitune_converter.training_extensions_dir == training_extensions_dir + + def test_get_converter_rejects_unknown_library(self, facade_converter): + """_get_converter raises for unsupported libraries.""" + with pytest.raises(ValueError, match="Unsupported model_library"): + facade_converter._get_converter("unknown") + + def test_process_model_config_defaults_to_torchvision(self, facade_converter): + """process_model_config uses torchvision when model_library is omitted.""" + config = {"model_short_name": "test_model"} + + with patch.object(TorchvisionConverter, "process_model_config", return_value=True) as mock_process: + result = facade_converter.process_model_config(config) + + assert result is True + mock_process.assert_called_once_with(config) + + def test_process_model_config_routes_to_timm(self, facade_converter): + """process_model_config dispatches by model_library.""" + config = {"model_short_name": "test_model", "model_library": "timm"} + + with patch.object(TimmConverter, "process_model_config", return_value=True) as mock_process: + result = facade_converter.process_model_config(config) + + assert result is True + mock_process.assert_called_once_with(config) + + def test_process_model_config_returns_false_for_unknown_library(self, facade_converter): + """process_model_config returns False when dispatch fails.""" + assert facade_converter.process_model_config({"model_library": "unsupported"}) is False class TestGetLabels: - """Tests for ModelConverter.get_labels.""" + """Tests for PyTorchConverter.get_labels via TorchvisionConverter.""" def test_imagenet1k_v1(self, converter): - """get_labels returns ImageNet1K labels.""" - mock_categories = ["tabby cat", "golden retriever", "great white shark"] - with ( - patch("model_converter.cli.importlib") as _, - patch.dict( - "sys.modules", - {"torchvision.models._meta": MagicMock(_IMAGENET_CATEGORIES=mock_categories)}, - ), - patch("model_converter.cli.ModelConverter.get_labels", wraps=converter.get_labels), + """get_labels returns ImageNet1K labels with underscores.""" + with patch( + "torchvision.models._meta._IMAGENET_CATEGORIES", + ["tabby cat", "golden retriever"], + create=True, ): - # Directly test the code path - pass - - # Test with actual mocking of the import - with patch("torchvision.models._meta._IMAGENET_CATEGORIES", ["tabby cat", "golden retriever"], create=True): result = converter.get_labels("IMAGENET1K_V1") - assert result is not None - assert " " not in result.split()[0] or "_" in result.split()[0] # spaces replaced with underscores + + assert result == "tabby_cat golden_retriever" def test_imagenet21k(self, converter): """get_labels returns ImageNet21K labels.""" mock_info = MagicMock() mock_info.label_descriptions.return_value = ["tabby, tabby cat", "golden retriever, dog"] - mock_imagenet_info_cls = MagicMock(return_value=mock_info) - with patch("timm.data.ImageNetInfo", mock_imagenet_info_cls): + with patch("timm.data.ImageNetInfo", return_value=mock_info): result = converter.get_labels("IMAGENET21K") assert result == "tabby golden_retriever" @@ -99,27 +147,22 @@ def test_coco_v1(self, converter): mock_weights = MagicMock() mock_weights.COCO_V1.meta = {"categories": ["person", "bicycle", "car"]} - with patch( - "torchvision.models.detection.MaskRCNN_ResNet50_FPN_Weights", - mock_weights, - ): + with patch("torchvision.models.detection.MaskRCNN_ResNet50_FPN_Weights", mock_weights): result = converter.get_labels("COCO_V1") assert result == "person bicycle car" def test_unknown_label_set(self, converter): """get_labels returns None for unknown label sets.""" - result = converter.get_labels("NONEXISTENT_LABELS") - assert result is None + assert converter.get_labels("NONEXISTENT_LABELS") is None class TestLoadModelClass: - """Tests for ModelConverter.load_model_class.""" + """Tests for PyTorchConverter.load_model_class via TorchvisionConverter.""" def test_successful_import(self, converter): """load_model_class dynamically imports a class.""" - result = converter.load_model_class("torch.nn.Linear") - assert result is torch.nn.Linear + assert converter.load_model_class("torch.nn.Linear") is torch.nn.Linear def test_import_failure(self, converter): """load_model_class raises on invalid path.""" @@ -128,7 +171,7 @@ def test_import_failure(self, converter): class TestLoadCheckpoint: - """Tests for ModelConverter.load_checkpoint.""" + """Tests for PyTorchConverter.load_checkpoint via TorchvisionConverter.""" def test_successful_load(self, converter, tmp_path): """load_checkpoint loads a PyTorch checkpoint.""" @@ -136,116 +179,222 @@ def test_successful_load(self, converter, tmp_path): checkpoint_path.touch() state_dict = {"layer.weight": torch.randn(10, 10)} - with patch("torch.load", return_value=state_dict): + with patch("model_converter.converters.pytorch.torch.load", return_value=state_dict) as mock_load: result = converter.load_checkpoint(checkpoint_path) - assert "layer.weight" in result + assert result == state_dict + mock_load.assert_called_once_with(checkpoint_path, map_location="cpu", weights_only=True) def test_load_failure(self, converter, tmp_path): - """load_checkpoint raises on invalid file.""" - bad_path = tmp_path / "nonexistent.pth" - with pytest.raises(FileNotFoundError): - converter.load_checkpoint(bad_path) + """load_checkpoint re-raises underlying torch loading errors.""" + with ( + patch( + "model_converter.converters.pytorch.torch.load", + side_effect=FileNotFoundError("missing"), + ), + pytest.raises(FileNotFoundError, match="missing"), + ): + converter.load_checkpoint(tmp_path / "nonexistent.pth") class TestLoadHuggingfaceModel: - """Tests for ModelConverter.load_huggingface_model.""" + """Tests for TimmConverter.load_huggingface_model.""" - def test_timm_model(self, converter): + def test_timm_model(self, timm_converter): """load_huggingface_model loads timm model.""" mock_model = MagicMock() mock_model.eval.return_value = mock_model with patch("timm.create_model", return_value=mock_model) as mock_create: - result = converter.load_huggingface_model( + result = timm_converter.load_huggingface_model( repo_id="timm/resnet50", revision="abc123", model_library="timm", ) assert result is mock_model - mock_create.assert_called_once() + mock_create.assert_called_once_with( + "hf-hub:timm/resnet50@abc123", + pretrained=True, + cache_dir=timm_converter.cache_dir, + ) mock_model.eval.assert_called_once() - def test_timm_model_with_params(self, converter): + def test_timm_model_with_params(self, timm_converter): """load_huggingface_model passes model_params to timm.""" mock_model = MagicMock() mock_model.eval.return_value = mock_model with patch("timm.create_model", return_value=mock_model) as mock_create: - converter.load_huggingface_model( + timm_converter.load_huggingface_model( repo_id="timm/resnet50", revision="abc123", model_library="timm", model_params={"num_classes": 10}, ) - call_kwargs = mock_create.call_args[1] - assert call_kwargs["num_classes"] == 10 + assert mock_create.call_args.kwargs["num_classes"] == 10 - def test_transformers_model(self, converter): + def test_transformers_model(self, timm_converter): """load_huggingface_model loads transformers model.""" mock_model = MagicMock() mock_model.eval.return_value = mock_model with patch("transformers.AutoModel.from_pretrained", return_value=mock_model) as mock_from: - result = converter.load_huggingface_model( + result = timm_converter.load_huggingface_model( repo_id="bert-base-uncased", revision="abc123", model_library="transformers", ) assert result is mock_model - mock_from.assert_called_once() + mock_from.assert_called_once_with( + "bert-base-uncased", + revision="abc123", + cache_dir=timm_converter.cache_dir, + ) - def test_unsupported_library(self, converter): + def test_unsupported_library(self, timm_converter): """load_huggingface_model raises for unsupported library.""" with pytest.raises(ValueError, match="Unsupported model library"): - converter.load_huggingface_model( + timm_converter.load_huggingface_model( repo_id="some/model", revision="abc123", model_library="unsupported_lib", ) - def test_load_failure(self, converter): - """load_huggingface_model raises on failure.""" + def test_load_failure(self, timm_converter): + """load_huggingface_model re-raises loading errors.""" with ( patch("timm.create_model", side_effect=RuntimeError("Connection error")), pytest.raises(RuntimeError, match="Connection error"), ): - converter.load_huggingface_model( + timm_converter.load_huggingface_model( repo_id="timm/resnet50", revision="abc123", model_library="timm", ) +class TestApplyTimmDataConfig: + """Tests for TimmConverter._apply_timm_data_config.""" + + def test_overrides_mean_scale_and_shape(self, timm_converter): + """Resolved timm mean/std (0..1) are scaled to 0..255 and shape is applied.""" + model = MagicMock() + config = { + "mean_values": "123.675 116.28 103.53", + "scale_values": "58.395 57.12 57.375", + "input_shape": [1, 3, 224, 224], + } + resolved = { + "mean": (0.5, 0.5, 0.5), + "std": (0.5, 0.5, 0.5), + "input_size": (3, 256, 256), + } + + with patch("timm.data.resolve_data_config", return_value=resolved): + timm_converter._apply_timm_data_config(model, config) + + assert config["mean_values"] == "127.5 127.5 127.5" + assert config["scale_values"] == "127.5 127.5 127.5" + assert config["input_shape"] == [1, 3, 256, 256] + assert config["reverse_input_channels"] is True + + def test_keeps_imagenet_values_when_model_matches(self, timm_converter): + """Standard ImageNet normalization round-trips to the canonical strings.""" + model = MagicMock() + config = {} + resolved = { + "mean": (0.485, 0.456, 0.406), + "std": (0.229, 0.224, 0.225), + "input_size": (3, 224, 224), + } + + with patch("timm.data.resolve_data_config", return_value=resolved): + timm_converter._apply_timm_data_config(model, config) + + assert config["mean_values"] == "123.675 116.28 103.53" + assert config["scale_values"] == "58.395 57.12 57.375" + assert config["input_shape"] == [1, 3, 224, 224] + + def test_ignores_missing_fields(self, timm_converter): + """Fields timm does not provide leave the existing config untouched.""" + model = MagicMock() + config = {"mean_values": "1 2 3", "scale_values": "4 5 6", "input_shape": [1, 3, 8, 8]} + + with patch("timm.data.resolve_data_config", return_value={}): + timm_converter._apply_timm_data_config(model, config) + + assert config == { + "mean_values": "1 2 3", + "scale_values": "4 5 6", + "input_shape": [1, 3, 8, 8], + "reverse_input_channels": True, + } + + def test_forces_reverse_input_channels(self, timm_converter): + """reverse_input_channels is always forced to True for timm (RGB) models.""" + model = MagicMock() + config = {"reverse_input_channels": False} + + with patch("timm.data.resolve_data_config", return_value={}): + timm_converter._apply_timm_data_config(model, config) + + assert config["reverse_input_channels"] is True + + def test_ignores_malformed_input_size(self, timm_converter): + """A non 3-tuple input_size is ignored, leaving input_shape untouched.""" + model = MagicMock() + config = {"input_shape": [1, 3, 8, 8]} + + with patch("timm.data.resolve_data_config", return_value={"input_size": (224, 224)}): + timm_converter._apply_timm_data_config(model, config) + + assert config["input_shape"] == [1, 3, 8, 8] + + def test_handles_resolution_failure(self, timm_converter): + """A failure inside resolve_data_config keeps the configured values.""" + model = MagicMock() + config = {"mean_values": "1 2 3"} + + with patch("timm.data.resolve_data_config", side_effect=RuntimeError("boom")): + timm_converter._apply_timm_data_config(model, config) + + assert config["mean_values"] == "1 2 3" + + def test_handles_missing_timm(self, timm_converter): + """If timm.data cannot be imported, configured values are kept.""" + model = MagicMock() + config = {"mean_values": "1 2 3"} + + with patch.dict(sys.modules, {"timm.data": None}): + timm_converter._apply_timm_data_config(model, config) + + assert config["mean_values"] == "1 2 3" + assert config["reverse_input_channels"] is True + + class TestCreateModel: - """Tests for ModelConverter.create_model.""" + """Tests for PyTorchConverter.create_model via TorchvisionConverter.""" def test_nn_module_with_model_key(self, converter): - """create_model extracts model from checkpoint 'model' key.""" + """create_model extracts model from checkpoint model key.""" mock_model = MagicMock(spec=nn.Module) mock_model.eval.return_value = mock_model - checkpoint = {"model": mock_model} - result = converter.create_model(torch.nn.Module, checkpoint) + result = converter.create_model(torch.nn.Module, {"model": mock_model}) assert result is mock_model def test_nn_module_with_state_dict_only_raises(self, converter): - """create_model raises when nn.Module with only state_dict.""" - checkpoint = {"state_dict": {"layer.weight": torch.randn(10)}} - + """create_model raises when nn.Module is used with only state_dict.""" with pytest.raises(ValueError, match="state_dict"): - converter.create_model(torch.nn.Module, checkpoint) + converter.create_model(torch.nn.Module, {"state_dict": {"layer.weight": torch.randn(10)}}) def test_nn_module_direct_model_as_checkpoint(self, converter): - """create_model handles nn.Module passed directly (not in dict) gracefully.""" - # When an nn.Module is passed as checkpoint, "model" in checkpoint raises TypeError - # which is caught by the except block and re-raised - model = nn.Linear(10, 10) + """create_model surfaces invalid checkpoint types.""" with pytest.raises(TypeError): - converter.create_model(torch.nn.Module, model) + converter.create_model(torch.nn.Module, nn.Linear(10, 10)) def test_nn_module_invalid_checkpoint(self, converter): """create_model raises when checkpoint is not a valid model.""" @@ -258,13 +407,11 @@ def test_model_class_with_state_dict(self, converter): mock_instance = MagicMock() mock_instance.eval.return_value = mock_instance mock_class.return_value = mock_instance - state_dict = {"layer.weight": torch.randn(10)} - checkpoint = {"state_dict": state_dict} - result = converter.create_model(mock_class, checkpoint) + result = converter.create_model(mock_class, {"state_dict": state_dict}) assert result is mock_instance - mock_instance.load_state_dict.assert_called_once() + mock_instance.load_state_dict.assert_called_once_with(state_dict, strict=False) def test_model_class_with_model_params(self, converter): """create_model passes model_params to class constructor.""" @@ -273,30 +420,29 @@ def test_model_class_with_model_params(self, converter): mock_instance.eval.return_value = mock_instance mock_class.return_value = mock_instance - checkpoint = {"state_dict": {"w": torch.randn(10)}} - - converter.create_model(mock_class, checkpoint, model_params={"num_classes": 5}) + converter.create_model( + mock_class, + {"state_dict": {"w": torch.randn(10)}}, + model_params={"num_classes": 5}, + ) mock_class.assert_called_once_with(num_classes=5) def test_model_class_with_model_key_as_module(self, converter): - """create_model returns checkpoint['model'] if it's an nn.Module.""" + """create_model returns checkpoint model when it already is an nn.Module.""" mock_model = MagicMock(spec=nn.Module) mock_class = MagicMock() - checkpoint = {"model": mock_model} - result = converter.create_model(mock_class, checkpoint) + result = converter.create_model(mock_class, {"model": mock_model}) assert result is mock_model def test_model_class_with_model_key_as_dict(self, converter): - """create_model uses checkpoint['model'] as state_dict.""" + """create_model uses checkpoint model dict as state_dict.""" mock_class = MagicMock() mock_instance = MagicMock() mock_instance.eval.return_value = mock_instance mock_class.return_value = mock_instance - checkpoint = {"model": {"layer.weight": torch.randn(10)}} - - result = converter.create_model(mock_class, checkpoint) + result = converter.create_model(mock_class, {"model": {"layer.weight": torch.randn(10)}}) assert result is mock_instance mock_instance.load_state_dict.assert_called_once() @@ -307,953 +453,346 @@ def test_model_class_bare_state_dict(self, converter): mock_instance.eval.return_value = mock_instance mock_class.return_value = mock_instance - checkpoint = {"layer.weight": torch.randn(10)} - - result = converter.create_model(mock_class, checkpoint) + result = converter.create_model(mock_class, {"layer.weight": torch.randn(10)}) assert result is mock_instance mock_instance.load_state_dict.assert_called_once() def test_create_model_failure(self, converter): """create_model raises on instantiation failure.""" - mock_class = MagicMock(side_effect=RuntimeError("init failed")) - with pytest.raises(RuntimeError, match="init failed"): - converter.create_model(mock_class, {}) + converter.create_model(MagicMock(side_effect=RuntimeError("init failed")), {}) -class TestCopyReadme: - """Tests for ModelConverter.copy_readme.""" - - def test_successful_copy(self, converter, tmp_path): - """copy_readme copies and fills README template.""" - output_folder = tmp_path / "model-fp16-ov" - output_folder.mkdir() +class TestExportToOpenvino: + """Tests for PyTorchConverter.export_to_openvino via TorchvisionConverter.""" - # We mock the template file reading - template_content = "# <>\nLicense: <>\nLink: <>\nDocs: <>" + def test_successful_export(self, converter, sample_model_config): + """export_to_openvino exports model to OV format.""" + mock_model = MagicMock(spec=nn.Module) + mock_model.eval.return_value = mock_model - config = { - "model_short_name": "test_model", - "license": "Apache-2.0", - "license_link": "https://apache.org/licenses/LICENSE-2.0", - "docs": "https://docs.example.com", - "model_library": "timm", - } + mock_ov_model = MagicMock() + mock_input = MagicMock() + mock_input.get_names.return_value = {"input"} + mock_ov_model.input.return_value = mock_input + mock_ov_model.inputs = [mock_input] + mock_ov_model.outputs = [MagicMock()] + mock_ov_model.get_rt_info.return_value = MagicMock(value={"model_type": "Classification"}) with ( + patch("openvino.convert_model", return_value=mock_ov_model), + patch("openvino.save_model") as mock_save, patch.object(Path, "exists", return_value=True), - patch.object(Path, "read_text", return_value=template_content), + patch("model_converter.converters.pytorch.shutil.copy2"), + patch.object(converter, "copy_readme"), ): - converter.copy_readme(config, output_folder, variant="fp16") - - readme = output_folder / "README.md" - assert readme.exists() - content = readme.read_text() - assert "test_model" in content - assert "Apache-2.0" in content - - def test_template_not_found(self, converter, tmp_path): - """copy_readme handles missing template gracefully.""" - output_folder = tmp_path / "model-fp16-ov" - output_folder.mkdir() - - config = { - "model_short_name": "test_model", - "license": "MIT", - "license_link": "https://mit.edu", - "docs": "", - "model_library": "timm", - } + fp16_path, fp32_path = converter.export_to_openvino( + model=mock_model, + input_shape=[1, 3, 224, 224], + output_path=converter.output_dir / "test_model", + model_config=sample_model_config, + input_names=["input"], + output_names=["result"], + metadata={("model_info", "model_type"): "Classification"}, + ) - # Template path doesn't exist - with patch.object(Path, "exists", return_value=False): - converter.copy_readme(config, output_folder, variant="fp16") + assert fp16_path == converter.output_dir / "test_model-fp16-ov" / "test_model.xml" + assert fp32_path == converter.output_dir / "test_model-fp16-ov" / "test_model_fp32.xml" + assert mock_save.call_count == 2 - # No README should be created - assert not (output_folder / "README.md").exists() + def test_export_failure(self, converter, sample_model_config): + """export_to_openvino raises on conversion failure.""" + mock_model = MagicMock(spec=nn.Module) + mock_model.eval.return_value = mock_model - def test_missing_model_short_name(self, converter, tmp_path): - """copy_readme warns when model_short_name is empty.""" - output_folder = tmp_path / "model-fp16-ov" - output_folder.mkdir() + with ( + patch("openvino.convert_model", side_effect=RuntimeError("Conversion failed")), + pytest.raises(RuntimeError, match="Conversion failed"), + ): + converter.export_to_openvino( + model=mock_model, + input_shape=[1, 3, 224, 224], + output_path=converter.output_dir / "test_model", + model_config=sample_model_config, + ) - config = { - "model_short_name": "", - "license": "MIT", - "license_link": "https://mit.edu", - } - # Should not raise but log warning - converter.copy_readme(config, output_folder) +class TestPrepareModelForExport: + """Tests for PyTorchConverter._prepare_model_for_export via TorchvisionConverter.""" - def test_missing_license_link(self, converter, tmp_path): - """copy_readme warns when license_link is empty.""" - output_folder = tmp_path / "model-fp16-ov" - output_folder.mkdir() + def test_with_adapter(self, converter): + """_prepare_model_for_export applies adapter for known model type.""" + mock_model = MagicMock(spec=nn.Module) - config = { - "model_short_name": "test", - "license": "MIT", - "license_link": "", - } + with patch("model_converter.converters.pytorch.get_adapter") as mock_get_adapter: + mock_adapted = MagicMock() + mock_get_adapter.return_value = mock_adapted + result = converter._prepare_model_for_export(mock_model, {"model_type": "MaskRCNN"}) - converter.copy_readme(config, output_folder) + assert result is mock_adapted - def test_missing_license(self, converter, tmp_path): - """copy_readme warns when license is empty.""" - output_folder = tmp_path / "model-fp16-ov" - output_folder.mkdir() + def test_without_adapter(self, converter): + """_prepare_model_for_export returns model unchanged without adapter.""" + mock_model = MagicMock(spec=nn.Module) - config = { - "model_short_name": "test", - "license": "", - "license_link": "https://mit.edu", - } + with patch("model_converter.converters.pytorch.get_adapter", return_value=mock_model): + result = converter._prepare_model_for_export(mock_model, {"model_type": "Classification"}) - converter.copy_readme(config, output_folder) + assert result is mock_model - def test_missing_docs_field(self, converter, tmp_path): - """copy_readme handles missing docs field.""" - output_folder = tmp_path / "model-fp16-ov" - output_folder.mkdir() - template_content = "# <>" - config = { - "model_short_name": "test_model", - "license": "Apache-2.0", - "license_link": "https://apache.org", - "model_library": "timm", - } +class TestCreateExampleInput: + """Tests for PyTorchConverter._create_example_input via TorchvisionConverter.""" - with ( - patch.object(Path, "exists", return_value=True), - patch.object(Path, "read_text", return_value=template_content), - ): - converter.copy_readme(config, output_folder, variant="fp16") + def test_maskrcnn_input(self, converter): + """_create_example_input uses rand for MaskRCNN.""" + result = converter._create_example_input([1, 3, 224, 224], {"model_type": "MaskRCNN"}) + assert result.shape == (1, 3, 224, 224) + assert result.min() >= 0 - def test_none_value_in_config(self, converter, tmp_path): - """copy_readme skips None values in config placeholders.""" - output_folder = tmp_path / "model-fp16-ov" - output_folder.mkdir() + def test_default_input(self, converter): + """_create_example_input uses randn for other model types.""" + assert converter._create_example_input([1, 3, 224, 224], {"model_type": "Classification"}).shape == ( + 1, + 3, + 224, + 224, + ) - template_content = "# <>" - config = { - "model_short_name": "test_model", - "license": "Apache-2.0", - "license_link": "https://apache.org", - "model_library": "timm", - "optional_field": None, - } - with ( - patch.object(Path, "exists", return_value=True), - patch.object(Path, "read_text", return_value=template_content), - ): - converter.copy_readme(config, output_folder, variant="fp16") - - -class TestCollectDatasetEntries: - """Tests for ModelConverter._collect_dataset_entries.""" - - def test_collects_entries(self, converter, dataset_dir): - """_collect_dataset_entries finds images with class labels.""" - entries = converter._collect_dataset_entries(dataset_dir) - assert len(entries) == 2 - # Entries are (path, class_label) tuples - assert entries[0][1] == 0 - assert entries[1][1] == 1 - - def test_empty_directory(self, converter, tmp_path): - """_collect_dataset_entries returns empty list for empty dir.""" - empty_dir = tmp_path / "empty" - empty_dir.mkdir() - entries = converter._collect_dataset_entries(empty_dir) - assert entries == [] - - -class TestPreprocessCalibrationImage: - """Tests for ModelConverter._preprocess_calibration_image.""" - - def test_valid_image(self, converter, tmp_path): - """_preprocess_calibration_image processes image correctly.""" - # Create a dummy image - img = np.zeros((100, 100, 3), dtype=np.uint8) - img_path = tmp_path / "test.jpg" - cv2.imwrite(str(img_path), img) - - result = converter._preprocess_calibration_image( - img_path=img_path, - width=224, - height=224, - mean=np.array([123.675, 116.28, 103.53]), - scale=np.array([58.395, 57.12, 57.375]), - reverse_input_channels=True, - ) +class TestPostprocessOpenvinoModel: + """Tests for PyTorchConverter._postprocess_openvino_model via TorchvisionConverter.""" - assert result is not None - assert result.shape == (1, 3, 224, 224) + def test_set_input_names(self, converter, mock_ov_model): + """_postprocess_openvino_model sets input tensor names.""" + converter._postprocess_openvino_model(mock_ov_model, input_names=["images"]) + mock_ov_model.input(0).set_names.assert_called_with({"images"}) - def test_no_channel_reversal(self, converter, tmp_path): - """_preprocess_calibration_image without channel reversal.""" - img = np.zeros((100, 100, 3), dtype=np.uint8) - img_path = tmp_path / "test.jpg" - cv2.imwrite(str(img_path), img) - - result = converter._preprocess_calibration_image( - img_path=img_path, - width=224, - height=224, - mean=np.array([0, 0, 0]), - scale=np.array([1, 1, 1]), - reverse_input_channels=False, - ) + def test_set_output_names(self, converter, mock_ov_model): + """_postprocess_openvino_model sets output tensor names.""" + converter._postprocess_openvino_model(mock_ov_model, output_names=["predictions"]) + mock_ov_model.output(0).set_names.assert_called_with({"predictions"}) - assert result is not None - assert result.shape == (1, 3, 224, 224) + def test_set_metadata(self, converter, mock_ov_model): + """_postprocess_openvino_model adds metadata.""" + metadata = {("model_info", "model_type"): "Classification", ("model_info", "labels"): "cat dog"} + converter._postprocess_openvino_model(mock_ov_model, metadata=metadata) + assert mock_ov_model.set_rt_info.call_count == 2 - def test_invalid_image(self, converter, tmp_path): - """_preprocess_calibration_image returns None for invalid image.""" - bad_path = tmp_path / "notanimage.txt" - bad_path.write_text("not an image") - - result = converter._preprocess_calibration_image( - img_path=bad_path, - width=224, - height=224, - mean=np.array([0, 0, 0]), - scale=np.array([1, 1, 1]), - reverse_input_channels=True, + def test_no_operations(self, converter, mock_ov_model): + """_postprocess_openvino_model handles None params.""" + assert converter._postprocess_openvino_model(mock_ov_model) is mock_ov_model + + +class TestBuildMetadata: + """Tests for PyTorchConverter._build_metadata via TorchvisionConverter.""" + + def test_builds_default_fields(self, converter, sample_model_config): + """_build_metadata includes shared Model API metadata.""" + metadata = converter._build_metadata({**sample_model_config, "labels": None}) + + assert metadata["model_info", "model_type"] == "Classification" + assert metadata["model_info", "model_short_name"] == "test_model" + assert metadata["model_info", "reverse_input_channels"] == "True" + assert metadata["model_info", "mean_values"] == "123.675 116.28 103.53" + assert metadata["model_info", "scale_values"] == "58.395 57.12 57.375" + + def test_includes_optional_fields(self, converter, sample_model_config): + """_build_metadata includes optional metadata fields when configured.""" + metadata = converter._build_metadata( + { + **sample_model_config, + "labels": None, + "confidence_threshold": 0.5, + "iou_threshold": 0.45, + "resize_type": "standard", + }, ) - assert result is None + assert metadata["model_info", "confidence_threshold"] == "0.5" + assert metadata["model_info", "iou_threshold"] == "0.45" + assert metadata["model_info", "resize_type"] == "standard" + def test_adds_resolved_labels(self, converter, sample_model_config): + """_build_metadata adds labels when they can be resolved.""" + with patch.object(converter, "get_labels", return_value="cat dog"): + metadata = converter._build_metadata(sample_model_config) -class TestCreateCalibrationDataset: - """Tests for ModelConverter.create_calibration_dataset.""" - - def test_no_dataset_path(self, tmp_path): - """create_calibration_dataset returns empty when no dataset path.""" - converter = ModelConverter( - output_dir=tmp_path / "out", - cache_dir=tmp_path / "cache", - dataset_path=None, - ) - result = converter.create_calibration_dataset(input_shape=[1, 3, 224, 224]) - assert result == [] + assert metadata["model_info", "labels"] == "cat dog" - def test_nonexistent_dataset_path(self, tmp_path): - """create_calibration_dataset returns empty for missing path.""" - converter = ModelConverter( - output_dir=tmp_path / "out", - cache_dir=tmp_path / "cache", - dataset_path=tmp_path / "nonexistent", - ) - result = converter.create_calibration_dataset(input_shape=[1, 3, 224, 224]) - assert result == [] + def test_skips_unknown_label_set(self, converter, sample_model_config): + """_build_metadata omits labels when lookup fails.""" + with patch.object(converter, "get_labels", return_value=None): + metadata = converter._build_metadata(sample_model_config) - def test_with_return_labels(self, tmp_path, dataset_dir): - """create_calibration_dataset returns images and labels.""" - converter = ModelConverter( - output_dir=tmp_path / "out", - cache_dir=tmp_path / "cache", - dataset_path=dataset_dir, - ) - result = converter.create_calibration_dataset( - input_shape=[1, 3, 224, 224], - return_labels=True, - ) - images, labels = result - assert len(images) == 2 - assert len(labels) == 2 - assert labels[0] == 0 - assert labels[1] == 1 - - def test_without_return_labels(self, tmp_path, dataset_dir): - """create_calibration_dataset returns images without labels flag.""" - converter = ModelConverter( - output_dir=tmp_path / "out", - cache_dir=tmp_path / "cache", - dataset_path=dataset_dir, - ) - result = converter.create_calibration_dataset( - input_shape=[1, 3, 224, 224], - return_labels=False, - ) - images, labels = result - assert len(images) == 2 - assert labels == [] + assert ("model_info", "labels") not in metadata - def test_with_mean_scale(self, tmp_path, dataset_dir): - """create_calibration_dataset uses mean and scale values.""" - converter = ModelConverter( - output_dir=tmp_path / "out", - cache_dir=tmp_path / "cache", - dataset_path=dataset_dir, - ) - result = converter.create_calibration_dataset( - input_shape=[1, 3, 224, 224], - mean_values="123.675 116.28 103.53", - scale_values="58.395 57.12 57.375", - return_labels=True, - ) - images, _labels = result - assert len(images) == 2 - def test_empty_dataset(self, tmp_path): - """create_calibration_dataset handles empty dataset directory.""" - empty_dataset = tmp_path / "empty_dataset" - empty_dataset.mkdir() - converter = ModelConverter( - output_dir=tmp_path / "out", - cache_dir=tmp_path / "cache", - dataset_path=empty_dataset, - ) - result = converter.create_calibration_dataset( - input_shape=[1, 3, 224, 224], - return_labels=True, - ) - assert result == ([], []) +class TestQuantizeAndCleanup: + """Tests for PyTorchConverter._quantize_and_cleanup via TorchvisionConverter.""" - def test_subset_size(self, tmp_path, dataset_dir): - """create_calibration_dataset respects subset_size.""" - converter = ModelConverter( - output_dir=tmp_path / "out", - cache_dir=tmp_path / "cache", - dataset_path=dataset_dir, - ) - result = converter.create_calibration_dataset( - input_shape=[1, 3, 224, 224], - subset_size=1, - return_labels=True, - ) - images, _labels = result - assert len(images) == 1 - - def test_image_processing_error(self, tmp_path): - """create_calibration_dataset skips images that raise exceptions (with labels).""" - dataset_path = tmp_path / "dataset" - class_dir = dataset_path / "0" - class_dir.mkdir(parents=True) - # Create a valid image file that will trigger an exception in preprocessing - img = np.zeros((10, 10, 3), dtype=np.uint8) - cv2.imwrite(str(class_dir / "image_001.jpg"), img) + def test_with_classification_labels(self, converter, tmp_path, sample_model_config): + """_quantize_and_cleanup validates classification models with labels.""" + fp32_path = tmp_path / "model_fp32.xml" + fp32_path.write_text("") + fp32_bin = tmp_path / "model_fp32.bin" + fp32_bin.write_text("weights") + validation_data = [np.zeros((1, 3, 224, 224))] + validation_labels = [0] - converter = ModelConverter( - output_dir=tmp_path / "out", - cache_dir=tmp_path / "cache", - dataset_path=dataset_path, - ) - # Mock _preprocess_calibration_image to raise an exception - with patch.object(converter, "_preprocess_calibration_image", side_effect=ValueError("bad image")): - result = converter.create_calibration_dataset( + with ( + patch.object(converter, "create_calibration_dataset", return_value=(validation_data, validation_labels)), + patch.object(converter, "quantize_model") as mock_quantize, + ): + converter._quantize_and_cleanup( + sample_model_config, + fp32_path, + model_type="Classification", input_shape=[1, 3, 224, 224], - return_labels=True, + mean_values="123.675 116.28 103.53", + scale_values="58.395 57.12 57.375", + reverse_input_channels=True, ) - images, _labels = result - assert len(images) == 0 - def test_image_processing_error_no_labels(self, tmp_path): - """create_calibration_dataset skips images that raise exceptions (without labels).""" - dataset_path = tmp_path / "dataset" - class_dir = dataset_path / "0" - class_dir.mkdir(parents=True) - img = np.zeros((10, 10, 3), dtype=np.uint8) - cv2.imwrite(str(class_dir / "image_001.jpg"), img) + call_kwargs = mock_quantize.call_args.kwargs + assert call_kwargs["validation_data"] == validation_data + assert call_kwargs["validation_labels"] == validation_labels + assert not fp32_path.exists() + assert not fp32_bin.exists() - converter = ModelConverter( - output_dir=tmp_path / "out", - cache_dir=tmp_path / "cache", - dataset_path=dataset_path, - ) - with patch.object(converter, "_preprocess_calibration_image", side_effect=OSError("read error")): - result = converter.create_calibration_dataset( - input_shape=[1, 3, 224, 224], - return_labels=False, - ) - images, _labels = result - assert len(images) == 0 + def test_measures_original_accuracy_with_torch_model(self, converter, tmp_path, sample_model_config): + """_quantize_and_cleanup measures the original PyTorch model accuracy when a model is given.""" + from model_converter.reporting import AccuracyResults - def test_image_returns_none_with_labels(self, tmp_path): - """create_calibration_dataset skips None images (with labels).""" - dataset_path = tmp_path / "dataset" - class_dir = dataset_path / "0" - class_dir.mkdir(parents=True) - img = np.zeros((10, 10, 3), dtype=np.uint8) - cv2.imwrite(str(class_dir / "image_001.jpg"), img) + fp32_path = tmp_path / "model_fp32.xml" + fp32_path.write_text("") + validation_data = [np.zeros((1, 3, 224, 224))] + validation_labels = [0] + torch_model = MagicMock() - converter = ModelConverter( - output_dir=tmp_path / "out", - cache_dir=tmp_path / "cache", - dataset_path=dataset_path, - ) - with patch.object(converter, "_preprocess_calibration_image", return_value=None): - result = converter.create_calibration_dataset( + with ( + patch.object(converter, "create_calibration_dataset", return_value=(validation_data, validation_labels)), + patch.object(converter, "validate_torch_model", return_value=0.91) as mock_validate, + patch.object(converter, "quantize_model") as mock_quantize, + ): + accuracy = converter._quantize_and_cleanup( + sample_model_config, + fp32_path, + model_type="Classification", input_shape=[1, 3, 224, 224], - return_labels=True, + mean_values="123.675 116.28 103.53", + scale_values="58.395 57.12 57.375", + reverse_input_channels=True, + torch_model=torch_model, ) - images, _labels = result - assert len(images) == 0 - def test_image_returns_none_without_labels(self, tmp_path): - """create_calibration_dataset skips None images (without labels).""" - dataset_path = tmp_path / "dataset" - class_dir = dataset_path / "0" - class_dir.mkdir(parents=True) - img = np.zeros((10, 10, 3), dtype=np.uint8) - cv2.imwrite(str(class_dir / "image_001.jpg"), img) + mock_validate.assert_called_once_with(torch_model, validation_data, validation_labels) + assert isinstance(accuracy, AccuracyResults) + assert accuracy.original_accuracy == pytest.approx(0.91) + assert accuracy.measured is True + # quantize_model still receives the validation collector. + assert mock_quantize.call_args.kwargs["accuracy_results"] is accuracy - converter = ModelConverter( - output_dir=tmp_path / "out", - cache_dir=tmp_path / "cache", - dataset_path=dataset_path, - ) - with patch.object(converter, "_preprocess_calibration_image", return_value=None): - result = converter.create_calibration_dataset( + def test_skips_original_accuracy_when_torch_validation_fails(self, converter, tmp_path, sample_model_config): + """A failed PyTorch validation leaves the original accuracy unset.""" + fp32_path = tmp_path / "model_fp32.xml" + fp32_path.write_text("") + validation_data = [np.zeros((1, 3, 224, 224))] + validation_labels = [0] + + with ( + patch.object(converter, "create_calibration_dataset", return_value=(validation_data, validation_labels)), + patch.object(converter, "validate_torch_model", return_value=None) as mock_validate, + patch.object(converter, "quantize_model"), + ): + accuracy = converter._quantize_and_cleanup( + sample_model_config, + fp32_path, + model_type="Classification", input_shape=[1, 3, 224, 224], - return_labels=False, + mean_values="123.675 116.28 103.53", + scale_values="58.395 57.12 57.375", + reverse_input_channels=True, + torch_model=MagicMock(), ) - images, _labels = result - assert len(images) == 0 - - def test_progress_logging_with_labels(self, tmp_path): - """create_calibration_dataset logs progress every 50 images (with labels).""" - dataset_path = tmp_path / "dataset" - class_dir = dataset_path / "0" - class_dir.mkdir(parents=True) - # Create 51 images to trigger the progress logging at i=49 (i+1=50) - img = np.zeros((10, 10, 3), dtype=np.uint8) - for i in range(51): - cv2.imwrite(str(class_dir / f"image_{i:03d}.jpg"), img) + mock_validate.assert_called_once() + assert accuracy.original_accuracy is None - converter = ModelConverter( - output_dir=tmp_path / "out", - cache_dir=tmp_path / "cache", - dataset_path=dataset_path, - verbose=True, - ) - result = converter.create_calibration_dataset( - input_shape=[1, 3, 10, 10], - return_labels=True, - ) - images, _labels = result - assert len(images) == 51 - - def test_progress_logging_without_labels(self, tmp_path): - """create_calibration_dataset logs progress every 50 images (without labels).""" - dataset_path = tmp_path / "dataset" - class_dir = dataset_path / "0" - class_dir.mkdir(parents=True) - - img = np.zeros((10, 10, 3), dtype=np.uint8) - for i in range(51): - cv2.imwrite(str(class_dir / f"image_{i:03d}.jpg"), img) - - converter = ModelConverter( - output_dir=tmp_path / "out", - cache_dir=tmp_path / "cache", - dataset_path=dataset_path, - verbose=True, - ) - result = converter.create_calibration_dataset( - input_shape=[1, 3, 10, 10], - return_labels=False, - ) - images, _labels = result - assert len(images) == 51 - - def test_dataset_dir_removed_after_init_check(self, tmp_path): - """create_calibration_dataset handles dir removed between checks.""" - - dataset_path = tmp_path / "dataset" - dataset_path.mkdir() - - converter = ModelConverter( - output_dir=tmp_path / "out", - cache_dir=tmp_path / "cache", - dataset_path=dataset_path, - ) - - # Remove the directory after converter init but before calibration runs - # Patch exists() to return True on first call (line 415) then False on second (line 428) - original_exists = Path.exists - call_count = [0] - - def mock_exists(self_path): - if self_path == dataset_path: - call_count[0] += 1 - return call_count[0] == 1 - return original_exists(self_path) - - with patch.object(Path, "exists", mock_exists): - result = converter.create_calibration_dataset( - input_shape=[1, 3, 224, 224], - return_labels=True, - ) - assert result == ([], []) - - -class TestValidateModel: - """Tests for ModelConverter.validate_model.""" - - def test_correct_predictions(self, converter, tmp_path): - """validate_model computes accuracy correctly.""" - mock_output_layer = MagicMock() - mock_compiled = MagicMock() - mock_compiled.outputs = [mock_output_layer] - - # Setup callable to return predictions matching labels - mock_compiled.return_value = {mock_output_layer: np.array([[0.1, 0.9, 0.0]])} - call_results = [ - {mock_output_layer: np.array([[0.1, 0.9, 0.0]])}, # pred = 1 - {mock_output_layer: np.array([[0.0, 0.0, 0.9]])}, # pred = 2 - ] - mock_compiled.side_effect = call_results - - mock_core = MagicMock() - mock_core.read_model.return_value = MagicMock() - mock_core.compile_model.return_value = mock_compiled - - with patch("openvino.Core", return_value=mock_core): - accuracy = converter.validate_model( - model_path=tmp_path / "model.xml", - validation_data=[np.zeros((1, 3, 224, 224)), np.zeros((1, 3, 224, 224))], - labels=[1, 2], - ) - - assert accuracy == pytest.approx(1.0) - - def test_partial_accuracy(self, converter, tmp_path): - """validate_model returns partial accuracy.""" - mock_output_layer = MagicMock() - mock_compiled = MagicMock() - mock_compiled.outputs = [mock_output_layer] - - call_results = [ - {mock_output_layer: np.array([[0.9, 0.1]])}, # pred = 0 (correct) - {mock_output_layer: np.array([[0.9, 0.1]])}, # pred = 0 (wrong, label=1) - ] - mock_compiled.side_effect = call_results - - mock_core = MagicMock() - mock_core.read_model.return_value = MagicMock() - mock_core.compile_model.return_value = mock_compiled - - with patch("openvino.Core", return_value=mock_core): - accuracy = converter.validate_model( - model_path=tmp_path / "model.xml", - validation_data=[np.zeros((1, 3, 224, 224)), np.zeros((1, 3, 224, 224))], - labels=[0, 1], - ) - - assert accuracy == pytest.approx(0.5) - - def test_validation_failure(self, converter, tmp_path): - """validate_model returns 0.0 on error.""" - with patch("openvino.Core", side_effect=RuntimeError("OV error")): - accuracy = converter.validate_model( - model_path=tmp_path / "model.xml", - validation_data=[np.zeros((1, 3, 224, 224))], - labels=[0], - ) - - assert accuracy == pytest.approx(0.0) - - -class TestQuantizeModel: - """Tests for ModelConverter.quantize_model.""" - - def test_no_calibration_data(self, converter, sample_model_config, tmp_path): - """quantize_model returns model_path when no calibration data.""" - model_path = tmp_path / "model.xml" - result = converter.quantize_model( - model_path=model_path, - calibration_data=[], - model_config=sample_model_config, - ) - assert result == model_path - - def test_successful_quantization(self, converter, sample_model_config, tmp_path): - """quantize_model performs INT8 quantization.""" - # Setup model path - fp16_dir = tmp_path / "test_model-fp16-ov" - fp16_dir.mkdir(parents=True) - model_path = fp16_dir / "test_model_fp32.xml" - model_path.write_text("") - - mock_ov_model = MagicMock() - mock_quantized = MagicMock() - mock_quantized.get_rt_info.return_value = MagicMock(value={"model_type": "Classification"}) - - mock_core = MagicMock() - mock_core.read_model.return_value = mock_ov_model - - calibration_data = [np.zeros((1, 3, 224, 224))] - - def consume_dataset(gen): - """Mock nncf.Dataset that consumes the generator.""" - list(gen) # Consume the generator to cover lines 572-573 - return MagicMock() - - with ( - patch("openvino.Core", return_value=mock_core), - patch("nncf.quantize", return_value=mock_quantized), - patch("nncf.Dataset", side_effect=consume_dataset), - patch("nncf.QuantizationPreset") as mock_preset, - patch("openvino.save_model"), - patch.object(Path, "exists", return_value=True), - patch("shutil.copy2"), - patch.object(converter, "copy_readme"), - ): - mock_preset.MIXED = "mixed" - mock_preset.PERFORMANCE = "performance" - result = converter.quantize_model( - model_path=model_path, - calibration_data=calibration_data, - model_config=sample_model_config, - preset="mixed", - ) - - assert result != model_path # Should return new quantized path - - def test_nncf_not_installed(self, converter, sample_model_config, tmp_path): - """quantize_model handles missing NNCF.""" - model_path = tmp_path / "model.xml" - calibration_data = [np.zeros((1, 3, 224, 224))] - - with ( - patch.dict("sys.modules", {"nncf": None}), - patch("builtins.__import__", side_effect=ImportError("No module named 'nncf'")), - ): - result = converter.quantize_model( - model_path=model_path, - calibration_data=calibration_data, - model_config=sample_model_config, - ) - - assert result == model_path - - def test_quantization_with_validation(self, converter, sample_model_config, tmp_path): - """quantize_model validates accuracy when validation data provided.""" - fp16_dir = tmp_path / "test_model-fp16-ov" - fp16_dir.mkdir(parents=True) - model_path = fp16_dir / "test_model_fp32.xml" - model_path.write_text("") - - mock_quantized = MagicMock() - mock_quantized.get_rt_info.return_value = MagicMock(value={"model_type": "Classification"}) - - mock_core = MagicMock() - mock_core.read_model.return_value = MagicMock() - - calibration_data = [np.zeros((1, 3, 224, 224))] - validation_data = [np.zeros((1, 3, 224, 224))] - validation_labels = [0] - - def consume_dataset(gen): - list(gen) - return MagicMock() - - with ( - patch("openvino.Core", return_value=mock_core), - patch("nncf.quantize", return_value=mock_quantized), - patch("nncf.Dataset", side_effect=consume_dataset), - patch("nncf.QuantizationPreset") as mock_preset, - patch("openvino.save_model"), - patch.object(Path, "exists", return_value=True), - patch("shutil.copy2"), - patch.object(converter, "copy_readme"), - patch.object(converter, "validate_model", return_value=0.95), - ): - mock_preset.MIXED = "mixed" - converter.quantize_model( - model_path=model_path, - calibration_data=calibration_data, - model_config=sample_model_config, - validation_data=validation_data, - validation_labels=validation_labels, - ) - - def test_quantization_runtime_error(self, converter, sample_model_config, tmp_path): - """quantize_model handles runtime errors gracefully.""" - model_path = tmp_path / "model.xml" - calibration_data = [np.zeros((1, 3, 224, 224))] - - with patch("openvino.Core", side_effect=RuntimeError("OV error")): - result = converter.quantize_model( - model_path=model_path, - calibration_data=calibration_data, - model_config=sample_model_config, - ) - - assert result == model_path - - -class TestExportToOpenvino: - """Tests for ModelConverter.export_to_openvino.""" - - def test_successful_export(self, converter, sample_model_config, tmp_path): - """export_to_openvino exports model to OV format.""" - mock_model = MagicMock(spec=nn.Module) - mock_model.eval.return_value = mock_model - - mock_ov_model = MagicMock() - mock_input = MagicMock() - mock_input.get_names.return_value = {"input"} - mock_ov_model.input.return_value = mock_input - mock_ov_model.inputs = [mock_input] - mock_ov_model.outputs = [MagicMock()] - mock_ov_model.get_rt_info.return_value = MagicMock(value={"model_type": "Classification"}) - - output_path = converter.output_dir / "test_model" - - with ( - patch("openvino.convert_model", return_value=mock_ov_model), - patch("openvino.save_model"), - patch.object(Path, "exists", return_value=True), - patch("shutil.copy2"), - patch.object(converter, "copy_readme"), - ): - fp16_path, _fp32_path = converter.export_to_openvino( - model=mock_model, - input_shape=[1, 3, 224, 224], - output_path=output_path, - model_config=sample_model_config, - input_names=["input"], - output_names=["result"], - metadata={("model_info", "model_type"): "Classification"}, - ) - - assert "fp16" in str(fp16_path.parent) or fp16_path.name == "test_model.xml" - - def test_export_failure(self, converter, sample_model_config, tmp_path): - """export_to_openvino raises on conversion failure.""" - mock_model = MagicMock(spec=nn.Module) - mock_model.eval.return_value = mock_model - - output_path = converter.output_dir / "test_model" - - with ( - patch("openvino.convert_model", side_effect=RuntimeError("Conversion failed")), - pytest.raises(RuntimeError, match="Conversion failed"), - ): - converter.export_to_openvino( - model=mock_model, - input_shape=[1, 3, 224, 224], - output_path=output_path, - model_config=sample_model_config, - ) - - -class TestPrepareModelForExport: - """Tests for ModelConverter._prepare_model_for_export.""" - - def test_with_adapter(self, converter): - """_prepare_model_for_export applies adapter for known model type.""" - mock_model = MagicMock(spec=nn.Module) - config = {"model_type": "MaskRCNN"} - - with patch("model_converter.cli.get_adapter") as mock_get_adapter: - mock_adapted = MagicMock() - mock_get_adapter.return_value = mock_adapted - result = converter._prepare_model_for_export(mock_model, config) - - assert result is mock_adapted - - def test_without_adapter(self, converter): - """_prepare_model_for_export returns model unchanged for no adapter.""" - mock_model = MagicMock(spec=nn.Module) - config = {"model_type": "Classification"} - - with patch("model_converter.cli.get_adapter", return_value=mock_model): - result = converter._prepare_model_for_export(mock_model, config) - - assert result is mock_model - - -class TestCreateExampleInput: - """Tests for ModelConverter._create_example_input.""" - - def test_maskrcnn_input(self, converter): - """_create_example_input uses rand for maskrcnn.""" - config = {"model_type": "MaskRCNN"} - result = converter._create_example_input([1, 3, 224, 224], config) - assert result.shape == (1, 3, 224, 224) - assert result.min() >= 0 # rand produces [0, 1) - - def test_default_input(self, converter): - """_create_example_input uses randn for non-maskrcnn.""" - config = {"model_type": "Classification"} - result = converter._create_example_input([1, 3, 224, 224], config) - assert result.shape == (1, 3, 224, 224) - - -class TestPostprocessOpenvinoModel: - """Tests for ModelConverter._postprocess_openvino_model.""" - - def test_set_input_names(self, converter, mock_ov_model): - """_postprocess_openvino_model sets input tensor names.""" - converter._postprocess_openvino_model( - mock_ov_model, - input_names=["images"], - ) - mock_ov_model.input(0).set_names.assert_called_with({"images"}) - - def test_set_output_names(self, converter, mock_ov_model): - """_postprocess_openvino_model sets output tensor names.""" - converter._postprocess_openvino_model( - mock_ov_model, - output_names=["predictions"], - ) - mock_ov_model.output(0).set_names.assert_called_with({"predictions"}) - - def test_set_metadata(self, converter, mock_ov_model): - """_postprocess_openvino_model adds metadata.""" - metadata = { - ("model_info", "model_type"): "Classification", - ("model_info", "labels"): "cat dog", - } - converter._postprocess_openvino_model(mock_ov_model, metadata=metadata) - assert mock_ov_model.set_rt_info.call_count == 2 - - def test_no_operations(self, converter, mock_ov_model): - """_postprocess_openvino_model handles None params.""" - result = converter._postprocess_openvino_model(mock_ov_model) - assert result is mock_ov_model - - -class TestLoadModelFromConfig: - """Tests for ModelConverter._load_model_from_config.""" - - def test_huggingface_path(self, converter): - """_load_model_from_config loads from HuggingFace.""" - config = { - "huggingface_repo": "timm/resnet50", - "huggingface_revision": "abc123", - "model_library": "timm", - } - mock_model = MagicMock() - with patch.object(converter, "load_huggingface_model", return_value=mock_model): - result = converter._load_model_from_config(config) - assert result is mock_model - - def test_huggingface_missing_revision(self, converter): - """_load_model_from_config raises when HF revision is missing.""" - config = { - "huggingface_repo": "timm/resnet50", - } - with pytest.raises(ValueError, match="huggingface_revision"): - converter._load_model_from_config(config) - - def test_url_path(self, converter, tmp_path): - """_load_model_from_config loads from URL.""" - config = { - "weights_url": "https://example.com/weights.pth", - "model_class_name": "torch.nn.Module", - } - - mock_model = MagicMock(spec=nn.Module) - mock_model.eval.return_value = mock_model - - with ( - patch.object(converter._url_downloader, "download", return_value=tmp_path / "weights.pth"), - patch.object(converter, "load_model_class", return_value=torch.nn.Module), - patch.object(converter, "load_checkpoint", return_value={"model": mock_model}), - patch.object(converter, "create_model", return_value=mock_model), - ): - result = converter._load_model_from_config(config) - - assert result is mock_model - - def test_url_path_default_class(self, converter, tmp_path): - """_load_model_from_config uses torch.nn.Module as default class.""" - config = { - "weights_url": "https://example.com/weights.pth", - } - - mock_model = MagicMock(spec=nn.Module) - - with ( - patch.object(converter._url_downloader, "download", return_value=tmp_path / "weights.pth"), - patch.object(converter, "load_model_class", return_value=torch.nn.Module) as mock_load_class, - patch.object(converter, "load_checkpoint", return_value={}), - patch.object(converter, "create_model", return_value=mock_model), - ): - converter._load_model_from_config(config) - - mock_load_class.assert_called_once_with("torch.nn.Module") - - -class TestQuantizeAndCleanup: - """Tests for ModelConverter._quantize_and_cleanup.""" - - def test_with_classification_labels(self, converter, tmp_path, sample_model_config): - """_quantize_and_cleanup runs validation for classification with labels.""" + def test_without_classification_labels(self, converter, tmp_path, sample_model_config): + """_quantize_and_cleanup skips validation for non-classification models.""" fp32_path = tmp_path / "model_fp32.xml" fp32_path.write_text("") - fp32_bin = tmp_path / "model_fp32.bin" - fp32_bin.write_text("weights") - - validation_data = [np.zeros((1, 3, 224, 224))] - validation_labels = [0] with ( - patch.object( - converter, - "create_calibration_dataset", - return_value=(validation_data, validation_labels), - ), - patch.object(converter, "quantize_model"), + patch.object(converter, "create_calibration_dataset", return_value=([np.zeros((1, 3, 224, 224))], [])), + patch.object(converter, "quantize_model") as mock_quantize, ): converter._quantize_and_cleanup( - sample_model_config, + {**sample_model_config, "labels": None}, fp32_path, - model_type="Classification", + model_type="Detection", input_shape=[1, 3, 224, 224], mean_values="123.675 116.28 103.53", scale_values="58.395 57.12 57.375", reverse_input_channels=True, ) - # FP32 files should be cleaned up - assert not fp32_path.exists() - assert not fp32_bin.exists() + call_kwargs = mock_quantize.call_args.kwargs + assert call_kwargs["validation_data"] is None + assert call_kwargs["validation_labels"] is None - def test_without_classification_labels(self, converter, tmp_path, sample_model_config): - """_quantize_and_cleanup skips validation for non-classification.""" + def test_empty_calibration_data(self, converter, tmp_path, sample_model_config): + """_quantize_and_cleanup skips quantization when no calibration data is produced.""" fp32_path = tmp_path / "model_fp32.xml" fp32_path.write_text("") - config = {**sample_model_config, "labels": None} - validation_data = [np.zeros((1, 3, 224, 224))] - with ( - patch.object(converter, "create_calibration_dataset", return_value=(validation_data, [])), + patch.object(converter, "create_calibration_dataset", return_value=([], [])), patch.object(converter, "quantize_model") as mock_quantize, ): converter._quantize_and_cleanup( - config, + sample_model_config, fp32_path, - model_type="Detection", + model_type="Classification", input_shape=[1, 3, 224, 224], mean_values="123.675 116.28 103.53", scale_values="58.395 57.12 57.375", reverse_input_channels=True, ) - # Quantize should be called with no validation data/labels - mock_quantize.assert_called_once() - call_kwargs = mock_quantize.call_args[1] - assert call_kwargs["validation_data"] is None - assert call_kwargs["validation_labels"] is None + mock_quantize.assert_not_called() + + def test_collects_validation_samples_for_non_top1_metric( + self, + converter, + tmp_path, + sample_model_config, + ): + """When the dispatched metric is non-Top1, raw image samples are forwarded to quantize_model.""" + from model_converter.datasets import CalibrationSample + from model_converter.metrics import MultilabelMAP - def test_empty_calibration_data(self, converter, tmp_path, sample_model_config): - """_quantize_and_cleanup skips quantization when no data.""" fp32_path = tmp_path / "model_fp32.xml" fp32_path.write_text("") + samples = [CalibrationSample(image_path=tmp_path / "x.jpg", label=0)] with ( - patch.object(converter, "create_calibration_dataset", return_value=([], [])), + patch.object( + converter, + "_metric_for_config", + return_value=MultilabelMAP(num_labels=3), + ), + patch.object( + converter, + "_collect_validation_samples", + return_value=samples, + ) as mock_collect, + patch.object( + converter, + "create_calibration_dataset", + return_value=([np.zeros((1, 3, 224, 224))], []), + ), patch.object(converter, "quantize_model") as mock_quantize, ): converter._quantize_and_cleanup( @@ -1266,10 +805,12 @@ def test_empty_calibration_data(self, converter, tmp_path, sample_model_config): reverse_input_channels=True, ) - mock_quantize.assert_not_called() + mock_collect.assert_called_once() + assert mock_quantize.call_args.kwargs["validation_samples"] == samples + assert isinstance(mock_quantize.call_args.kwargs["metric"], MultilabelMAP) def test_cleanup_failure(self, converter, tmp_path, sample_model_config): - """_quantize_and_cleanup handles cleanup failure gracefully.""" + """_quantize_and_cleanup swallows cleanup errors.""" fp32_path = tmp_path / "model_fp32.xml" fp32_path.write_text("") @@ -1278,7 +819,6 @@ def test_cleanup_failure(self, converter, tmp_path, sample_model_config): patch.object(Path, "exists", return_value=True), patch.object(Path, "unlink", side_effect=OSError("Permission denied")), ): - # Should not raise converter._quantize_and_cleanup( sample_model_config, fp32_path, @@ -1290,12 +830,11 @@ def test_cleanup_failure(self, converter, tmp_path, sample_model_config): ) -class TestProcessModelConfig: - """Tests for ModelConverter.process_model_config.""" +class TestTorchvisionProcessModelConfig: + """Tests for TorchvisionConverter.process_model_config.""" def test_already_exists(self, converter, sample_model_config): """process_model_config skips when both models already exist.""" - # Create existing model files fp16_dir = converter.output_dir / "test_model-fp16-ov" fp16_dir.mkdir(parents=True) (fp16_dir / "test_model.xml").write_text("") @@ -1304,343 +843,485 @@ def test_already_exists(self, converter, sample_model_config): int8_dir.mkdir(parents=True) (int8_dir / "test_model.xml").write_text("") - result = converter.process_model_config(sample_model_config) - assert result is True + assert converter.process_model_config(sample_model_config) is True def test_missing_license(self, converter): """process_model_config fails when license is missing.""" - config = {"model_short_name": "test", "license_link": "https://example.com"} - result = converter.process_model_config(config) - assert result is False + assert ( + converter.process_model_config({"model_short_name": "test", "license_link": "https://example.com"}) is False + ) def test_missing_license_link(self, converter): """process_model_config fails when license_link is missing.""" - config = {"model_short_name": "test", "license": "MIT"} - result = converter.process_model_config(config) - assert result is False + assert converter.process_model_config({"model_short_name": "test", "license": "MIT"}) is False - def test_successful_conversion(self, converter, sample_model_config): - """process_model_config successfully converts a model.""" + def test_successful_conversion(self, converter, sample_model_config, tmp_path): + """process_model_config runs the torchvision conversion workflow.""" mock_model = MagicMock(spec=nn.Module) mock_model.eval.return_value = mock_model - fp16_path = converter.output_dir / "test_model-fp16-ov" / "test_model.xml" fp32_path = converter.output_dir / "test_model-fp16-ov" / "test_model_fp32.xml" with ( - patch.object(converter, "_load_model_from_config", return_value=mock_model), - patch.object(converter, "get_labels", return_value="cat dog"), - patch.object(converter, "export_to_openvino", return_value=(fp16_path, fp32_path)), + patch.object(converter._url_downloader, "download", return_value=tmp_path / "weights.pth"), + patch.object(converter, "load_model_class", return_value=torch.nn.Module), + patch.object(converter, "load_checkpoint", return_value={"model": mock_model}), + patch.object(converter, "create_model", return_value=mock_model), + patch.object(converter, "_build_metadata", return_value={("model_info", "model_type"): "Classification"}), + patch.object(converter, "export_to_openvino", return_value=(fp16_path, fp32_path)) as mock_export, ): result = converter.process_model_config(sample_model_config) assert result is True + assert mock_export.call_args.kwargs["output_path"] == converter.output_dir / "test_model" def test_successful_conversion_with_dataset(self, tmp_path, sample_model_config, dataset_dir): - """process_model_config quantizes when dataset is available.""" - conv = ModelConverter( + """process_model_config quantizes when a dataset is available.""" + converter = TorchvisionConverter( output_dir=tmp_path / "out", cache_dir=tmp_path / "cache", - dataset_path=dataset_dir, + dataset_registry=dataset_dir, ) mock_model = MagicMock(spec=nn.Module) mock_model.eval.return_value = mock_model - - fp16_path = conv.output_dir / "test_model-fp16-ov" / "test_model.xml" - fp32_path = conv.output_dir / "test_model-fp16-ov" / "test_model_fp32.xml" + fp16_path = converter.output_dir / "test_model-fp16-ov" / "test_model.xml" + fp32_path = converter.output_dir / "test_model-fp16-ov" / "test_model_fp32.xml" with ( - patch.object(conv, "_load_model_from_config", return_value=mock_model), - patch.object(conv, "get_labels", return_value="cat dog"), - patch.object(conv, "export_to_openvino", return_value=(fp16_path, fp32_path)), - patch.object(conv, "_quantize_and_cleanup"), + patch.object(converter._url_downloader, "download", return_value=tmp_path / "weights.pth"), + patch.object(converter, "load_model_class", return_value=torch.nn.Module), + patch.object(converter, "load_checkpoint", return_value={"model": mock_model}), + patch.object(converter, "create_model", return_value=mock_model), + patch.object(converter, "_build_metadata", return_value={("model_info", "model_type"): "Classification"}), + patch.object(converter, "export_to_openvino", return_value=(fp16_path, fp32_path)), + patch.object(converter, "_quantize_and_cleanup", return_value=AccuracyResults()) as mock_quantize, ): - result = conv.process_model_config(sample_model_config) + result = converter.process_model_config(sample_model_config) assert result is True + mock_quantize.assert_called_once() def test_conversion_failure(self, converter, sample_model_config): """process_model_config returns False on failure.""" - with patch.object(converter, "_load_model_from_config", side_effect=RuntimeError("load failed")): - result = converter.process_model_config(sample_model_config) - - assert result is False + with patch.object(converter._url_downloader, "download", side_effect=RuntimeError("download failed")): + assert converter.process_model_config(sample_model_config) is False - def test_no_labels_configured(self, converter): - """process_model_config works without labels in config.""" - config = { - "model_short_name": "test_model", - "license": "Apache-2.0", - "license_link": "https://apache.org", - "weights_url": "https://example.com/weights.pth", - "model_class_name": "torch.nn.Module", - "input_shape": [1, 3, 224, 224], - "model_type": "Classification", - } - mock_model = MagicMock(spec=nn.Module) - fp16_path = converter.output_dir / "test_model-fp16-ov" / "test_model.xml" - fp32_path = converter.output_dir / "test_model-fp16-ov" / "test_model_fp32.xml" +class TestTimmProcessModelConfig: + """Tests for TimmConverter.process_model_config.""" - with ( - patch.object(converter, "_load_model_from_config", return_value=mock_model), - patch.object(converter, "export_to_openvino", return_value=(fp16_path, fp32_path)), - ): - result = converter.process_model_config(config) + def test_missing_huggingface_repo(self, timm_converter, sample_timm_config): + """process_model_config fails when huggingface_repo is missing.""" + assert ( + timm_converter.process_model_config({ + k: v for k, v in sample_timm_config.items() if k != "huggingface_repo" + }) + is False + ) - assert result is True + def test_missing_huggingface_revision(self, timm_converter, sample_timm_config): + """process_model_config fails when huggingface_revision is missing.""" + assert ( + timm_converter.process_model_config({ + k: v for k, v in sample_timm_config.items() if k != "huggingface_revision" + }) + is False + ) - def test_labels_not_found(self, converter, sample_model_config): - """process_model_config handles unknown label set.""" + def test_successful_conversion(self, timm_converter, sample_timm_config): + """process_model_config runs the timm conversion workflow.""" mock_model = MagicMock(spec=nn.Module) - fp16_path = converter.output_dir / "test_model-fp16-ov" / "test_model.xml" - fp32_path = converter.output_dir / "test_model-fp16-ov" / "test_model_fp32.xml" + mock_model.eval.return_value = mock_model + fp16_path = timm_converter.output_dir / "test_timm_model-fp16-ov" / "test_timm_model.xml" + fp32_path = timm_converter.output_dir / "test_timm_model-fp16-ov" / "test_timm_model_fp32.xml" with ( - patch.object(converter, "_load_model_from_config", return_value=mock_model), - patch.object(converter, "get_labels", return_value=None), - patch.object(converter, "export_to_openvino", return_value=(fp16_path, fp32_path)), + patch.object(timm_converter, "load_huggingface_model", return_value=mock_model), + patch.object( + timm_converter, + "_build_metadata", + return_value={("model_info", "model_type"): "Classification"}, + ), + patch.object(timm_converter, "export_to_openvino", return_value=(fp16_path, fp32_path)) as mock_export, ): - result = converter.process_model_config(sample_model_config) + result = timm_converter.process_model_config(sample_timm_config) assert result is True - - def test_metadata_fields(self, converter): - """process_model_config includes optional metadata fields.""" - config = { - "model_short_name": "test_model", - "license": "Apache-2.0", - "license_link": "https://apache.org", - "weights_url": "https://example.com/weights.pth", - "input_shape": [1, 3, 224, 224], - "model_type": "Detection", - "confidence_threshold": "0.5", - "iou_threshold": "0.45", - "resize_type": "standard", - } - - mock_model = MagicMock(spec=nn.Module) - fp16_path = converter.output_dir / "test_model-fp16-ov" / "test_model.xml" - fp32_path = converter.output_dir / "test_model-fp16-ov" / "test_model_fp32.xml" - - with ( - patch.object(converter, "_load_model_from_config", return_value=mock_model), - patch.object(converter, "export_to_openvino", return_value=(fp16_path, fp32_path)) as mock_export, - ): - converter.process_model_config(config) - - # Check metadata was passed - call_kwargs = mock_export.call_args[1] - metadata = call_kwargs["metadata"] - assert ("model_info", "confidence_threshold") in metadata - assert ("model_info", "iou_threshold") in metadata + assert mock_export.call_args.kwargs["output_path"] == timm_converter.output_dir / "test_timm_model" class TestMetadataValue: - """Tests for ModelConverter._metadata_value.""" - - def test_string(self): - """_metadata_value converts string to string.""" - assert ModelConverter._metadata_value("hello") == "hello" + """Tests for BaseConverter._metadata_value via TorchvisionConverter.""" - def test_integer(self): - """_metadata_value converts int to string.""" - assert ModelConverter._metadata_value(42) == "42" + @pytest.mark.parametrize( + ("value", "expected"), + [("hello", "hello"), (42, "42"), (0.5, "0.5"), (True, "True"), ([1, 2, 3], "1 2 3"), (("a", "b"), "a b")], + ) + def test_value_conversion(self, converter, value, expected): + """_metadata_value converts scalars and iterables to metadata strings.""" + assert converter._metadata_value(value) == expected - def test_float(self): - """_metadata_value converts float to string.""" - assert ModelConverter._metadata_value(0.5) == "0.5" - def test_boolean(self): - """_metadata_value converts bool to string.""" - assert ModelConverter._metadata_value(True) == "True" +class TestProcessConfigFile: + """Tests for facade ModelConverter.process_config_file.""" - def test_list(self): - """_metadata_value joins list with spaces.""" - assert ModelConverter._metadata_value([1, 2, 3]) == "1 2 3" + def test_collect_results_aggregates_converters(self, facade_converter): + """collect_results concatenates results from every instantiated converter.""" + from model_converter.reporting import ConversionResult - def test_tuple(self): - """_metadata_value joins tuple with spaces.""" - assert ModelConverter._metadata_value(("a", "b")) == "a b" + assert facade_converter.collect_results() == [] + converter = facade_converter._get_converter("torchvision") + result = ConversionResult( + model_short_name="m", + model_full_name="M", + model_type="Classification", + model_library="torchvision", + ) + converter.results.append(result) -class TestProcessConfigFile: - """Tests for ModelConverter.process_config_file.""" + assert facade_converter.collect_results() == [result] - def test_multiple_models(self, converter, tmp_path): + def test_multiple_models(self, facade_converter, tmp_path): """process_config_file processes multiple models.""" - config = { - "models": [ - { - "model_short_name": "model1", - "license": "MIT", - "license_link": "https://mit.edu", - "weights_url": "https://example.com/1.pth", - }, + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( { - "model_short_name": "model2", - "license": "MIT", - "license_link": "https://mit.edu", - "weights_url": "https://example.com/2.pth", + "models": [ + {"model_short_name": "model1", "license": "MIT", "license_link": "https://mit.edu"}, + {"model_short_name": "model2", "license": "MIT", "license_link": "https://mit.edu"}, + ], }, - ], - } - config_path = tmp_path / "config.json" - config_path.write_text(json.dumps(config)) + ), + ) - with patch.object(converter, "process_model_config", return_value=True) as mock_process: - successful, failed = converter.process_config_file(config_path) + with patch.object(facade_converter, "process_model_config", return_value=True) as mock_process: + successful, failed = facade_converter.process_config_file(config_path) assert successful == 2 assert failed == 0 assert mock_process.call_count == 2 - def test_filter_match(self, converter, tmp_path): - """process_config_file filters to specific model.""" - config = { - "models": [ - {"model_short_name": "model1", "license": "MIT", "license_link": "x"}, - {"model_short_name": "model2", "license": "MIT", "license_link": "x"}, - ], - } + def test_filter_match(self, facade_converter, tmp_path): + """process_config_file filters to a specific model.""" config_path = tmp_path / "config.json" - config_path.write_text(json.dumps(config)) + config_path.write_text( + json.dumps( + { + "models": [ + {"model_short_name": "model1", "license": "MIT", "license_link": "x"}, + {"model_short_name": "model2", "license": "MIT", "license_link": "x"}, + ], + }, + ), + ) - with patch.object(converter, "process_model_config", return_value=True) as mock_process: - successful, _failed = converter.process_config_file(config_path, model_filter="model2") + with patch.object(facade_converter, "process_model_config", return_value=True) as mock_process: + successful, _failed = facade_converter.process_config_file(config_path, model_filter="model2") assert successful == 1 mock_process.assert_called_once() - def test_filter_no_match(self, converter, tmp_path): - """process_config_file returns 0,0 when filter doesn't match.""" - config = { - "models": [ - {"model_short_name": "model1"}, - ], - } + def test_filter_no_match(self, facade_converter, tmp_path): + """process_config_file returns 0,0 when filter does not match.""" + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"models": [{"model_short_name": "model1"}]})) + + successful, failed = facade_converter.process_config_file(config_path, model_filter="nonexistent") + assert successful == 0 + assert failed == 0 + + def test_empty_models(self, facade_converter, tmp_path): + """process_config_file handles empty models list.""" config_path = tmp_path / "config.json" - config_path.write_text(json.dumps(config)) + config_path.write_text(json.dumps({"models": []})) - successful, failed = converter.process_config_file(config_path, model_filter="nonexistent") + successful, failed = facade_converter.process_config_file(config_path) assert successful == 0 assert failed == 0 - def test_empty_models(self, converter, tmp_path): - """process_config_file handles empty models list.""" - config = {"models": []} - config_path = tmp_path / "config.json" - config_path.write_text(json.dumps(config)) + def test_invalid_json(self, facade_converter, tmp_path): + """process_config_file raises on invalid JSON.""" + config_path = tmp_path / "bad.json" + config_path.write_text("not valid json {{{") + + with pytest.raises(json.JSONDecodeError): + facade_converter.process_config_file(config_path) + + def test_model_failure(self, facade_converter, tmp_path): + """process_config_file counts failed models.""" + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"models": [{"model_short_name": "model1"}]})) + + with patch.object(facade_converter, "process_model_config", return_value=False): + successful, failed = facade_converter.process_config_file(config_path) + + assert successful == 0 + assert failed == 1 + + +class TestListModels: + """Tests for list_models function.""" + + def test_normal_output(self, tmp_path, capsys): + """list_models prints model information.""" + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "models": [ + { + "model_short_name": "resnet50", + "model_full_name": "ResNet-50", + "model_library": "torchvision", + "model_type": "Classification", + }, + ], + }, + ), + ) + + list_models(config_path) + + captured = capsys.readouterr() + assert "resnet50" in captured.out + assert "ResNet-50" in captured.out + assert "torchvision" in captured.out + assert "Classification" in captured.out + + def test_file_not_found(self, tmp_path, capsys): + """list_models handles missing config file.""" + list_models(tmp_path / "nonexistent.json") + assert "Error" in capsys.readouterr().err + + def test_empty_models(self, tmp_path, capsys): + """list_models handles empty models list.""" + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"models": []})) + + list_models(config_path) + assert "No models found" in capsys.readouterr().out + + def test_invalid_json(self, tmp_path, capsys): + """list_models handles invalid JSON.""" + config_path = tmp_path / "bad.json" + config_path.write_text("not json") + + list_models(config_path) + assert "Error" in capsys.readouterr().err + + +class TestMain: + """Tests for main() CLI entry point.""" + + def test_missing_config_file(self, tmp_path, monkeypatch): + """main returns 1 when config file does not exist.""" + monkeypatch.setattr(sys, "argv", ["model_converter", str(tmp_path / "nonexistent.json")]) + assert main() == 1 + + def test_list_flag(self, tmp_path, monkeypatch, capsys): + """main --list lists models and exits.""" + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"models": [{"model_short_name": "test", "model_type": "cls"}]})) + + monkeypatch.setattr(sys, "argv", ["model_converter", str(config_path), "--list"]) + assert main() == 0 + assert "test" in capsys.readouterr().out + + def test_successful_run(self, tmp_path, monkeypatch): + """main runs conversion successfully.""" + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"models": [{"model_short_name": "m1"}]})) + + monkeypatch.setattr( + sys, + "argv", + [ + "model_converter", + str(config_path), + "-o", + str(tmp_path / "output"), + "-c", + str(tmp_path / "cache"), + ], + ) + + with patch.object(ModelConverter, "process_config_file", return_value=(1, 0)): + assert main() == 0 + + def test_datasets_config_not_found_warns(self, tmp_path, monkeypatch, caplog): + """main warns and continues when --datasets-config path does not exist.""" + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"models": [{"model_short_name": "m1"}]})) + missing_datasets = tmp_path / "missing_datasets.json" + + monkeypatch.setattr( + sys, + "argv", + [ + "model_converter", + str(config_path), + "-o", + str(tmp_path / "output"), + "-c", + str(tmp_path / "cache"), + "--datasets-config", + str(missing_datasets), + ], + ) + + with patch.object(ModelConverter, "process_config_file", return_value=(1, 0)): + assert main() == 0 + assert "Dataset configuration not found" in caplog.text - successful, failed = converter.process_config_file(config_path) - assert successful == 0 - assert failed == 0 + def test_datasets_config_invalid_warns(self, tmp_path, monkeypatch, caplog): + """main warns and continues when --datasets-config contains invalid content.""" + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"models": [{"model_short_name": "m1"}]})) + bad_datasets = tmp_path / "bad_datasets.json" + bad_datasets.write_text('{"not_datasets": {}}') - def test_invalid_json(self, converter, tmp_path): - """process_config_file raises on invalid JSON.""" - config_path = tmp_path / "bad.json" - config_path.write_text("not valid json {{{") + monkeypatch.setattr( + sys, + "argv", + [ + "model_converter", + str(config_path), + "-o", + str(tmp_path / "output"), + "-c", + str(tmp_path / "cache"), + "--datasets-config", + str(bad_datasets), + ], + ) - with pytest.raises(json.JSONDecodeError): - converter.process_config_file(config_path) + with patch.object(ModelConverter, "process_config_file", return_value=(1, 0)): + assert main() == 0 + assert "Failed to load dataset registry" in caplog.text - def test_model_failure(self, converter, tmp_path): - """process_config_file counts failed models.""" - config = { - "models": [ - {"model_short_name": "model1", "license": "MIT", "license_link": "x"}, - ], - } + def test_report_default_path(self, tmp_path, monkeypatch, capsys): + """main prints console table even without --report (reporting is on by default).""" config_path = tmp_path / "config.json" - config_path.write_text(json.dumps(config)) + config_path.write_text(json.dumps({"models": [{"model_short_name": "m1"}]})) - with patch.object(converter, "process_model_config", return_value=False): - successful, failed = converter.process_config_file(config_path) + monkeypatch.setattr( + sys, + "argv", + ["model_converter", str(config_path), "-o", str(tmp_path / "output"), "-c", str(tmp_path / "cache")], + ) - assert successful == 0 - assert failed == 1 + from model_converter.reporting import STATUS_OK, ConversionResult + result = ConversionResult( + model_short_name="m1", + model_full_name="M1", + model_type="Classification", + model_library="torchvision", + status=STATUS_OK, + ) -class TestListModels: - """Tests for list_models function.""" + with ( + patch.object(ModelConverter, "process_config_file", return_value=(1, 0)), + patch.object(ModelConverter, "collect_results", return_value=[result]), + ): + assert main() == 0 - def test_normal_output(self, tmp_path, capsys): - """list_models prints model information.""" - config = { - "models": [ - { - "model_short_name": "resnet50", - "model_full_name": "ResNet-50", - "model_type": "Classification", - }, - ], - } - config_path = tmp_path / "config.json" - config_path.write_text(json.dumps(config)) + assert "Conversion Summary Report" in capsys.readouterr().out - list_models(config_path) + def test_no_report_flag_disables_reporting(self, tmp_path, monkeypatch, capsys): + """main --no-report suppresses the console table output.""" + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"models": [{"model_short_name": "m1"}]})) - captured = capsys.readouterr() - assert "resnet50" in captured.out - assert "ResNet-50" in captured.out - assert "Classification" in captured.out + monkeypatch.setattr( + sys, + "argv", + [ + "model_converter", + str(config_path), + "-o", + str(tmp_path / "output"), + "-c", + str(tmp_path / "cache"), + "--no-report", + ], + ) - def test_file_not_found(self, tmp_path, capsys): - """list_models handles missing config file.""" - list_models(tmp_path / "nonexistent.json") + with patch.object(ModelConverter, "process_config_file", return_value=(1, 0)): + assert main() == 0 - captured = capsys.readouterr() - assert "Error" in captured.err + assert "Conversion Summary Report" not in capsys.readouterr().out - def test_empty_models(self, tmp_path, capsys): - """list_models handles empty models list.""" - config = {"models": []} + def test_report_custom_path(self, tmp_path, monkeypatch, capsys): + """main --report PATH passes the custom report path to ModelConverter.""" config_path = tmp_path / "config.json" - config_path.write_text(json.dumps(config)) + config_path.write_text(json.dumps({"models": [{"model_short_name": "m1"}]})) + report_path = tmp_path / "custom" / "my_report.md" - list_models(config_path) + monkeypatch.setattr( + sys, + "argv", + [ + "model_converter", + str(config_path), + "-o", + str(tmp_path / "output"), + "-c", + str(tmp_path / "cache"), + "--report", + str(report_path), + ], + ) - captured = capsys.readouterr() - assert "No models found" in captured.out + captured: dict = {} + original_init = ModelConverter.__init__ - def test_invalid_json(self, tmp_path, capsys): - """list_models handles invalid JSON.""" - config_path = tmp_path / "bad.json" - config_path.write_text("not json") + def spy_init(self, *args, **kwargs): + captured.update(kwargs) + return original_init(self, *args, **kwargs) - list_models(config_path) + with ( + patch.object(ModelConverter, "__init__", autospec=True, side_effect=spy_init), + patch.object(ModelConverter, "process_config_file", return_value=(0, 0)), + patch.object(ModelConverter, "collect_results", return_value=[]), + ): + assert main() == 0 - captured = capsys.readouterr() - assert "Error" in captured.err + assert captured.get("report_path") == report_path + def test_measure_accuracy_default_true(self, tmp_path, monkeypatch): + """main passes measure_accuracy=True to ModelConverter by default.""" + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"models": []})) -class TestMain: - """Tests for main() CLI entry point.""" + monkeypatch.setattr( + sys, + "argv", + ["model_converter", str(config_path), "-o", str(tmp_path / "output"), "-c", str(tmp_path / "cache")], + ) - def test_missing_config_file(self, tmp_path, monkeypatch): - """main returns 1 when config file doesn't exist.""" - monkeypatch.setattr(sys, "argv", ["model_converter", str(tmp_path / "nonexistent.json")]) - result = main() - assert result == 1 + captured: dict = {} + original_init = ModelConverter.__init__ - def test_list_flag(self, tmp_path, monkeypatch, capsys): - """main --list flag lists models and exits.""" - config = {"models": [{"model_short_name": "test", "model_type": "cls"}]} - config_path = tmp_path / "config.json" - config_path.write_text(json.dumps(config)) + def spy_init(self, *args, **kwargs): + captured.update(kwargs) + return original_init(self, *args, **kwargs) - monkeypatch.setattr(sys, "argv", ["model_converter", str(config_path), "--list"]) - result = main() - assert result == 0 + with ( + patch.object(ModelConverter, "__init__", autospec=True, side_effect=spy_init), + patch.object(ModelConverter, "process_config_file", return_value=(0, 0)), + patch.object(ModelConverter, "collect_results", return_value=[]), + ): + assert main() == 0 - captured = capsys.readouterr() - assert "test" in captured.out + assert captured.get("measure_accuracy") is True - def test_successful_run(self, tmp_path, monkeypatch): - """main runs conversion successfully.""" - config = {"models": [{"model_short_name": "m1", "license": "MIT", "license_link": "x"}]} + def test_no_measure_accuracy_flag(self, tmp_path, monkeypatch): + """main --no-measure-accuracy passes measure_accuracy=False.""" config_path = tmp_path / "config.json" - config_path.write_text(json.dumps(config)) + config_path.write_text(json.dumps({"models": []})) monkeypatch.setattr( sys, @@ -1652,21 +1333,30 @@ def test_successful_run(self, tmp_path, monkeypatch): str(tmp_path / "output"), "-c", str(tmp_path / "cache"), - "-d", - str(tmp_path / "dataset"), + "--no-measure-accuracy", ], ) - with patch.object(ModelConverter, "process_config_file", return_value=(1, 0)): - result = main() + captured: dict = {} + original_init = ModelConverter.__init__ - assert result == 0 + def spy_init(self, *args, **kwargs): + captured.update(kwargs) + return original_init(self, *args, **kwargs) + + with ( + patch.object(ModelConverter, "__init__", autospec=True, side_effect=spy_init), + patch.object(ModelConverter, "process_config_file", return_value=(0, 0)), + patch.object(ModelConverter, "collect_results", return_value=[]), + ): + assert main() == 0 + + assert captured.get("measure_accuracy") is False def test_failed_run(self, tmp_path, monkeypatch): """main returns 1 when models fail.""" - config = {"models": [{"model_short_name": "m1"}]} config_path = tmp_path / "config.json" - config_path.write_text(json.dumps(config)) + config_path.write_text(json.dumps({"models": [{"model_short_name": "m1"}]})) monkeypatch.setattr( sys, @@ -1682,15 +1372,12 @@ def test_failed_run(self, tmp_path, monkeypatch): ) with patch.object(ModelConverter, "process_config_file", return_value=(0, 1)): - result = main() - - assert result == 1 + assert main() == 1 def test_verbose_flag(self, tmp_path, monkeypatch): - """main enables verbose logging with -v flag.""" - config = {"models": []} + """main enables verbose logging with -v.""" config_path = tmp_path / "config.json" - config_path.write_text(json.dumps(config)) + config_path.write_text(json.dumps({"models": []})) monkeypatch.setattr( sys, @@ -1707,15 +1394,12 @@ def test_verbose_flag(self, tmp_path, monkeypatch): ) with patch.object(ModelConverter, "process_config_file", return_value=(0, 0)): - result = main() - - assert result == 0 + assert main() == 0 def test_model_filter(self, tmp_path, monkeypatch): - """main passes --model filter to process_config_file.""" - config = {"models": [{"model_short_name": "target"}]} + """main passes --model filter through to process_config_file.""" config_path = tmp_path / "config.json" - config_path.write_text(json.dumps(config)) + config_path.write_text(json.dumps({"models": [{"model_short_name": "target"}]})) monkeypatch.setattr( sys, @@ -1735,13 +1419,12 @@ def test_model_filter(self, tmp_path, monkeypatch): with patch.object(ModelConverter, "process_config_file", return_value=(1, 0)) as mock_process: main() - mock_process.assert_called_once_with(config_path=config_path, model_filter="target") + mock_process.assert_called_once_with(config_path=config_path, model_filter="target", library_filter=None) def test_exception_during_processing(self, tmp_path, monkeypatch): - """main returns 1 on unhandled exception.""" - config = {"models": [{"model_short_name": "m1"}]} + """main returns 1 on processing exceptions.""" config_path = tmp_path / "config.json" - config_path.write_text(json.dumps(config)) + config_path.write_text(json.dumps({"models": [{"model_short_name": "m1"}]})) monkeypatch.setattr( sys, @@ -1757,6 +1440,293 @@ def test_exception_during_processing(self, tmp_path, monkeypatch): ) with patch.object(ModelConverter, "process_config_file", side_effect=ValueError("bad config")): - result = main() + assert main() == 1 + + +class TestLibraryFilter: + """Tests for --library filter functionality.""" + + def test_single_library_filter(self, facade_converter, tmp_path): + """process_config_file filters models by a single library.""" + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "models": [ + {"model_short_name": "m1", "model_library": "torchvision"}, + {"model_short_name": "m2", "model_library": "getitune"}, + {"model_short_name": "m3", "model_library": "timm"}, + ], + }, + ), + ) + + with patch.object(facade_converter, "process_model_config", return_value=True) as mock_process: + successful, failed = facade_converter.process_config_file( + config_path, + library_filter=["getitune"], + ) - assert result == 1 + assert successful == 1 + assert failed == 0 + mock_process.assert_called_once_with({"model_short_name": "m2", "model_library": "getitune"}) + + def test_multiple_library_filter(self, facade_converter, tmp_path): + """process_config_file filters models by multiple libraries.""" + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "models": [ + {"model_short_name": "m1", "model_library": "torchvision"}, + {"model_short_name": "m2", "model_library": "getitune"}, + {"model_short_name": "m3", "model_library": "timm"}, + {"model_short_name": "m4", "model_library": "yolo"}, + ], + }, + ), + ) + + with patch.object(facade_converter, "process_model_config", return_value=True) as mock_process: + successful, failed = facade_converter.process_config_file( + config_path, + library_filter=["getitune", "timm"], + ) + + assert successful == 2 + assert failed == 0 + assert mock_process.call_count == 2 + + def test_library_filter_no_match(self, facade_converter, tmp_path): + """process_config_file returns 0,0 when no models match library filter.""" + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "models": [ + {"model_short_name": "m1", "model_library": "torchvision"}, + ], + }, + ), + ) + + successful, failed = facade_converter.process_config_file(config_path, library_filter=["yolo"]) + assert successful == 0 + assert failed == 0 + + def test_library_and_model_filter_combined(self, facade_converter, tmp_path): + """process_config_file applies both library and model filters (AND logic).""" + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "models": [ + {"model_short_name": "m1", "model_library": "getitune"}, + {"model_short_name": "m2", "model_library": "getitune"}, + {"model_short_name": "m1", "model_library": "torchvision"}, + ], + }, + ), + ) + + with patch.object(facade_converter, "process_model_config", return_value=True) as mock_process: + successful, failed = facade_converter.process_config_file( + config_path, + model_filter="m1", + library_filter=["getitune"], + ) + + assert successful == 1 + assert failed == 0 + mock_process.assert_called_once_with({"model_short_name": "m1", "model_library": "getitune"}) + + def test_no_library_filter_processes_all(self, facade_converter, tmp_path): + """process_config_file processes all models when no library filter is given.""" + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "models": [ + {"model_short_name": "m1", "model_library": "torchvision"}, + {"model_short_name": "m2", "model_library": "getitune"}, + ], + }, + ), + ) + + with patch.object(facade_converter, "process_model_config", return_value=True) as mock_process: + successful, failed = facade_converter.process_config_file(config_path, library_filter=None) + + assert successful == 2 + assert failed == 0 + assert mock_process.call_count == 2 + + def test_list_models_with_library_filter(self, tmp_path, capsys): + """list_models filters output by library.""" + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "models": [ + { + "model_short_name": "resnet50", + "model_full_name": "ResNet-50", + "model_library": "torchvision", + "model_type": "Classification", + }, + { + "model_short_name": "dino_v2", + "model_full_name": "DINOv2", + "model_library": "getitune", + "model_type": "Classification", + }, + ], + }, + ), + ) + + list_models(config_path, library_filter=["getitune"]) + + captured = capsys.readouterr() + assert "dino_v2" in captured.out + assert "resnet50" not in captured.out + + def test_list_models_no_library_filter(self, tmp_path, capsys): + """list_models shows all models when no library filter.""" + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "models": [ + { + "model_short_name": "resnet50", + "model_full_name": "ResNet-50", + "model_library": "torchvision", + "model_type": "Classification", + }, + { + "model_short_name": "dino_v2", + "model_full_name": "DINOv2", + "model_library": "getitune", + "model_type": "Classification", + }, + ], + }, + ), + ) + + list_models(config_path, library_filter=None) + + captured = capsys.readouterr() + assert "dino_v2" in captured.out + assert "resnet50" in captured.out + + def test_unknown_library_warns(self, facade_converter, tmp_path, caplog): + """process_config_file warns on unknown library names.""" + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "models": [ + {"model_short_name": "m1", "model_library": "torchvision"}, + ], + }, + ), + ) + + with patch.object(facade_converter, "process_model_config", return_value=True): + facade_converter.process_config_file(config_path, library_filter=["torchvision", "nonexistent"]) + + assert "Unknown library 'nonexistent'" in caplog.text + + def test_list_models_unknown_library_warns(self, tmp_path, capsys): + """list_models prints warning for unknown library names.""" + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"models": [{"model_short_name": "m1", "model_library": "torchvision"}]})) + + list_models(config_path, library_filter=["nonexistent"]) + + captured = capsys.readouterr() + assert "Unknown library 'nonexistent'" in captured.err + assert "No models found" in captured.out + + def test_main_library_flag(self, tmp_path, monkeypatch): + """main passes --library filter through to process_config_file.""" + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"models": [{"model_short_name": "m1", "model_library": "getitune"}]})) + + monkeypatch.setattr( + sys, + "argv", + [ + "model_converter", + str(config_path), + "-o", + str(tmp_path / "output"), + "-c", + str(tmp_path / "cache"), + "--library", + "getitune,timm", + ], + ) + + with patch.object(ModelConverter, "process_config_file", return_value=(1, 0)) as mock_process: + main() + + mock_process.assert_called_once_with( + config_path=config_path, + model_filter=None, + library_filter=["getitune", "timm"], + ) + + def test_main_library_and_model_flags(self, tmp_path, monkeypatch): + """main passes both --library and --model filters through.""" + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"models": [{"model_short_name": "m1", "model_library": "getitune"}]})) + + monkeypatch.setattr( + sys, + "argv", + [ + "model_converter", + str(config_path), + "-o", + str(tmp_path / "output"), + "-c", + str(tmp_path / "cache"), + "--library", + "getitune", + "--model", + "m1", + ], + ) + + with patch.object(ModelConverter, "process_config_file", return_value=(1, 0)) as mock_process: + main() + + mock_process.assert_called_once_with( + config_path=config_path, + model_filter="m1", + library_filter=["getitune"], + ) + + def test_main_list_with_library_filter(self, tmp_path, monkeypatch, capsys): + """main --list with --library filters listed models.""" + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "models": [ + {"model_short_name": "m1", "model_library": "torchvision", "model_type": "cls"}, + {"model_short_name": "m2", "model_library": "getitune", "model_type": "cls"}, + ], + }, + ), + ) + + monkeypatch.setattr(sys, "argv", ["model_converter", str(config_path), "--list", "--library", "getitune"]) + assert main() == 0 + + captured = capsys.readouterr() + assert "m2" in captured.out + assert "m1" not in captured.out diff --git a/model_converter/tests/unit/test_converters.py b/model_converter/tests/unit/test_converters.py new file mode 100644 index 000000000..8ea5ea829 --- /dev/null +++ b/model_converter/tests/unit/test_converters.py @@ -0,0 +1,1744 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Tests for shared converter helpers via TorchvisionConverter.""" + +import json +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import MagicMock, patch + +import cv2 +import numpy as np +import pytest +from model_converter.converters.torchvision import TorchvisionConverter + + +def _set_template_root(monkeypatch: pytest.MonkeyPatch, template_dir: Path) -> None: + """Redirect BaseConverter template lookup to a test template directory.""" + import model_converter.converters.base as base_module + + monkeypatch.setattr(base_module, "__file__", str(template_dir.parent / "converters" / "base.py")) + + +class TestCopyReadme: + """Tests for BaseConverter.copy_readme via TorchvisionConverter.""" + + def test_renders_placeholders_and_tags_yaml(self, converter, template_dir, tmp_path, monkeypatch): + """copy_readme fills placeholders from config and renders tags as YAML.""" + _set_template_root(monkeypatch, template_dir) + (template_dir / "README-torchvision-fp16.md").write_text( + "# <>\n" + "License: <>\n" + "License link: <>\n" + "Docs: <>\n" + "Variant: <>\n" + "Tags:\n" + "<>", + ) + output_folder = tmp_path / "test_model-fp16-ov" + output_folder.mkdir() + + converter.copy_readme( + { + "model_short_name": "test_model", + "model_library": "torchvision", + "license": "Apache-2.0", + "license_link": "https://apache.org/licenses/LICENSE-2.0", + "docs": "https://docs.example.com", + "tags": ["vision", "classification"], + }, + output_folder, + variant="fp16", + ) + + content = (output_folder / "README.md").read_text() + assert "# test_model" in content + assert "License: Apache-2.0" in content + assert "License link: https://apache.org/licenses/LICENSE-2.0" in content + assert "Docs: https://docs.example.com" in content + assert "Variant: fp16" in content + assert " - vision" in content + assert " - classification" in content + + @pytest.mark.parametrize( + ("config", "expected"), + [ + ( + {"model_short_name": "", "license": "MIT", "license_link": "https://mit.edu"}, + "non-empty model_short_name", + ), + ({"model_short_name": "test_model", "license": "MIT", "license_link": ""}, "non-empty license_link"), + ({"model_short_name": "test_model", "license": "", "license_link": "https://mit.edu"}, "non-empty license"), + ], + ) + def test_logs_validation_failures(self, converter, tmp_path, caplog, config, expected): + """copy_readme logs warnings when required metadata is missing.""" + output_folder = tmp_path / "test_model-fp16-ov" + output_folder.mkdir() + + converter.copy_readme(config, output_folder) + + assert expected in caplog.text + assert not (output_folder / "README.md").exists() + + def test_missing_docs_logs_warning_and_still_writes_readme( + self, + converter, + template_dir, + tmp_path, + monkeypatch, + caplog, + ): + """copy_readme leaves the docs placeholder empty when docs are omitted.""" + _set_template_root(monkeypatch, template_dir) + (template_dir / "README-torchvision-fp16.md").write_text("Docs: <>") + output_folder = tmp_path / "test_model-fp16-ov" + output_folder.mkdir() + + converter.copy_readme( + { + "model_short_name": "test_model", + "model_library": "torchvision", + "license": "Apache-2.0", + "license_link": "https://apache.org/licenses/LICENSE-2.0", + }, + output_folder, + variant="fp16", + ) + + assert "does not define 'docs' field" in caplog.text + assert (output_folder / "README.md").read_text() == "Docs: " + + def test_logs_warning_when_template_not_found(self, converter, tmp_path, monkeypatch, caplog): + """copy_readme logs a warning and does nothing when the template file doesn't exist.""" + import model_converter.converters.base as base_module + + empty_templates = tmp_path / "empty_templates" + empty_templates.mkdir() + monkeypatch.setattr(base_module, "__file__", str(empty_templates.parent / "converters" / "base.py")) + + output_folder = tmp_path / "test_model-fp16-ov" + output_folder.mkdir() + + converter.copy_readme( + { + "model_short_name": "test_model", + "model_library": "nonexistent_library", + "license": "Apache-2.0", + "license_link": "https://apache.org/licenses/LICENSE-2.0", + }, + output_folder, + variant="fp16", + ) + + assert "README template not found" in caplog.text + assert not (output_folder / "README.md").exists() + + def test_skips_none_config_values_in_placeholders(self, converter, template_dir, tmp_path, monkeypatch): + """copy_readme skips None values in model_config without crashing.""" + _set_template_root(monkeypatch, template_dir) + (template_dir / "README-torchvision-fp16.md").write_text("<>") + output_folder = tmp_path / "test_model-fp16-ov" + output_folder.mkdir() + + converter.copy_readme( + { + "model_short_name": "test_model", + "model_library": "torchvision", + "license": "Apache-2.0", + "license_link": "https://apache.org/licenses/LICENSE-2.0", + "some_field": None, + }, + output_folder, + variant="fp16", + ) + + content = (output_folder / "README.md").read_text() + assert "test_model" in content + + +class TestCollectDatasetEntries: + """Tests for BaseConverter._collect_dataset_entries via TorchvisionConverter.""" + + def test_collects_supported_images_with_numeric_labels(self, converter, dataset_dir): + """Dataset collection returns supported image paths paired with class labels.""" + extra_class_dir = dataset_dir / "2" + extra_class_dir.mkdir() + ignored_file = extra_class_dir / "ignore.bmp" + ignored_file.write_bytes(b"bmp") + (dataset_dir / "README.txt").write_text("ignore me") + + entries = converter._collect_dataset_entries(dataset_dir) + + assert [(path.name, label) for path, label in entries] == [ + ("image_001.jpg", 0), + ("image_002.jpg", 1), + ] + + def test_dataset_type_dispatches_to_coco_reader(self, converter, tmp_path): + """Passing dataset_type='coco-detection' uses the COCO images reader.""" + images = tmp_path / "images" + annotations = tmp_path / "annotations" + images.mkdir() + annotations.mkdir() + (images / "img1.jpg").write_bytes(b"") + (annotations / "instances_val2017.json").write_text("{}") + + entries = converter._collect_dataset_entries(tmp_path, dataset_type="coco-detection") + + assert len(entries) == 1 + assert entries[0][0].name == "img1.jpg" + assert entries[0][1] == 0 # placeholder label for non-classification layouts + + def test_dataset_type_dispatches_to_ade20k_reader(self, converter, tmp_path): + """Passing dataset_type='ade20k' uses the ADE20K image/mask reader.""" + images = tmp_path / "images" + annotations = tmp_path / "annotations" + images.mkdir() + annotations.mkdir() + (images / "ADE_val_00000001.jpg").write_bytes(b"") + (annotations / "ADE_val_00000001.png").write_bytes(b"") + + entries = converter._collect_dataset_entries(tmp_path, dataset_type="ade20k") + + assert len(entries) == 1 + assert entries[0][0].name == "ADE_val_00000001.jpg" + + +class TestCropResize: + """Tests for BaseConverter._crop_resize.""" + + def test_square_target_tall_image(self, converter): + """Tall image is center-cropped to a square before resizing.""" + img = np.zeros((60, 40, 3), dtype=np.uint8) + # Mark the center column so we can verify cropping happened + img[:, 10:30, :] = 255 + result = converter._crop_resize(img, width=40, height=40) + assert result.shape == (40, 40, 3) + + def test_square_target_wide_image(self, converter): + """Wide image is center-cropped to a square before resizing.""" + img = np.zeros((40, 80, 3), dtype=np.uint8) + result = converter._crop_resize(img, width=20, height=20) + assert result.shape == (20, 20, 3) + + def test_square_target_already_square(self, converter): + """Square image is resized directly without cropping.""" + img = np.zeros((100, 100, 3), dtype=np.uint8) + result = converter._crop_resize(img, width=50, height=50) + assert result.shape == (50, 50, 3) + + def test_non_square_wide_target(self, converter): + """Target wider than source aspect ratio crops height.""" + img = np.zeros((100, 100, 3), dtype=np.uint8) + result = converter._crop_resize(img, width=200, height=100) + assert result.shape == (100, 200, 3) + + def test_non_square_tall_target(self, converter): + """Target taller than source aspect ratio crops width.""" + img = np.zeros((100, 100, 3), dtype=np.uint8) + result = converter._crop_resize(img, width=100, height=200) + assert result.shape == (200, 100, 3) + + +class TestPreprocessCalibrationImage: + """Tests for BaseConverter._preprocess_calibration_image via TorchvisionConverter.""" + + def test_preprocesses_image_and_reverses_channels(self, converter, tmp_path): + """Image preprocessing resizes, normalizes, and switches BGR to RGB when requested.""" + img_path = tmp_path / "pixel.png" + img = np.array([[[10, 20, 30]]], dtype=np.uint8) + cv2.imwrite(str(img_path), img) + + result = converter._preprocess_calibration_image( + img_path=img_path, + width=1, + height=1, + mean=np.array([0.0, 0.0, 0.0], dtype=np.float32), + scale=np.array([1.0, 1.0, 1.0], dtype=np.float32), + reverse_input_channels=True, + ) + + assert result is not None + assert result.shape == (1, 3, 1, 1) + assert result.dtype == np.float32 + np.testing.assert_allclose(result[:, :, 0, 0], np.array([[30.0, 20.0, 10.0]], dtype=np.float32)) + + def test_returns_none_for_invalid_image(self, converter, tmp_path): + """Image preprocessing returns None when cv2 cannot decode the file.""" + bad_path = tmp_path / "not-an-image.txt" + bad_path.write_text("not an image") + + result = converter._preprocess_calibration_image( + img_path=bad_path, + width=8, + height=8, + mean=np.array([0.0, 0.0, 0.0], dtype=np.float32), + scale=np.array([1.0, 1.0, 1.0], dtype=np.float32), + reverse_input_channels=False, + ) + + assert result is None + + def test_crop_resize_type_produces_square_output(self, converter, tmp_path): + """resize_type='crop' center-crops a wide image before resizing to target.""" + img_path = tmp_path / "wide.png" + # 4x8 wide image (H=4, W=8) + img = np.zeros((4, 8, 3), dtype=np.uint8) + cv2.imwrite(str(img_path), img) + + result = converter._preprocess_calibration_image( + img_path=img_path, + width=2, + height=2, + mean=np.array([0.0, 0.0, 0.0], dtype=np.float32), + scale=np.array([1.0, 1.0, 1.0], dtype=np.float32), + reverse_input_channels=False, + resize_type="crop", + ) + + assert result is not None + assert result.shape == (1, 3, 2, 2) + + def test_standard_resize_type_is_default(self, converter, tmp_path): + """Omitting resize_type uses the 'standard' plain-resize path.""" + img_path = tmp_path / "img.png" + img = np.zeros((8, 8, 3), dtype=np.uint8) + cv2.imwrite(str(img_path), img) + + result = converter._preprocess_calibration_image( + img_path=img_path, + width=4, + height=4, + mean=np.array([0.0, 0.0, 0.0], dtype=np.float32), + scale=np.array([1.0, 1.0, 1.0], dtype=np.float32), + reverse_input_channels=False, + ) + + assert result is not None + assert result.shape == (1, 3, 4, 4) + + +class TestCreateCalibrationDataset: + """Tests for BaseConverter.create_calibration_dataset via TorchvisionConverter.""" + + def test_returns_empty_list_when_dataset_path_is_missing(self, tmp_path): + """Calibration dataset creation is skipped when dataset_path is None.""" + converter = TorchvisionConverter( + output_dir=tmp_path / "out", + cache_dir=tmp_path / "cache", + dataset_registry=None, + ) + + assert converter.create_calibration_dataset(input_shape=[1, 3, 224, 224]) == ([], []) + + def test_returns_empty_tuple_when_dataset_dir_does_not_exist(self, tmp_path): + """Calibration dataset returns empty tuple when image_dir disappears between checks.""" + converter = TorchvisionConverter( + output_dir=tmp_path / "out", + cache_dir=tmp_path / "cache", + dataset_registry=None, + ) + # Mock dataset_path to simulate a race condition: exists() returns True first then False + mock_path = MagicMock(spec=Path) + mock_path.exists.side_effect = [True, False] + + result = converter.create_calibration_dataset(input_shape=[1, 3, 224, 224], dataset_path=mock_path) + + assert result == ([], []) + + def test_returns_empty_tuple_when_no_images_found(self, tmp_path): + """Calibration dataset returns empty tuple when dataset directory has no images.""" + empty_dataset = tmp_path / "empty_dataset" + empty_dataset.mkdir() + + converter = TorchvisionConverter( + output_dir=tmp_path / "out", + cache_dir=tmp_path / "cache", + dataset_registry=None, + ) + + result = converter.create_calibration_dataset(input_shape=[1, 3, 224, 224], dataset_path=empty_dataset) + + assert result == ([], []) + + def test_returns_images_and_labels(self, tmp_path, dataset_dir): + """Calibration dataset returns preprocessed images and class labels when requested.""" + converter = TorchvisionConverter( + output_dir=tmp_path / "out", + cache_dir=tmp_path / "cache", + dataset_registry=None, + ) + + images, labels = converter.create_calibration_dataset( + input_shape=[1, 3, 224, 224], + return_labels=True, + dataset_path=dataset_dir, + ) + + assert len(images) == 2 + assert labels == [0, 1] + assert all(image.shape == (1, 3, 224, 224) for image in images) + + def test_passes_normalization_options_and_subset_size(self, tmp_path, dataset_dir): + """Calibration dataset forwards normalization options to per-image preprocessing.""" + converter = TorchvisionConverter( + output_dir=tmp_path / "out", + cache_dir=tmp_path / "cache", + dataset_registry=None, + ) + seen_calls: list[dict[str, object]] = [] + + def fake_preprocess(**kwargs): + seen_calls.append(kwargs) + return np.zeros((1, 3, kwargs["height"], kwargs["width"]), dtype=np.float32) + + converter._preprocess_calibration_image = fake_preprocess # type: ignore[method-assign] + + images, labels = converter.create_calibration_dataset( + input_shape=[1, 3, 224, 224], + mean_values="1 2 3", + scale_values="4 5 6", + reverse_input_channels=False, + subset_size=1, + return_labels=False, + dataset_path=dataset_dir, + ) + + assert len(images) == 1 + assert labels == [] + assert len(seen_calls) == 1 + np.testing.assert_array_equal(seen_calls[0]["mean"], np.array([1.0, 2.0, 3.0])) + np.testing.assert_array_equal(seen_calls[0]["scale"], np.array([4.0, 5.0, 6.0])) + assert seen_calls[0]["reverse_input_channels"] is False + + def test_skips_unreadable_images_in_return_labels_path(self, tmp_path): + """Calibration dataset skips images that fail preprocessing when return_labels=True.""" + dataset_path = tmp_path / "dataset" + class_dir = dataset_path / "0" + class_dir.mkdir(parents=True) + + (class_dir / "bad_image.jpg").write_text("not an image") + img = np.zeros((10, 10, 3), dtype=np.uint8) + cv2.imwrite(str(class_dir / "good_image.png"), img) + + converter = TorchvisionConverter( + output_dir=tmp_path / "out", + cache_dir=tmp_path / "cache", + dataset_registry=None, + ) + + images, labels = converter.create_calibration_dataset( + input_shape=[1, 3, 8, 8], + return_labels=True, + dataset_path=dataset_path, + ) + + assert len(images) == 1 + assert labels == [0] + + def test_handles_preprocess_exception_in_return_labels_path(self, tmp_path, dataset_dir): + """Calibration dataset logs warning and continues when preprocessing raises.""" + converter = TorchvisionConverter( + output_dir=tmp_path / "out", + cache_dir=tmp_path / "cache", + dataset_registry=None, + ) + + def raising_preprocess(**kwargs): + disk_error_msg = "disk error" + raise OSError(disk_error_msg) + + converter._preprocess_calibration_image = raising_preprocess # type: ignore[method-assign] + + images, labels = converter.create_calibration_dataset( + input_shape=[1, 3, 224, 224], + return_labels=True, + dataset_path=dataset_dir, + ) + + assert images == [] + assert labels == [] + + def test_skips_unreadable_images_in_non_return_labels_path(self, tmp_path): + """Calibration dataset skips images that fail preprocessing when return_labels=False.""" + dataset_path = tmp_path / "dataset" + class_dir = dataset_path / "0" + class_dir.mkdir(parents=True) + + (class_dir / "bad_image.jpg").write_text("not an image") + img = np.zeros((10, 10, 3), dtype=np.uint8) + cv2.imwrite(str(class_dir / "good_image.png"), img) + + converter = TorchvisionConverter( + output_dir=tmp_path / "out", + cache_dir=tmp_path / "cache", + dataset_registry=None, + ) + + images, labels = converter.create_calibration_dataset( + input_shape=[1, 3, 8, 8], + return_labels=False, + dataset_path=dataset_path, + ) + + assert len(images) == 1 + assert labels == [] + + def test_handles_preprocess_exception_in_non_return_labels_path(self, tmp_path, dataset_dir): + """Calibration dataset logs warning and continues when preprocessing raises (no labels).""" + converter = TorchvisionConverter( + output_dir=tmp_path / "out", + cache_dir=tmp_path / "cache", + dataset_registry=None, + ) + + def raising_preprocess(**kwargs): + bad_pixel_msg = "bad pixel data" + raise ValueError(bad_pixel_msg) + + converter._preprocess_calibration_image = raising_preprocess # type: ignore[method-assign] + + images, labels = converter.create_calibration_dataset( + input_shape=[1, 3, 224, 224], + return_labels=False, + dataset_path=dataset_dir, + ) + + assert images == [] + assert labels == [] + + def test_logs_progress_every_50_images_return_labels(self, tmp_path): + """Calibration dataset logs progress every 50 images in return_labels path.""" + dataset_path = tmp_path / "dataset" + class_dir = dataset_path / "0" + class_dir.mkdir(parents=True) + # Create 50 image files to trigger the modulo-50 logging + for i in range(50): + (class_dir / f"image_{i:04d}.jpg").write_bytes(b"fake") + + converter = TorchvisionConverter( + output_dir=tmp_path / "out", + cache_dir=tmp_path / "cache", + dataset_registry=None, + ) + + call_count = [0] + + def fake_preprocess(**kwargs): + call_count[0] += 1 + return np.zeros((1, 3, kwargs["height"], kwargs["width"]), dtype=np.float32) + + converter._preprocess_calibration_image = fake_preprocess # type: ignore[method-assign] + + images, _ = converter.create_calibration_dataset( + input_shape=[1, 3, 8, 8], + return_labels=True, + subset_size=50, + dataset_path=dataset_path, + ) + + assert len(images) == 50 + assert call_count[0] == 50 + + def test_logs_progress_every_50_images_no_labels(self, tmp_path): + """Calibration dataset logs progress every 50 images in non-return_labels path.""" + dataset_path = tmp_path / "dataset" + class_dir = dataset_path / "0" + class_dir.mkdir(parents=True) + for i in range(50): + (class_dir / f"image_{i:04d}.jpg").write_bytes(b"fake") + + converter = TorchvisionConverter( + output_dir=tmp_path / "out", + cache_dir=tmp_path / "cache", + dataset_registry=None, + ) + + def fake_preprocess(**kwargs): + return np.zeros((1, 3, kwargs["height"], kwargs["width"]), dtype=np.float32) + + converter._preprocess_calibration_image = fake_preprocess # type: ignore[method-assign] + + images, labels = converter.create_calibration_dataset( + input_shape=[1, 3, 8, 8], + return_labels=False, + subset_size=50, + dataset_path=dataset_path, + ) + + assert len(images) == 50 + assert labels == [] + + def test_forwards_resize_type_to_preprocess(self, tmp_path, dataset_dir): + """create_calibration_dataset passes resize_type to _preprocess_calibration_image.""" + converter = TorchvisionConverter( + output_dir=tmp_path / "out", + cache_dir=tmp_path / "cache", + dataset_registry=None, + ) + seen_resize_types: list[str] = [] + + def fake_preprocess(**kwargs): + seen_resize_types.append(kwargs.get("resize_type", "standard")) + return np.zeros((1, 3, kwargs["height"], kwargs["width"]), dtype=np.float32) + + converter._preprocess_calibration_image = fake_preprocess # type: ignore[method-assign] + + converter.create_calibration_dataset( + input_shape=[1, 3, 8, 8], + resize_type="crop", + return_labels=False, + dataset_path=dataset_dir, + ) + + assert all(rt == "crop" for rt in seen_resize_types) + + def test_forwards_resize_type_to_preprocess_with_labels(self, tmp_path, dataset_dir): + """create_calibration_dataset passes resize_type to _preprocess_calibration_image (return_labels=True).""" + converter = TorchvisionConverter( + output_dir=tmp_path / "out", + cache_dir=tmp_path / "cache", + dataset_registry=None, + ) + seen_resize_types: list[str] = [] + + def fake_preprocess(**kwargs): + seen_resize_types.append(kwargs.get("resize_type", "standard")) + return np.zeros((1, 3, kwargs["height"], kwargs["width"]), dtype=np.float32) + + converter._preprocess_calibration_image = fake_preprocess # type: ignore[method-assign] + + converter.create_calibration_dataset( + input_shape=[1, 3, 8, 8], + resize_type="crop", + return_labels=True, + dataset_path=dataset_dir, + ) + + assert all(rt == "crop" for rt in seen_resize_types) + + """Tests for BaseConverter.validate_model via TorchvisionConverter.""" + + def test_computes_top1_accuracy(self, converter, tmp_path): + """Model validation returns the fraction of correct predictions.""" + output_layer = object() + compiled_model = MagicMock() + compiled_model.outputs = [output_layer] + compiled_model.side_effect = [ + {output_layer: np.array([[0.1, 0.9, 0.0]])}, + {output_layer: np.array([[0.8, 0.1, 0.1]])}, + ] + fake_ov = ModuleType("openvino") + fake_ov.Core = MagicMock( + return_value=SimpleNamespace( + read_model=MagicMock(return_value="model"), + compile_model=MagicMock(return_value=compiled_model), + ), + ) + + with patch.dict(sys.modules, {"openvino": fake_ov}): + accuracy = converter.validate_model( + model_path=tmp_path / "model.xml", + validation_data=[np.zeros((1, 3, 224, 224)), np.zeros((1, 3, 224, 224))], + labels=[1, 0], + ) + + assert accuracy == pytest.approx(1.0) + + def test_returns_zero_when_openvino_fails(self, converter, tmp_path): + """Model validation returns 0.0 when OpenVINO raises an error.""" + fake_ov = ModuleType("openvino") + fake_ov.Core = MagicMock(side_effect=RuntimeError("OV error")) + + with patch.dict(sys.modules, {"openvino": fake_ov}): + accuracy = converter.validate_model( + model_path=tmp_path / "model.xml", + validation_data=[np.zeros((1, 3, 224, 224))], + labels=[0], + ) + + assert accuracy == pytest.approx(0.0) + + +class TestValidateTorchModel: + """Tests for PyTorchConverter.validate_torch_model via TorchvisionConverter.""" + + def test_computes_top1_accuracy(self, converter): + """Top-1 accuracy is computed from the original PyTorch model outputs.""" + import torch + + model = MagicMock() + model.side_effect = [ + (torch.tensor([[0.1, 0.9, 0.0]]),), + torch.tensor([[0.8, 0.1, 0.1]]), + ] + + accuracy = converter.validate_torch_model( + model, + [np.zeros((1, 3, 224, 224), dtype=np.float64), np.zeros((1, 3, 224, 224), dtype=np.float64)], + [1, 0], + ) + + model.eval.assert_called_once() + assert accuracy == pytest.approx(1.0) + # Inputs are cast to float32 to match PyTorch model weights. + assert model.call_args.args[0].dtype == torch.float32 + + def test_returns_none_when_inference_fails(self, converter): + """Validation returns None when the PyTorch model raises an error.""" + model = MagicMock(side_effect=RuntimeError("forward failed")) + + accuracy = converter.validate_torch_model( + model, + [np.zeros((1, 3, 224, 224), dtype=np.float32)], + [0], + ) + + assert accuracy is None + + """Tests for BaseConverter.quantize_model via TorchvisionConverter.""" + + def test_returns_original_model_when_no_calibration_data(self, converter, sample_model_config, tmp_path): + """Quantization is skipped when calibration data is empty.""" + model_path = tmp_path / "model.xml" + + assert converter.quantize_model(model_path, [], sample_model_config) == model_path + + def test_quantizes_model_and_writes_artifacts( + self, + converter, + sample_model_config, + template_dir, + tmp_path, + monkeypatch, + ): + """Quantization saves the INT8 model plus metadata, README, and gitattributes.""" + _set_template_root(monkeypatch, template_dir) + (template_dir / "README-torchvision-int8.md").write_text( + "# <>\nVariant: <>\nTags:\n<>\nDocs: <>\n", + ) + fp16_dir = tmp_path / "test_model-fp16-ov" + fp16_dir.mkdir(parents=True) + model_path = fp16_dir / "test_model_fp32.xml" + model_path.write_text("") + calibration_data = [ + np.zeros((1, 3, 224, 224), dtype=np.float32), + np.ones((1, 3, 224, 224), dtype=np.float32), + ] + + quantized_model = MagicMock() + quantized_model.get_rt_info.return_value = SimpleNamespace(value={"model_type": "Classification"}) + core = SimpleNamespace(read_model=MagicMock(return_value="ov_model")) + save_model = MagicMock( + side_effect=lambda _model, path, compress_to_fp16=True: Path(path).write_text("") or None, + ) + fake_ov = ModuleType("openvino") + fake_ov.Core = MagicMock(return_value=core) + fake_ov.save_model = save_model + + dataset_factory = MagicMock(side_effect=lambda generator: list(generator)) + quantize = MagicMock(return_value=quantized_model) + fake_nncf = ModuleType("nncf") + fake_nncf.Dataset = dataset_factory + fake_nncf.quantize = quantize + fake_nncf.QuantizationPreset = SimpleNamespace(PERFORMANCE="performance", MIXED="mixed") + + model_config = sample_model_config | {"tags": ["vision", "classification"]} + + with patch.dict(sys.modules, {"openvino": fake_ov, "nncf": fake_nncf}): + result = converter.quantize_model( + model_path=model_path, + calibration_data=calibration_data, + model_config=model_config, + preset="performance", + ) + + output_folder = tmp_path / "test_model-int8-ov" + assert result == output_folder / "test_model.xml" + assert result.read_text() == "" + assert json.loads((output_folder / "config.json").read_text()) == {"model_type": "Classification"} + assert (output_folder / ".gitattributes").read_text() == (template_dir / ".gitattributes").read_text() + readme = (output_folder / "README.md").read_text() + assert "# test_model" in readme + assert "Variant: int8" in readme + assert " - vision" in readme + assert "Docs: https://docs.example.com" in readme + dataset_factory.assert_called_once() + quantize.assert_called_once_with( + "ov_model", + calibration_dataset=calibration_data, + preset="performance", + subset_size=2, + ) + save_model.assert_called_once() + + def test_passes_transformer_model_type_when_configured( + self, + converter, + sample_model_config, + template_dir, + tmp_path, + monkeypatch, + ): + """Transformer models forward NNCF's TRANSFORMER model type to quantization.""" + _set_template_root(monkeypatch, template_dir) + (template_dir / "README-torchvision-int8.md").write_text("# <>") + fp16_dir = tmp_path / "test_model-fp16-ov" + fp16_dir.mkdir(parents=True) + model_path = fp16_dir / "test_model_fp32.xml" + model_path.write_text("") + calibration_data = [np.zeros((1, 3, 224, 224), dtype=np.float32)] + + quantized_model = MagicMock() + quantized_model.get_rt_info.return_value = SimpleNamespace(value={"model_type": "Classification"}) + core = SimpleNamespace(read_model=MagicMock(return_value="ov_model")) + fake_ov = ModuleType("openvino") + fake_ov.Core = MagicMock(return_value=core) + fake_ov.save_model = MagicMock( + side_effect=lambda _model, path, compress_to_fp16=True: Path(path).write_text("") or None, + ) + quantize = MagicMock(return_value=quantized_model) + fake_nncf = ModuleType("nncf") + fake_nncf.Dataset = MagicMock(side_effect=lambda generator: list(generator)) + fake_nncf.quantize = quantize + fake_nncf.QuantizationPreset = SimpleNamespace(PERFORMANCE="performance", MIXED="mixed") + fake_nncf.ModelType = SimpleNamespace(TRANSFORMER="transformer") + + model_config = sample_model_config | {"quantization_model_type": "transformer"} + + with patch.dict(sys.modules, {"openvino": fake_ov, "nncf": fake_nncf}): + converter.quantize_model( + model_path=model_path, + calibration_data=calibration_data, + model_config=model_config, + preset="mixed", + ) + + assert quantize.call_args.kwargs["model_type"] == "transformer" + + def test_validates_fp32_and_int8_outputs_when_requested( + self, + converter, + sample_model_config, + template_dir, + tmp_path, + monkeypatch, + ): + """Quantization validates both source and INT8 models when validation data is provided.""" + _set_template_root(monkeypatch, template_dir) + (template_dir / "README-torchvision-int8.md").write_text("# <>") + fp16_dir = tmp_path / "test_model-fp16-ov" + fp16_dir.mkdir(parents=True) + model_path = fp16_dir / "test_model_fp32.xml" + model_path.write_text("") + + quantized_model = MagicMock() + quantized_model.get_rt_info.return_value = SimpleNamespace(value={"model_type": "Classification"}) + fake_ov = ModuleType("openvino") + fake_ov.Core = MagicMock(return_value=SimpleNamespace(read_model=MagicMock(return_value="ov_model"))) + fake_ov.save_model = MagicMock() + fake_nncf = ModuleType("nncf") + fake_nncf.Dataset = MagicMock(side_effect=lambda generator: list(generator)) + fake_nncf.quantize = MagicMock(return_value=quantized_model) + fake_nncf.QuantizationPreset = SimpleNamespace(PERFORMANCE="performance", MIXED="mixed") + + calibration_data = [np.zeros((1, 3, 224, 224), dtype=np.float32)] + validation_data = [np.zeros((1, 3, 224, 224), dtype=np.float32)] + validation_labels = [0] + converter.validate_model = MagicMock(side_effect=[0.97, 0.95]) + + with patch.dict(sys.modules, {"openvino": fake_ov, "nncf": fake_nncf}): + converter.quantize_model( + model_path=model_path, + calibration_data=calibration_data, + model_config=sample_model_config, + validation_data=validation_data, + validation_labels=validation_labels, + ) + + call_args = converter.validate_model.call_args_list + assert call_args[0].args == (model_path, validation_data, validation_labels) + assert call_args[1].args == ( + tmp_path / "test_model-int8-ov" / "test_model.xml", + validation_data, + validation_labels, + ) + assert converter.validate_model.call_count == 2 + + def test_measures_fp16_accuracy_and_populates_results( + self, + converter, + sample_model_config, + tmp_path, + template_dir, + monkeypatch, + ): + """Quantization measures FP16 accuracy and fills the AccuracyResults collector.""" + from model_converter.reporting import AccuracyResults + + _set_template_root(monkeypatch, template_dir) + (template_dir / "README-torchvision-int8.md").write_text("# <>") + fp16_dir = tmp_path / "test_model-fp16-ov" + fp16_dir.mkdir(parents=True) + model_path = fp16_dir / "test_model_fp32.xml" + model_path.write_text("") + (fp16_dir / "test_model.xml").write_text("") + + quantized_model = MagicMock() + quantized_model.get_rt_info.return_value = SimpleNamespace(value={"model_type": "Classification"}) + fake_ov = ModuleType("openvino") + fake_ov.Core = MagicMock(return_value=SimpleNamespace(read_model=MagicMock(return_value="ov_model"))) + fake_ov.save_model = MagicMock() + fake_nncf = ModuleType("nncf") + fake_nncf.Dataset = MagicMock(side_effect=lambda generator: list(generator)) + fake_nncf.quantize = MagicMock(return_value=quantized_model) + fake_nncf.QuantizationPreset = SimpleNamespace(PERFORMANCE="performance", MIXED="mixed") + + calibration_data = [np.zeros((1, 3, 224, 224), dtype=np.float32)] + validation_data = [np.zeros((1, 3, 224, 224), dtype=np.float32)] + validation_labels = [0] + converter.validate_model = MagicMock(side_effect=[0.97, 0.96, 0.95]) + accuracy = AccuracyResults() + + with patch.dict(sys.modules, {"openvino": fake_ov, "nncf": fake_nncf}): + converter.quantize_model( + model_path=model_path, + calibration_data=calibration_data, + model_config=sample_model_config, + validation_data=validation_data, + validation_labels=validation_labels, + accuracy_results=accuracy, + ) + + assert converter.validate_model.call_count == 3 + assert accuracy.measured is True + assert accuracy.int8_succeeded is True + assert accuracy.fp32_accuracy == pytest.approx(0.97) + assert accuracy.fp16_accuracy == pytest.approx(0.96) + assert accuracy.int8_accuracy == pytest.approx(0.95) + + def test_record_result_copies_measured_accuracy(self, converter, sample_model_config): + """_record_result copies measured accuracies onto the conversion result.""" + from model_converter.reporting import STATUS_OK, AccuracyResults + + accuracy = AccuracyResults( + fp32_accuracy=0.9, + fp16_accuracy=0.89, + int8_accuracy=0.88, + int8_succeeded=True, + measured=True, + ) + result = converter._record_result( + converter._build_result(sample_model_config), + converted=True, + quantized=True, + accuracy=accuracy, + ) + + assert result.fp32_accuracy == pytest.approx(0.9) + assert result.fp16_accuracy == pytest.approx(0.89) + assert result.int8_accuracy == pytest.approx(0.88) + assert result.status == STATUS_OK + assert converter.results[-1] is result + + def test_record_result_calls_upsert_when_report_path_set( + self, + tmp_output_dir, + tmp_cache_dir, + sample_model_config, + tmp_path, + ): + """_record_result calls upsert_result when report_path is set and not skipped.""" + from unittest.mock import patch + + from model_converter.converters.torchvision import TorchvisionConverter + + report_path = tmp_path / "report.md" + conv = TorchvisionConverter(output_dir=tmp_output_dir, cache_dir=tmp_cache_dir, report_path=report_path) + + with patch("model_converter.converters.base.upsert_result") as mock_upsert: + conv._record_result(conv._build_result(sample_model_config), converted=True, quantized=True) + + mock_upsert.assert_called_once() + assert mock_upsert.call_args.args[1] == report_path + + def test_record_result_skips_upsert_when_skipped( + self, + tmp_output_dir, + tmp_cache_dir, + sample_model_config, + tmp_path, + ): + """_record_result does not call upsert_result when the model export is skipped.""" + from unittest.mock import patch + + from model_converter.converters.torchvision import TorchvisionConverter + + report_path = tmp_path / "report.md" + conv = TorchvisionConverter(output_dir=tmp_output_dir, cache_dir=tmp_cache_dir, report_path=report_path) + + with patch("model_converter.converters.base.upsert_result") as mock_upsert: + conv._record_result(conv._build_result(sample_model_config), converted=False, quantized=False, skipped=True) + + mock_upsert.assert_not_called() + + def test_record_result_skips_upsert_when_no_report_path(self, converter, sample_model_config, tmp_path): + """_record_result does not call upsert_result when report_path is None.""" + from unittest.mock import patch + + with patch("model_converter.converters.base.upsert_result") as mock_upsert: + converter._record_result(converter._build_result(sample_model_config), converted=True, quantized=True) + + mock_upsert.assert_not_called() + + def test_returns_original_path_when_nncf_not_installed(self, converter, sample_model_config, tmp_path): + """Quantization returns original path when nncf is not installed.""" + model_path = tmp_path / "model.xml" + model_path.write_text("") + calibration_data = [np.zeros((1, 3, 224, 224), dtype=np.float32)] + + with patch.dict(sys.modules, {"nncf": None, "openvino": None}): + result = converter.quantize_model(model_path, calibration_data, sample_model_config) + + assert result == model_path + + def test_returns_original_path_when_quantize_raises_runtime_error( + self, + converter, + sample_model_config, + tmp_path, + ): + """Quantization returns original path when a runtime error occurs.""" + model_path = tmp_path / "test_model-fp16-ov" / "test_model_fp32.xml" + model_path.parent.mkdir(parents=True) + model_path.write_text("") + calibration_data = [np.zeros((1, 3, 224, 224), dtype=np.float32)] + + fake_ov = ModuleType("openvino") + fake_ov.Core = MagicMock(return_value=SimpleNamespace(read_model=MagicMock(return_value="ov_model"))) + fake_nncf = ModuleType("nncf") + fake_nncf.Dataset = MagicMock(side_effect=lambda generator: list(generator)) + fake_nncf.quantize = MagicMock(side_effect=RuntimeError("quantization failed")) + fake_nncf.QuantizationPreset = SimpleNamespace(PERFORMANCE="performance", MIXED="mixed") + + with patch.dict(sys.modules, {"openvino": fake_ov, "nncf": fake_nncf}): + result = converter.quantize_model(model_path, calibration_data, sample_model_config) + + assert result == model_path + + +class TestMeasureMetric: + """Tests for BaseConverter._measure_metric (Model API → metric dispatch).""" + + @staticmethod + def _install_fake_model_api(monkeypatch, wrapper): + """Install a fake ``model_api.models`` module so ``Model.create_model`` returns ``wrapper``.""" + fake_models = ModuleType("model_api.models") + + class FakeModel: + @staticmethod + def create_model(_path): + return wrapper + + fake_models.Model = FakeModel + fake_root = ModuleType("model_api") + fake_root.models = fake_models + monkeypatch.setitem(sys.modules, "model_api", fake_root) + monkeypatch.setitem(sys.modules, "model_api.models", fake_models) + + def test_multilabel_dispatch_uses_raw_scores_and_returns_compute( + self, + converter, + tmp_path, + monkeypatch, + ): + from model_converter.datasets import CalibrationSample + from model_converter.metrics import MultilabelMAP + + img_path = tmp_path / "img.jpg" + cv2.imwrite(str(img_path), np.zeros((10, 10, 3), dtype=np.uint8)) + wrapper = MagicMock(return_value=SimpleNamespace(raw_scores=np.array([0.1, 0.9, 0.2], dtype=np.float32))) + self._install_fake_model_api(monkeypatch, wrapper) + + metric = MultilabelMAP(num_labels=3) + metric.update = MagicMock() + metric.compute = MagicMock(return_value=0.42) + metric.reset = MagicMock() + sample = CalibrationSample(image_path=img_path, label=1) + + value = converter._measure_metric(tmp_path / "model.xml", [sample], metric) + + assert value == pytest.approx(0.42) + metric.reset.assert_called_once() + update_kwargs = metric.update.call_args.kwargs + assert "prediction" in update_kwargs + assert "ground_truth" in update_kwargs + np.testing.assert_array_equal(update_kwargs["ground_truth"], np.array([0, 1, 0], dtype=np.int64)) + + def test_bbox_dispatch_converts_xyxy_to_xywh_and_uses_image_id( + self, + converter, + tmp_path, + monkeypatch, + ): + from model_converter.datasets import CalibrationSample + from model_converter.metrics import CocoDetectionMAP + + img_path = tmp_path / "img.jpg" + cv2.imwrite(str(img_path), np.zeros((10, 10, 3), dtype=np.uint8)) + wrapper = MagicMock( + return_value=SimpleNamespace( + bboxes=np.array([[10, 20, 30, 50]], dtype=np.int32), + labels=np.array([4], dtype=np.int32), + scores=np.array([0.9], dtype=np.float32), + ), + ) + self._install_fake_model_api(monkeypatch, wrapper) + + ann = tmp_path / "ann.json" + ann.write_text('{"images":[],"annotations":[],"categories":[]}') + metric = CocoDetectionMAP(annotation_file=ann, iou_type="bbox") + metric.update = MagicMock() + metric.compute = MagicMock(return_value=0.55) + metric.reset = MagicMock() + sample = CalibrationSample(image_path=img_path, label=0, image_id=42) + + value = converter._measure_metric(tmp_path / "model.xml", [sample], metric) + + assert value == pytest.approx(0.55) + preds = metric.update.call_args.kwargs["predictions"] + assert len(preds) == 1 + assert preds[0]["image_id"] == 42 + # XYXY (10, 20, 30, 50) → XYWH (10, 20, 20, 30) + assert preds[0]["bbox"] == [10.0, 20.0, 20.0, 30.0] + # class index 4 (airplane) → COCO80_TO_COCO91[4] = 5 + assert preds[0]["category_id"] == 5 + assert preds[0]["score"] == pytest.approx(0.9) + + def test_bbox_dispatch_uses_coco80_to_coco91_mapping_for_class_11( + self, + converter, + tmp_path, + monkeypatch, + ): + """Class index 11 (stop sign) must map to COCO category ID 13, not 12.""" + from model_converter.datasets import CalibrationSample + from model_converter.metrics import CocoDetectionMAP + + img_path = tmp_path / "img.jpg" + cv2.imwrite(str(img_path), np.zeros((10, 10, 3), dtype=np.uint8)) + wrapper = MagicMock( + return_value=SimpleNamespace( + bboxes=np.array([[0, 0, 5, 5]], dtype=np.int32), + labels=np.array([11], dtype=np.int32), # stop sign — first class where +1 diverges + scores=np.array([0.8], dtype=np.float32), + ), + ) + self._install_fake_model_api(monkeypatch, wrapper) + + ann = tmp_path / "ann.json" + ann.write_text('{"images":[],"annotations":[],"categories":[]}') + metric = CocoDetectionMAP(annotation_file=ann, iou_type="bbox") + metric.update = MagicMock() + metric.compute = MagicMock(return_value=0.0) + metric.reset = MagicMock() + sample = CalibrationSample(image_path=img_path, label=0, image_id=1) + + converter._measure_metric(tmp_path / "model.xml", [sample], metric) + + preds = metric.update.call_args.kwargs["predictions"] + assert preds[0]["category_id"] == 13 # COCO80_TO_COCO91[11] = 13 (not 12) + + def test_bbox_dispatch_skips_samples_without_image_id( + self, + converter, + tmp_path, + monkeypatch, + ): + from model_converter.datasets import CalibrationSample + from model_converter.metrics import CocoDetectionMAP + + img_path = tmp_path / "img.jpg" + cv2.imwrite(str(img_path), np.zeros((10, 10, 3), dtype=np.uint8)) + wrapper = MagicMock( + return_value=SimpleNamespace( + bboxes=np.array([[1, 2, 3, 4]], dtype=np.int32), + labels=np.array([0], dtype=np.int32), + scores=np.array([0.5], dtype=np.float32), + ), + ) + self._install_fake_model_api(monkeypatch, wrapper) + + ann = tmp_path / "ann.json" + ann.write_text('{"images":[],"annotations":[],"categories":[]}') + metric = CocoDetectionMAP(annotation_file=ann, iou_type="bbox") + metric.update = MagicMock() + metric.compute = MagicMock(return_value=0.0) + sample = CalibrationSample(image_path=img_path, label=0) # image_id=None + + converter._measure_metric(tmp_path / "model.xml", [sample], metric) + + metric.update.assert_not_called() + + def test_metric_update_errors_are_swallowed_per_sample( + self, + converter, + tmp_path, + monkeypatch, + ): + """A bad single-sample dispatch must not abort metric measurement.""" + from model_converter.datasets import CalibrationSample + from model_converter.metrics import MultilabelMAP + + img_path = tmp_path / "img.jpg" + cv2.imwrite(str(img_path), np.zeros((10, 10, 3), dtype=np.uint8)) + # Wrapper returns malformed raw_scores that will trip metric.update. + wrapper = MagicMock(return_value=SimpleNamespace(raw_scores=np.array([0.1, 0.2, 0.3], dtype=np.float32))) + self._install_fake_model_api(monkeypatch, wrapper) + + metric = MultilabelMAP(num_labels=3) + metric.update = MagicMock(side_effect=TypeError("boom")) + metric.compute = MagicMock(return_value=0.0) + sample = CalibrationSample(image_path=img_path, label=0) + + value = converter._measure_metric(tmp_path / "m.xml", [sample], metric) + + assert value == pytest.approx(0.0) + metric.compute.assert_called_once() + + def test_semseg_dispatch_shifts_ade20k_mask_and_updates_metric( + self, + converter, + tmp_path, + monkeypatch, + ): + from model_converter.datasets import CalibrationSample + from model_converter.metrics import SemSegMIoU + + img_path = tmp_path / "img.jpg" + mask_path = tmp_path / "mask.png" + cv2.imwrite(str(img_path), np.zeros((10, 10, 3), dtype=np.uint8)) + # ADE20K convention: 0 = unlabeled (ignore), 1..N = classes + cv2.imwrite(str(mask_path), np.array([[0, 1, 2], [1, 2, 3]], dtype=np.uint8)) + + pred_mask = np.array([[0, 0, 1], [0, 1, 2]], dtype=np.uint8) + wrapper = MagicMock(return_value=SimpleNamespace(resultImage=pred_mask)) + self._install_fake_model_api(monkeypatch, wrapper) + + metric = SemSegMIoU(num_classes=150, ignore_index=255) + metric.update = MagicMock() + metric.compute = MagicMock(return_value=0.3) + sample = CalibrationSample(image_path=img_path, label=0, mask_path=mask_path) + + converter._measure_metric(tmp_path / "model.xml", [sample], metric) + + call = metric.update.call_args + gt_arg = call.args[1] if len(call.args) > 1 else call.kwargs["ground_truth"] + # Mask is shifted: 0 → 255 (ignore), 1 → 0, 2 → 1, 3 → 2 + np.testing.assert_array_equal(gt_arg, np.array([[255, 0, 1], [0, 1, 2]], dtype=np.int32)) + + def test_semseg_skips_sample_without_mask_path( + self, + converter, + tmp_path, + monkeypatch, + ): + from model_converter.datasets import CalibrationSample + from model_converter.metrics import SemSegMIoU + + img_path = tmp_path / "img.jpg" + cv2.imwrite(str(img_path), np.zeros((10, 10, 3), dtype=np.uint8)) + wrapper = MagicMock(return_value=SimpleNamespace(resultImage=np.zeros((10, 10), dtype=np.uint8))) + self._install_fake_model_api(monkeypatch, wrapper) + + metric = SemSegMIoU(num_classes=150, ignore_index=255) + metric.update = MagicMock() + metric.compute = MagicMock(return_value=0.0) + sample = CalibrationSample(image_path=img_path, label=0) # mask_path=None + + converter._measure_metric(tmp_path / "model.xml", [sample], metric) + + metric.update.assert_not_called() + + def test_returns_none_when_model_api_unavailable(self, converter, tmp_path, monkeypatch): + """If ``model_api`` cannot be imported, return None rather than raising.""" + from model_converter.metrics import MultilabelMAP + + monkeypatch.setitem(sys.modules, "model_api", None) + + metric = MultilabelMAP(num_labels=3) + assert converter._measure_metric(tmp_path / "m.xml", [], metric) is None + + def test_returns_none_when_create_model_raises(self, converter, tmp_path, monkeypatch): + from model_converter.metrics import MultilabelMAP + + fake_models = ModuleType("model_api.models") + + class FakeModel: + @staticmethod + def create_model(_path): + msg = "bad model" + raise RuntimeError(msg) + + fake_models.Model = FakeModel + fake_root = ModuleType("model_api") + fake_root.models = fake_models + monkeypatch.setitem(sys.modules, "model_api", fake_root) + monkeypatch.setitem(sys.modules, "model_api.models", fake_models) + + metric = MultilabelMAP(num_labels=3) + assert converter._measure_metric(tmp_path / "m.xml", [], metric) is None + + def test_skips_samples_that_fail_to_load(self, converter, tmp_path, monkeypatch): + from model_converter.datasets import CalibrationSample + from model_converter.metrics import MultilabelMAP + + wrapper = MagicMock() + self._install_fake_model_api(monkeypatch, wrapper) + + metric = MultilabelMAP(num_labels=3) + metric.update = MagicMock() + metric.compute = MagicMock(return_value=0.0) + sample = CalibrationSample(image_path=tmp_path / "nonexistent.jpg", label=0) + + converter._measure_metric(tmp_path / "m.xml", [sample], metric) + + wrapper.assert_not_called() + metric.update.assert_not_called() + + def test_inference_errors_are_swallowed_per_sample(self, converter, tmp_path, monkeypatch): + from model_converter.datasets import CalibrationSample + from model_converter.metrics import MultilabelMAP + + img_path = tmp_path / "img.jpg" + cv2.imwrite(str(img_path), np.zeros((10, 10, 3), dtype=np.uint8)) + wrapper = MagicMock(side_effect=RuntimeError("boom")) + self._install_fake_model_api(monkeypatch, wrapper) + + metric = MultilabelMAP(num_labels=3) + metric.update = MagicMock() + metric.compute = MagicMock(return_value=0.0) + sample = CalibrationSample(image_path=img_path, label=0) + + value = converter._measure_metric(tmp_path / "m.xml", [sample], metric) + + assert value == pytest.approx(0.0) + metric.update.assert_not_called() + + def test_multilabel_skips_when_result_has_no_raw_scores( + self, + converter, + tmp_path, + monkeypatch, + ): + from model_converter.datasets import CalibrationSample + from model_converter.metrics import MultilabelMAP + + img_path = tmp_path / "img.jpg" + cv2.imwrite(str(img_path), np.zeros((10, 10, 3), dtype=np.uint8)) + # Result lacks the ``raw_scores`` attribute — wrapper returns a bare namespace. + wrapper = MagicMock(return_value=SimpleNamespace()) + self._install_fake_model_api(monkeypatch, wrapper) + + metric = MultilabelMAP(num_labels=3) + metric.update = MagicMock() + metric.compute = MagicMock(return_value=0.0) + sample = CalibrationSample(image_path=img_path, label=0) + + converter._measure_metric(tmp_path / "m.xml", [sample], metric) + + metric.update.assert_not_called() + + def test_coco_dispatch_unknown_iou_type_does_nothing( + self, + converter, + tmp_path, + monkeypatch, + ): + from model_converter.datasets import CalibrationSample + from model_converter.metrics import CocoDetectionMAP + + img_path = tmp_path / "img.jpg" + cv2.imwrite(str(img_path), np.zeros((10, 10, 3), dtype=np.uint8)) + wrapper = MagicMock(return_value=SimpleNamespace()) + self._install_fake_model_api(monkeypatch, wrapper) + + ann = tmp_path / "ann.json" + ann.write_text('{"images":[],"annotations":[],"categories":[]}') + metric = CocoDetectionMAP(annotation_file=ann, iou_type="bbox") + metric.iou_type = "segm" # Force an unsupported branch (not wired up here) + metric.update = MagicMock() + metric.compute = MagicMock(return_value=0.0) + sample = CalibrationSample(image_path=img_path, label=0, image_id=1) + + converter._measure_metric(tmp_path / "m.xml", [sample], metric) + + metric.update.assert_not_called() + + def test_bbox_dispatch_skips_when_result_missing_fields( + self, + converter, + tmp_path, + monkeypatch, + ): + from model_converter.datasets import CalibrationSample + from model_converter.metrics import CocoDetectionMAP + + img_path = tmp_path / "img.jpg" + cv2.imwrite(str(img_path), np.zeros((10, 10, 3), dtype=np.uint8)) + # Result without bboxes / labels / scores attrs. + wrapper = MagicMock(return_value=SimpleNamespace()) + self._install_fake_model_api(monkeypatch, wrapper) + + ann = tmp_path / "ann.json" + ann.write_text('{"images":[],"annotations":[],"categories":[]}') + metric = CocoDetectionMAP(annotation_file=ann, iou_type="bbox") + metric.update = MagicMock() + metric.compute = MagicMock(return_value=0.0) + sample = CalibrationSample(image_path=img_path, label=0, image_id=1) + + converter._measure_metric(tmp_path / "m.xml", [sample], metric) + + metric.update.assert_not_called() + + def test_bbox_dispatch_falls_back_for_out_of_range_label( + self, + converter, + tmp_path, + monkeypatch, + ): + """Labels outside 0-79 fall back to idx+1 without raising.""" + from model_converter.datasets import CalibrationSample + from model_converter.metrics import CocoDetectionMAP + + img_path = tmp_path / "img.jpg" + cv2.imwrite(str(img_path), np.zeros((10, 10, 3), dtype=np.uint8)) + wrapper = MagicMock( + return_value=SimpleNamespace( + bboxes=np.array([[0, 0, 5, 5]], dtype=np.int32), + labels=np.array([99], dtype=np.int32), # out of COCO-80 range + scores=np.array([0.5], dtype=np.float32), + ), + ) + self._install_fake_model_api(monkeypatch, wrapper) + + ann = tmp_path / "ann.json" + ann.write_text('{"images":[],"annotations":[],"categories":[]}') + metric = CocoDetectionMAP(annotation_file=ann, iou_type="bbox") + metric.update = MagicMock() + metric.compute = MagicMock(return_value=0.0) + metric.reset = MagicMock() + sample = CalibrationSample(image_path=img_path, label=0, image_id=1) + + converter._measure_metric(tmp_path / "model.xml", [sample], metric) + + preds = metric.update.call_args.kwargs["predictions"] + assert preds[0]["category_id"] == 100 # fallback: 99 + 1 + + from model_converter.datasets import CalibrationSample + from model_converter.metrics import SemSegMIoU + + img_path = tmp_path / "img.jpg" + cv2.imwrite(str(img_path), np.zeros((10, 10, 3), dtype=np.uint8)) + mask_path = tmp_path / "bad_mask.png" + mask_path.write_bytes(b"not a png") # cv2.imread will return None + + wrapper = MagicMock(return_value=SimpleNamespace(resultImage=np.zeros((10, 10), dtype=np.uint8))) + self._install_fake_model_api(monkeypatch, wrapper) + + metric = SemSegMIoU(num_classes=150, ignore_index=255) + metric.update = MagicMock() + metric.compute = MagicMock(return_value=0.0) + sample = CalibrationSample(image_path=img_path, label=0, mask_path=mask_path) + + converter._measure_metric(tmp_path / "m.xml", [sample], metric) + + metric.update.assert_not_called() + + def test_semseg_resizes_prediction_when_shape_mismatch( + self, + converter, + tmp_path, + monkeypatch, + ): + """Hard predictions in a different resolution are nearest-resized to match the GT mask.""" + from model_converter.datasets import CalibrationSample + from model_converter.metrics import SemSegMIoU + + img_path = tmp_path / "img.jpg" + cv2.imwrite(str(img_path), np.zeros((10, 10, 3), dtype=np.uint8)) + mask_path = tmp_path / "mask.png" + cv2.imwrite(str(mask_path), np.array([[1, 2], [1, 2]], dtype=np.uint8)) + + # Prediction at half the GT resolution → must be resized. + pred_mask = np.array([[0]], dtype=np.uint8) + wrapper = MagicMock(return_value=SimpleNamespace(resultImage=pred_mask)) + self._install_fake_model_api(monkeypatch, wrapper) + + metric = SemSegMIoU(num_classes=150, ignore_index=255) + metric.update = MagicMock() + metric.compute = MagicMock(return_value=0.0) + sample = CalibrationSample(image_path=img_path, label=0, mask_path=mask_path) + + converter._measure_metric(tmp_path / "m.xml", [sample], metric) + + # Prediction should have been resized to GT shape (2, 2) and filled with 0. + pred_arg = metric.update.call_args.args[0] + assert pred_arg.shape == (2, 2) + np.testing.assert_array_equal(pred_arg, np.zeros((2, 2), dtype=np.uint8)) + + +class TestQuantizeModelMetricPath: + """Tests for the new metric-aware accuracy path inside quantize_model.""" + + def test_uses_measure_metric_when_non_top1_metric_provided( + self, + converter, + sample_model_config, + template_dir, + tmp_path, + monkeypatch, + ): + from model_converter.datasets import CalibrationSample + from model_converter.metrics import MultilabelMAP + + _set_template_root(monkeypatch, template_dir) + (template_dir / "README-torchvision-int8.md").write_text("# <>") + fp16_dir = tmp_path / "test_model-fp16-ov" + fp16_dir.mkdir(parents=True) + model_path = fp16_dir / "test_model_fp32.xml" + model_path.write_text("") + (fp16_dir / "test_model.xml").write_text("") + + quantized_model = MagicMock() + quantized_model.get_rt_info.return_value = SimpleNamespace(value={"model_type": "Classification"}) + fake_ov = ModuleType("openvino") + fake_ov.Core = MagicMock(return_value=SimpleNamespace(read_model=MagicMock(return_value="ov_model"))) + fake_ov.save_model = MagicMock() + fake_nncf = ModuleType("nncf") + fake_nncf.Dataset = MagicMock(side_effect=lambda generator: list(generator)) + fake_nncf.quantize = MagicMock(return_value=quantized_model) + fake_nncf.QuantizationPreset = SimpleNamespace(PERFORMANCE="performance", MIXED="mixed") + + calibration_data = [np.zeros((1, 3, 224, 224), dtype=np.float32)] + samples = [CalibrationSample(image_path=tmp_path / "img.jpg", label=0)] + metric = MultilabelMAP(num_labels=3) + converter._measure_metric = MagicMock(side_effect=[0.81, 0.80, 0.78]) + + from model_converter.reporting import AccuracyResults + + accuracy = AccuracyResults() + with patch.dict(sys.modules, {"openvino": fake_ov, "nncf": fake_nncf}): + converter.quantize_model( + model_path=model_path, + calibration_data=calibration_data, + model_config=sample_model_config, + validation_samples=samples, + metric=metric, + accuracy_results=accuracy, + ) + + assert converter._measure_metric.call_count == 3 + assert accuracy.measured is True + assert accuracy.fp32_accuracy == pytest.approx(0.81) + assert accuracy.fp16_accuracy == pytest.approx(0.80) + assert accuracy.int8_accuracy == pytest.approx(0.78) + + def test_falls_back_to_validate_model_when_top1_metric( + self, + converter, + sample_model_config, + template_dir, + tmp_path, + monkeypatch, + ): + """A TopOneAccuracy metric must use the existing validate_model path (not _measure_metric).""" + from model_converter.metrics import TopOneAccuracy + + _set_template_root(monkeypatch, template_dir) + (template_dir / "README-torchvision-int8.md").write_text("# <>") + fp16_dir = tmp_path / "test_model-fp16-ov" + fp16_dir.mkdir(parents=True) + model_path = fp16_dir / "test_model_fp32.xml" + model_path.write_text("") + + quantized_model = MagicMock() + quantized_model.get_rt_info.return_value = SimpleNamespace(value={"model_type": "Classification"}) + fake_ov = ModuleType("openvino") + fake_ov.Core = MagicMock(return_value=SimpleNamespace(read_model=MagicMock(return_value="ov_model"))) + fake_ov.save_model = MagicMock() + fake_nncf = ModuleType("nncf") + fake_nncf.Dataset = MagicMock(side_effect=lambda generator: list(generator)) + fake_nncf.quantize = MagicMock(return_value=quantized_model) + fake_nncf.QuantizationPreset = SimpleNamespace(PERFORMANCE="performance", MIXED="mixed") + + calibration_data = [np.zeros((1, 3, 224, 224), dtype=np.float32)] + validation_data = [np.zeros((1, 3, 224, 224), dtype=np.float32)] + validation_labels = [0] + converter.validate_model = MagicMock(return_value=0.9) + converter._measure_metric = MagicMock() + + with patch.dict(sys.modules, {"openvino": fake_ov, "nncf": fake_nncf}): + converter.quantize_model( + model_path=model_path, + calibration_data=calibration_data, + model_config=sample_model_config, + validation_data=validation_data, + validation_labels=validation_labels, + metric=TopOneAccuracy(), + ) + + # Old path is used; _measure_metric is never called. + converter._measure_metric.assert_not_called() + assert converter.validate_model.call_count == 2 + + def test_metric_path_logs_warning_when_fp16_missing( + self, + converter, + sample_model_config, + template_dir, + tmp_path, + monkeypatch, + caplog, + ): + """When FP16 model file is absent, metric path skips fp16 and logs a warning.""" + import logging as _logging + + from model_converter.datasets import CalibrationSample + from model_converter.metrics import MultilabelMAP + + _set_template_root(monkeypatch, template_dir) + (template_dir / "README-torchvision-int8.md").write_text("# <>") + fp16_dir = tmp_path / "test_model-fp16-ov" + fp16_dir.mkdir(parents=True) + # Only FP32 exists; the FP16 .xml is missing → warning branch. + model_path = fp16_dir / "test_model_fp32.xml" + model_path.write_text("") + + quantized_model = MagicMock() + quantized_model.get_rt_info.return_value = SimpleNamespace(value={"model_type": "Classification"}) + fake_ov = ModuleType("openvino") + fake_ov.Core = MagicMock(return_value=SimpleNamespace(read_model=MagicMock(return_value="ov_model"))) + fake_ov.save_model = MagicMock() + fake_nncf = ModuleType("nncf") + fake_nncf.Dataset = MagicMock(side_effect=lambda generator: list(generator)) + fake_nncf.quantize = MagicMock(return_value=quantized_model) + fake_nncf.QuantizationPreset = SimpleNamespace(PERFORMANCE="performance", MIXED="mixed") + + calibration_data = [np.zeros((1, 3, 224, 224), dtype=np.float32)] + samples = [CalibrationSample(image_path=tmp_path / "img.jpg", label=0)] + metric = MultilabelMAP(num_labels=3) + converter._measure_metric = MagicMock(side_effect=[0.81, 0.78]) + + with ( + patch.dict(sys.modules, {"openvino": fake_ov, "nncf": fake_nncf}), + caplog.at_level(_logging.WARNING), + ): + converter.quantize_model( + model_path=model_path, + calibration_data=calibration_data, + model_config=sample_model_config, + validation_samples=samples, + metric=metric, + ) + + # FP32 + INT8 only (FP16 was skipped); warning logged. + assert converter._measure_metric.call_count == 2 + assert "FP16 model not found" in caplog.text + + +class TestMetricForConfig: + """Tests for BaseConverter._metric_for_config (config → metric dispatch).""" + + def test_imagenet_classification_returns_top1(self, converter, tmp_path): + from model_converter.metrics import TopOneAccuracy + + metric = converter._metric_for_config( + {"dataset_type": "imagenet-1k", "model_type": "Classification"}, + dataset_path=tmp_path, + ) + assert isinstance(metric, TopOneAccuracy) + + def test_imagenet_with_multilabel_task_returns_multilabel(self, converter, tmp_path): + from model_converter.metrics import MultilabelMAP + + metric = converter._metric_for_config( + { + "dataset_type": "imagenet-1k", + "model_type": "Classification", + "getitune_task": "MULTI_LABEL_CLS", + }, + dataset_path=tmp_path, + ) + assert isinstance(metric, MultilabelMAP) + + def test_coco_detection_resolves_annotation_file(self, converter, tmp_path): + from model_converter.metrics import CocoDetectionMAP + + ann_dir = tmp_path / "annotations" + ann_dir.mkdir() + (ann_dir / "instances_val2017.json").write_text( + '{"images":[],"annotations":[],"categories":[]}', + ) + metric = converter._metric_for_config( + {"dataset_type": "coco-detection", "model_type": "YOLOX"}, + dataset_path=tmp_path, + ) + assert isinstance(metric, CocoDetectionMAP) + assert metric.annotation_file == ann_dir / "instances_val2017.json" + + def test_unknown_dataset_returns_none(self, converter, tmp_path): + assert ( + converter._metric_for_config( + {"dataset_type": "unknown", "model_type": "Classification"}, + dataset_path=tmp_path, + ) + is None + ) + + def test_returns_none_when_dataset_path_missing_for_coco(self, converter): + # Without a dataset_path, COCO cannot resolve the annotation file → None + metric = converter._metric_for_config( + {"dataset_type": "coco-detection", "model_type": "YOLOX"}, + dataset_path=None, + ) + assert metric is None + + +class TestCollectValidationSamples: + """Tests for BaseConverter._collect_validation_samples.""" + + def test_returns_empty_when_dataset_path_missing(self, converter, tmp_path): + assert converter._collect_validation_samples(tmp_path / "nope", "imagenet-1k") == [] + + def test_returns_empty_when_dataset_type_unknown(self, converter, tmp_path): + assert converter._collect_validation_samples(tmp_path, "totally-fake") == [] + + def test_returns_empty_when_dataset_path_none(self, converter): + assert converter._collect_validation_samples(None, "imagenet-1k") == [] + + def test_respects_subset_size(self, converter, tmp_path): + class_dir = tmp_path / "0" + class_dir.mkdir() + for i in range(5): + (class_dir / f"img{i}.JPEG").write_bytes(b"") + samples = converter._collect_validation_samples(tmp_path, "imagenet-1k", subset_size=3) + assert len(samples) == 3 + + def test_collects_coco_samples_with_image_id(self, converter, tmp_path): + images = tmp_path / "images" + annotations = tmp_path / "annotations" + images.mkdir() + annotations.mkdir() + (images / "000001.jpg").write_bytes(b"") + (annotations / "instances_val2017.json").write_text( + '{"images":[{"id":1,"file_name":"000001.jpg"}],"annotations":[],"categories":[]}', + ) + samples = converter._collect_validation_samples(tmp_path, "coco-detection") + assert len(samples) == 1 + assert samples[0].image_id == 1 diff --git a/model_converter/tests/unit/test_dataset_registry.py b/model_converter/tests/unit/test_dataset_registry.py new file mode 100644 index 000000000..b2c4528c9 --- /dev/null +++ b/model_converter/tests/unit/test_dataset_registry.py @@ -0,0 +1,169 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Tests for DatasetRegistry.""" + +import json + +import pytest +from model_converter.dataset_registry import DatasetRegistry + + +@pytest.fixture +def datasets_config(tmp_path): + """Create a temporary datasets configuration file.""" + config = { + "datasets": { + "imagenet-1k": str(tmp_path / "imagenet"), + "imagenet-21k": str(tmp_path / "imagenet21k"), + "coco-detection": str(tmp_path / "coco"), + }, + } + config_file = tmp_path / "datasets.json" + with config_file.open("w") as f: + json.dump(config, f) + return config_file + + +def test_load_valid_config(datasets_config): + """Test loading a valid datasets configuration.""" + registry = DatasetRegistry(datasets_config) + assert registry.list_types() == ["coco-detection", "imagenet-1k", "imagenet-21k"] + + +def test_load_missing_file(tmp_path): + """Test loading a non-existent configuration file.""" + missing_file = tmp_path / "missing.json" + with pytest.raises(FileNotFoundError, match="Dataset configuration file not found"): + DatasetRegistry(missing_file) + + +def test_load_invalid_json(tmp_path): + """Test loading invalid JSON.""" + bad_json = tmp_path / "bad.json" + bad_json.write_text("{ invalid json }") + with pytest.raises(ValueError, match="Invalid JSON"): + DatasetRegistry(bad_json) + + +def test_load_missing_datasets_key(tmp_path): + """Test loading configuration without 'datasets' key.""" + config_file = tmp_path / "config.json" + with config_file.open("w") as f: + json.dump({"wrong_key": {}}, f) + with pytest.raises(ValueError, match='Expected JSON with "datasets" key'): + DatasetRegistry(config_file) + + +def test_get_path_exists(datasets_config, tmp_path): + """Test getting path for existing dataset type.""" + registry = DatasetRegistry(datasets_config) + path = registry.get_path("imagenet-1k") + assert path == tmp_path / "imagenet" + + +def test_get_path_missing(datasets_config): + """Test getting path for non-existent dataset type.""" + registry = DatasetRegistry(datasets_config) + with pytest.raises(ValueError, match="Dataset type 'missing' not found"): + registry.get_path("missing") + + +def test_get_path_validate_exists(datasets_config, tmp_path): + """Test path validation when dataset directory exists.""" + # Create the directory + imagenet_dir = tmp_path / "imagenet" + imagenet_dir.mkdir() + + registry = DatasetRegistry(datasets_config) + path = registry.get_path("imagenet-1k", validate_exists=True) + assert path == imagenet_dir + + +def test_get_path_validate_missing(datasets_config): + """Test path validation when dataset directory doesn't exist.""" + registry = DatasetRegistry(datasets_config) + with pytest.raises(FileNotFoundError, match=r"Dataset path.*does not exist"): + registry.get_path("imagenet-1k", validate_exists=True) + + +def test_has_type(datasets_config): + """Test checking if dataset type exists.""" + registry = DatasetRegistry(datasets_config) + assert registry.has_type("imagenet-1k") + assert registry.has_type("coco-detection") + assert not registry.has_type("missing") + + +def test_list_types(datasets_config): + """Test listing all dataset types.""" + registry = DatasetRegistry(datasets_config) + types = registry.list_types() + assert types == ["coco-detection", "imagenet-1k", "imagenet-21k"] + + +def test_resolve_from_config_with_dataset_type(datasets_config, tmp_path): + """Test resolving dataset path from model config.""" + registry = DatasetRegistry(datasets_config) + config = {"model_short_name": "resnet50", "dataset_type": "imagenet-1k"} + path = registry.resolve_from_config(config) + assert path == tmp_path / "imagenet" + + +def test_resolve_from_config_no_dataset_type(datasets_config): + """Test resolving when config has no dataset_type.""" + registry = DatasetRegistry(datasets_config) + config = {"model_short_name": "yolo11"} + path = registry.resolve_from_config(config) + assert path is None + + +def test_resolve_from_config_invalid_type(datasets_config): + """Test resolving with invalid dataset_type value.""" + registry = DatasetRegistry(datasets_config) + config = {"model_short_name": "model1", "dataset_type": ""} + with pytest.raises(ValueError, match="invalid dataset_type"): + registry.resolve_from_config(config) + + +def test_resolve_from_config_unknown_type(datasets_config): + """Test resolving with unknown dataset type.""" + registry = DatasetRegistry(datasets_config) + config = {"model_short_name": "model1", "dataset_type": "unknown"} + with pytest.raises(ValueError, match="Dataset type 'unknown' not found"): + registry.resolve_from_config(config) + + +def test_load_oserror(tmp_path): + """Test error when config file cannot be read due to OS error.""" + from unittest.mock import patch + + config_file = tmp_path / "datasets.json" + config_file.write_text('{"datasets": {}}') + + with ( + patch("pathlib.Path.open", side_effect=OSError("permission denied")), + pytest.raises( + ValueError, + match="Failed to read dataset configuration file", + ), + ): + DatasetRegistry(config_file) + + +def test_load_datasets_not_dict(tmp_path): + """Test error when 'datasets' value is not a dictionary.""" + config_file = tmp_path / "datasets.json" + config_file.write_text('{"datasets": "not-a-dict"}') + with pytest.raises(ValueError, match="Invalid datasets format"): + DatasetRegistry(config_file) + + +def test_load_non_string_path(tmp_path): + """Test error when a dataset path value is not a string.""" + config_file = tmp_path / "datasets.json" + config_file.write_text('{"datasets": {"imagenet-1k": 123}}') + with pytest.raises(ValueError, match="Invalid entry in datasets configuration"): + DatasetRegistry(config_file) diff --git a/model_converter/tests/unit/test_getitune.py b/model_converter/tests/unit/test_getitune.py new file mode 100644 index 000000000..f497d6448 --- /dev/null +++ b/model_converter/tests/unit/test_getitune.py @@ -0,0 +1,954 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Tests for GetituneConverter.""" + +import json +import logging +import types +from pathlib import Path +from unittest.mock import ANY, MagicMock, patch + +import numpy as np +import pytest +from model_converter.converters.getitune import GetituneConverter +from model_converter.reporting import AccuracyResults + + +def _write_openvino_model(xml_path: Path) -> None: + """Create paired OpenVINO XML and BIN files.""" + xml_path.parent.mkdir(parents=True, exist_ok=True) + xml_path.write_text("") + xml_path.with_suffix(".bin").write_bytes(b"bin") + + +def _make_openvino_module(model_info: dict[str, str] | None = None) -> types.ModuleType: + """Create a fake openvino module for metadata extraction.""" + ov_module = types.ModuleType("openvino") + core = MagicMock() + model = MagicMock() + model.get_rt_info.return_value = MagicMock(value=model_info or {"model_type": "Classification"}) + core.read_model.return_value = model + setattr(ov_module, "Core", MagicMock(return_value=core)) + return ov_module + + +@pytest.fixture +def sample_getitune_config(): + """Sample getitune model configuration.""" + return { + "model_short_name": "dino_v2_cls", + "model_full_name": "DINOv2 Classification", + "model_library": "getitune", + "getitune_task": "MULTI_CLASS_CLS", + "getitune_recipe": "dino_v2", + "model_type": "Classification", + "license": "apache-2.0", + "license_link": "https://www.apache.org/licenses/LICENSE-2.0", + "docs": "https://example.com", + "tags": ["classification", "openvino"], + "quantize": True, + "dataset_type": "imagenet-1k", + } + + +@pytest.fixture +def training_extensions_dir(tmp_path): + """Temporary training_extensions checkout with export script.""" + repo_dir = tmp_path / "training_extensions" + repo_dir.mkdir() + (repo_dir / "export_pretrained_models.py").write_text("print('export')\n") + library_dir = repo_dir / "library" + library_dir.mkdir() + (library_dir / "pyproject.toml").write_text('[project]\nname = "getitune"\n') + return repo_dir + + +@pytest.fixture +def getitune_converter(tmp_output_dir, tmp_cache_dir, training_extensions_dir, mock_dataset_registry): + """GetituneConverter with temporary directories.""" + return GetituneConverter( + output_dir=tmp_output_dir, + cache_dir=tmp_cache_dir, + verbose=True, + dataset_registry=mock_dataset_registry, + training_extensions_dir=training_extensions_dir, + ) + + +class TestProcessModelConfig: + """Tests for GetituneConverter.process_model_config.""" + + def test_process_model_config_succeeds( + self, + getitune_converter, + sample_getitune_config, + training_extensions_dir, + tmp_path, + ): + """process_model_config exports, repackages, and quantizes a valid model.""" + export_root = tmp_path / "getitune_export_process" + expected_xml = export_root / "multi_class_cls" / "dino_v2" / "exported_model.xml" + + def fake_run(*args, **kwargs): + _write_openvino_model(expected_xml) + return MagicMock(returncode=0, stdout="ok", stderr="") + + fake_openvino = _make_openvino_module() + + with ( + patch("model_converter.converters.getitune.tempfile.mkdtemp", return_value=str(export_root)), + patch("model_converter.converters.getitune.subprocess.run", side_effect=fake_run), + patch.object(getitune_converter, "copy_readme") as mock_copy_readme, + patch.object( + getitune_converter, + "_quantize_exported_model", + return_value=AccuracyResults(), + ) as mock_quantize, + patch.dict("sys.modules", {"openvino": fake_openvino}), + ): + assert getitune_converter.process_model_config(sample_getitune_config) is True + + output_folder = getitune_converter.output_dir / "dino_v2_cls-fp16-ov" + assert (output_folder / "dino_v2_cls.xml").exists() + assert (output_folder / "dino_v2_cls.bin").exists() + assert json.loads((output_folder / "config.json").read_text()) == {"model_type": "Classification"} + mock_copy_readme.assert_called_once_with(sample_getitune_config, output_folder, variant="fp16") + mock_quantize.assert_called_once_with(sample_getitune_config) + assert not export_root.exists() + assert training_extensions_dir.exists() + + def test_process_model_config_skips_when_outputs_exist(self, getitune_converter, sample_getitune_config): + """process_model_config skips work when FP16 and INT8 models already exist.""" + model_short_name = sample_getitune_config["model_short_name"] + fp16_model = getitune_converter.output_dir / f"{model_short_name}-fp16-ov" / f"{model_short_name}.xml" + int8_model = getitune_converter.output_dir / f"{model_short_name}-int8-ov" / f"{model_short_name}.xml" + fp16_model.parent.mkdir(parents=True) + int8_model.parent.mkdir(parents=True) + fp16_model.write_text("") + int8_model.write_text("") + + with patch.object(getitune_converter, "_run_export") as mock_run_export: + assert getitune_converter.process_model_config(sample_getitune_config) is True + + mock_run_export.assert_not_called() + + def test_process_model_config_returns_false_without_training_extensions_dir( + self, + tmp_output_dir, + tmp_cache_dir, + dataset_dir, + sample_getitune_config, + caplog, + ): + """process_model_config returns False when training_extensions_dir is not provided.""" + converter = GetituneConverter( + output_dir=tmp_output_dir, + cache_dir=tmp_cache_dir, + verbose=True, + dataset_registry=dataset_dir, + training_extensions_dir=None, + ) + + with caplog.at_level(logging.ERROR): + result = converter.process_model_config(sample_getitune_config) + + assert result is False + assert "training_extensions_dir is required" in caplog.text + + def test_process_model_config_returns_false_when_training_extensions_dir_missing( + self, + tmp_output_dir, + tmp_cache_dir, + dataset_dir, + sample_getitune_config, + tmp_path, + caplog, + ): + """process_model_config returns False when training_extensions_dir is missing.""" + converter = GetituneConverter( + output_dir=tmp_output_dir, + cache_dir=tmp_cache_dir, + verbose=True, + dataset_registry=dataset_dir, + training_extensions_dir=tmp_path / "missing_training_extensions", + ) + + with caplog.at_level(logging.ERROR): + result = converter.process_model_config(sample_getitune_config) + + assert result is False + assert "training_extensions directory not found" in caplog.text + + def test_process_model_config_returns_false_when_export_fails( + self, + getitune_converter, + sample_getitune_config, + caplog, + ): + """process_model_config logs and returns False on export failures.""" + with ( + patch.object(getitune_converter, "_run_export", side_effect=RuntimeError("export failed")), + patch.object(getitune_converter, "_repackage_model") as mock_repackage, + caplog.at_level(logging.ERROR), + ): + result = getitune_converter.process_model_config(sample_getitune_config) + + assert result is False + mock_repackage.assert_not_called() + assert "export failed" in caplog.text + + def test_process_model_config_skips_quantization_when_disabled( + self, + getitune_converter, + sample_getitune_config, + tmp_path, + ): + """process_model_config does not quantize when quantize is disabled.""" + config = sample_getitune_config | {"quantize": False} + exported_model_path = tmp_path / "exported_model.xml" + + with ( + patch.object(getitune_converter, "_run_export", return_value=exported_model_path), + patch.object(getitune_converter, "_repackage_model") as mock_repackage, + patch.object(getitune_converter, "_quantize_exported_model") as mock_quantize, + ): + assert getitune_converter.process_model_config(config) is True + + mock_repackage.assert_called_once_with(config, exported_model_path) + mock_quantize.assert_not_called() + + def test_process_model_config_returns_false_when_license_is_missing( + self, + getitune_converter, + sample_getitune_config, + caplog, + ): + """process_model_config returns False when license is not defined.""" + config = dict(sample_getitune_config) + config.pop("license") + + with caplog.at_level(logging.ERROR): + result = getitune_converter.process_model_config(config) + + assert result is False + assert "must define 'license'" in caplog.text + + def test_process_model_config_returns_false_when_license_link_is_missing( + self, + getitune_converter, + sample_getitune_config, + caplog, + ): + """process_model_config returns False when license_link is not defined.""" + config = dict(sample_getitune_config) + config.pop("license_link") + + with caplog.at_level(logging.ERROR): + result = getitune_converter.process_model_config(config) + + assert result is False + assert "must define 'license_link'" in caplog.text + + def test_process_model_config_logs_description_when_present( + self, + getitune_converter, + sample_getitune_config, + training_extensions_dir, + tmp_path, + caplog, + ): + """process_model_config logs model description when config includes it.""" + config = sample_getitune_config | {"description": "A test getitune model"} + export_root = tmp_path / "getitune_export_desc" + expected_xml = export_root / "multi_class_cls" / "dino_v2" / "exported_model.xml" + + def fake_run(*args, **kwargs): + expected_xml.parent.mkdir(parents=True, exist_ok=True) + expected_xml.write_text("") + expected_xml.with_suffix(".bin").write_bytes(b"bin") + return MagicMock(returncode=0, stdout="ok", stderr="") + + fake_openvino = _make_openvino_module() + + with ( + patch("model_converter.converters.getitune.tempfile.mkdtemp", return_value=str(export_root)), + patch("model_converter.converters.getitune.subprocess.run", side_effect=fake_run), + patch.object(getitune_converter, "copy_readme"), + patch.object(getitune_converter, "_quantize_exported_model", return_value=AccuracyResults()), + patch.dict("sys.modules", {"openvino": fake_openvino}), + caplog.at_level(logging.INFO), + ): + result = getitune_converter.process_model_config(config) + + assert result is True + assert "A test getitune model" in caplog.text + + +class TestRunExport: + """Tests for GetituneConverter._run_export.""" + + def test_run_export_builds_expected_command_and_returns_standard_path( + self, + getitune_converter, + sample_getitune_config, + tmp_path, + ): + """_run_export invokes export_pretrained_models.py via uv run and returns the standard output path.""" + export_root = tmp_path / "getitune_export_standard" + expected_xml = export_root / "multi_class_cls" / "dino_v2" / "exported_model.xml" + + def fake_run(*args, **kwargs): + _write_openvino_model(expected_xml) + return MagicMock(returncode=0, stdout="ok", stderr="") + + with ( + patch("model_converter.converters.getitune.tempfile.mkdtemp", return_value=str(export_root)), + patch("model_converter.converters.getitune.subprocess.run", side_effect=fake_run) as mock_run, + ): + result = getitune_converter._run_export(sample_getitune_config) + + assert result == expected_xml + command = mock_run.call_args.args[0] + library_dir = getitune_converter.training_extensions_dir / "library" + assert command == [ + "uv", + "run", + "--project", + str(library_dir), + "--extra", + "cpu", + "python", + str(getitune_converter.training_extensions_dir / "export_pretrained_models.py"), + "--task", + "MULTI_CLASS_CLS", + "--model", + "dino_v2", + "--output-dir", + str(export_root), + "--format", + "OPENVINO", + "--precision", + "FP16", + ] + assert mock_run.call_args.kwargs == { + "cwd": str(getitune_converter.training_extensions_dir), + "capture_output": True, + "text": True, + "check": False, + } + + def test_run_export_falls_back_to_first_found_xml(self, getitune_converter, sample_getitune_config, tmp_path): + """_run_export falls back to searching for any exported XML file.""" + export_root = tmp_path / "getitune_export_fallback" + fallback_xml = export_root / "other" / "nested" / "model.xml" + + def fake_run(*args, **kwargs): + _write_openvino_model(fallback_xml) + return MagicMock(returncode=0, stdout="ok", stderr="") + + with ( + patch("model_converter.converters.getitune.tempfile.mkdtemp", return_value=str(export_root)), + patch("model_converter.converters.getitune.subprocess.run", side_effect=fake_run), + ): + result = getitune_converter._run_export(sample_getitune_config) + + assert result == fallback_xml + + def test_run_export_raises_on_non_zero_return_code(self, getitune_converter, sample_getitune_config, tmp_path): + """_run_export raises RuntimeError when the export script fails.""" + export_root = tmp_path / "getitune_export_error" + + with ( + patch("model_converter.converters.getitune.tempfile.mkdtemp", return_value=str(export_root)), + patch( + "model_converter.converters.getitune.subprocess.run", + return_value=MagicMock(returncode=1, stdout="", stderr="boom"), + ), + pytest.raises(RuntimeError, match="return code 1"), + ): + getitune_converter._run_export(sample_getitune_config) + + def test_run_export_raises_when_task_is_missing(self, getitune_converter, sample_getitune_config): + """_run_export requires getitune_task.""" + config = dict(sample_getitune_config) + config.pop("getitune_task") + + with pytest.raises(ValueError, match="must define 'getitune_task'"): + getitune_converter._run_export(config) + + def test_run_export_raises_when_recipe_is_missing(self, getitune_converter, sample_getitune_config): + """_run_export requires getitune_recipe.""" + config = dict(sample_getitune_config) + config.pop("getitune_recipe") + + with pytest.raises(ValueError, match="must define 'getitune_recipe'"): + getitune_converter._run_export(config) + + def test_run_export_raises_when_export_script_is_missing( + self, + tmp_output_dir, + tmp_cache_dir, + dataset_dir, + sample_getitune_config, + tmp_path, + ): + """_run_export raises FileNotFoundError when export_pretrained_models.py is absent.""" + training_extensions_dir = tmp_path / "training_extensions" + training_extensions_dir.mkdir() + converter = GetituneConverter( + output_dir=tmp_output_dir, + cache_dir=tmp_cache_dir, + verbose=True, + dataset_registry=dataset_dir, + training_extensions_dir=training_extensions_dir, + ) + + with pytest.raises(FileNotFoundError, match="Export script not found"): + converter._run_export(sample_getitune_config) + + def test_run_export_raises_when_library_pyproject_is_missing( + self, + tmp_output_dir, + tmp_cache_dir, + dataset_dir, + sample_getitune_config, + tmp_path, + ): + """_run_export raises FileNotFoundError when library/pyproject.toml is absent.""" + training_extensions_dir = tmp_path / "training_extensions" + training_extensions_dir.mkdir() + (training_extensions_dir / "export_pretrained_models.py").write_text("print('export')\n") + converter = GetituneConverter( + output_dir=tmp_output_dir, + cache_dir=tmp_cache_dir, + verbose=True, + dataset_registry=dataset_dir, + training_extensions_dir=training_extensions_dir, + ) + + with pytest.raises(FileNotFoundError, match="getitune library project not found"): + converter._run_export(sample_getitune_config) + + def test_run_export_raises_when_no_xml_files_are_found( + self, + getitune_converter, + sample_getitune_config, + tmp_path, + ): + """_run_export raises FileNotFoundError when the export output contains no XML files.""" + export_root = tmp_path / "getitune_export_empty" + + with ( + patch("model_converter.converters.getitune.tempfile.mkdtemp", return_value=str(export_root)), + patch( + "model_converter.converters.getitune.subprocess.run", + return_value=MagicMock(returncode=0, stdout="ok", stderr=""), + ), + pytest.raises(FileNotFoundError, match="No exported model found"), + ): + getitune_converter._run_export(sample_getitune_config) + + +class TestRepackageModel: + """Tests for GetituneConverter._repackage_model.""" + + def test_repackage_model_creates_layout_and_cleans_up( + self, + getitune_converter, + sample_getitune_config, + tmp_path, + ): + """_repackage_model copies files, writes metadata, and removes the export directory.""" + export_root = tmp_path / "getitune_export_repackage" + exported_model_path = export_root / "multi_class_cls" / "dino_v2" / "exported_model.xml" + _write_openvino_model(exported_model_path) + fake_openvino = _make_openvino_module({"model_type": "Classification", "labels": "cat dog"}) + + with ( + patch.object(getitune_converter, "copy_readme") as mock_copy_readme, + patch.dict("sys.modules", {"openvino": fake_openvino}), + ): + getitune_converter._repackage_model(sample_getitune_config, exported_model_path) + + output_folder = getitune_converter.output_dir / "dino_v2_cls-fp16-ov" + assert (output_folder / "dino_v2_cls.xml").exists() + assert (output_folder / "dino_v2_cls.bin").exists() + assert (output_folder / "dino_v2_cls_fp32.xml").exists() + assert (output_folder / "dino_v2_cls_fp32.bin").exists() + assert json.loads((output_folder / "config.json").read_text()) == { + "model_type": "Classification", + "labels": "cat dog", + } + assert (output_folder / ".gitattributes").exists() + mock_copy_readme.assert_called_once_with(sample_getitune_config, output_folder, variant="fp16") + assert not export_root.exists() + + def test_repackage_model_handles_metadata_extraction_failure( + self, + getitune_converter, + sample_getitune_config, + tmp_path, + caplog, + ): + """_repackage_model continues when openvino metadata extraction fails.""" + export_root = tmp_path / "getitune_export_meta_fail" + exported_model_path = export_root / "multi_class_cls" / "dino_v2" / "exported_model.xml" + _write_openvino_model(exported_model_path) + + failing_ov = types.ModuleType("openvino") + core = MagicMock() + model = MagicMock() + model.get_rt_info.side_effect = RuntimeError("No rt_info found") + core.read_model.return_value = model + setattr(failing_ov, "Core", MagicMock(return_value=core)) + + with ( + patch.object(getitune_converter, "copy_readme"), + patch.dict("sys.modules", {"openvino": failing_ov}), + caplog.at_level(logging.WARNING), + ): + getitune_converter._repackage_model(sample_getitune_config, exported_model_path) + + output_folder = getitune_converter.output_dir / "dino_v2_cls-fp16-ov" + assert (output_folder / "dino_v2_cls.xml").exists() + assert "Could not extract model_info metadata" in caplog.text + assert not (output_folder / "config.json").exists() + + def test_repackage_model_handles_non_matching_temp_path( + self, + getitune_converter, + sample_getitune_config, + tmp_path, + ): + """_repackage_model handles exported path without getitune_export_ prefix.""" + export_dir = tmp_path / "some_other_dir" / "sub" + exported_model_path = export_dir / "exported_model.xml" + _write_openvino_model(exported_model_path) + + fake_openvino = _make_openvino_module() + + with ( + patch.object(getitune_converter, "copy_readme"), + patch.dict("sys.modules", {"openvino": fake_openvino}), + ): + getitune_converter._repackage_model(sample_getitune_config, exported_model_path) + + output_folder = getitune_converter.output_dir / "dino_v2_cls-fp16-ov" + assert (output_folder / "dino_v2_cls.xml").exists() + # The temp directory should still exist since it doesn't match the pattern + assert export_dir.exists() + + +class TestQuantizeExportedModel: + """Tests for GetituneConverter._quantize_exported_model.""" + + def test_quantize_exported_model_uses_fp32_model_and_cleans_up( + self, + getitune_converter, + sample_getitune_config, + ): + """_quantize_exported_model uses the FP32 copy and removes it afterward.""" + output_folder = getitune_converter.output_dir / "dino_v2_cls-fp16-ov" + fp32_xml = output_folder / "dino_v2_cls_fp32.xml" + fp32_bin = output_folder / "dino_v2_cls_fp32.bin" + _write_openvino_model(fp32_xml) + calibration_data = [np.zeros((1, 3, 224, 224), dtype=np.float32)] + + with ( + patch.object( + getitune_converter, + "_read_preprocessing_from_model", + return_value=([1, 3, 224, 224], "123.675 116.28 103.53", "58.395 57.12 57.375", True), + ) as mock_read_preproc, + patch.object( + getitune_converter, + "create_calibration_dataset", + return_value=(calibration_data, []), + ) as mock_dataset, + patch.object(getitune_converter, "quantize_model") as mock_quantize, + ): + getitune_converter._quantize_exported_model(sample_getitune_config) + + mock_read_preproc.assert_called_once() + mock_dataset.assert_called_once_with( + input_shape=[1, 3, 224, 224], + mean_values="123.675 116.28 103.53", + scale_values="58.395 57.12 57.375", + reverse_input_channels=True, + subset_size=500, + return_labels=True, + dataset_path=ANY, + dataset_type="imagenet-1k", + ) + mock_quantize.assert_called_once_with( + model_path=fp32_xml, + calibration_data=calibration_data, + model_config=sample_getitune_config, + preset="mixed", + validation_data=None, + validation_labels=None, + validation_samples=None, + metric=ANY, + accuracy_results=ANY, + ) + assert not fp32_xml.exists() + assert not fp32_bin.exists() + + def test_quantize_exported_model_falls_back_to_fp16_when_fp32_is_missing( + self, + getitune_converter, + sample_getitune_config, + ): + """_quantize_exported_model falls back to the FP16 model when no FP32 copy exists.""" + output_folder = getitune_converter.output_dir / "dino_v2_cls-fp16-ov" + fp16_xml = output_folder / "dino_v2_cls.xml" + _write_openvino_model(fp16_xml) + calibration_data = [np.zeros((1, 3, 224, 224), dtype=np.float32)] + + with ( + patch.object( + getitune_converter, + "_read_preprocessing_from_model", + return_value=([1, 3, 224, 224], "0 0 0", "1 1 1", True), + ), + patch.object( + getitune_converter, + "create_calibration_dataset", + return_value=(calibration_data, []), + ), + patch.object(getitune_converter, "quantize_model") as mock_quantize, + ): + getitune_converter._quantize_exported_model(sample_getitune_config) + + mock_quantize.assert_called_once_with( + model_path=fp16_xml, + calibration_data=calibration_data, + model_config=sample_getitune_config, + preset="mixed", + validation_data=None, + validation_labels=None, + validation_samples=None, + metric=ANY, + accuracy_results=ANY, + ) + + def test_quantize_exported_model_handles_cleanup_oserror( + self, + getitune_converter, + sample_getitune_config, + caplog, + ): + """_quantize_exported_model logs a warning when FP32 file cleanup fails.""" + output_folder = getitune_converter.output_dir / "dino_v2_cls-fp16-ov" + fp32_xml = output_folder / "dino_v2_cls_fp32.xml" + _write_openvino_model(fp32_xml) + calibration_data = [np.zeros((1, 3, 224, 224), dtype=np.float32)] + + with ( + patch.object( + getitune_converter, + "_read_preprocessing_from_model", + return_value=([1, 3, 224, 224], "0 0 0", "1 1 1", True), + ), + patch.object( + getitune_converter, + "create_calibration_dataset", + return_value=(calibration_data, []), + ), + patch.object(getitune_converter, "quantize_model"), + patch("pathlib.Path.unlink", side_effect=OSError("permission denied")), + caplog.at_level(logging.WARNING), + ): + getitune_converter._quantize_exported_model(sample_getitune_config) + + assert "Failed to remove temporary FP32 files" in caplog.text + + def test_quantize_exported_model_measures_accuracy_for_imagenet1k( + self, + getitune_converter, + sample_getitune_config, + ): + """Accuracy is measured for MULTI_CLASS_CLS IMAGENET1K_V1 models.""" + config = {**sample_getitune_config, "labels": "IMAGENET1K_V1"} + output_folder = getitune_converter.output_dir / "dino_v2_cls-fp16-ov" + fp32_xml = output_folder / "dino_v2_cls_fp32.xml" + _write_openvino_model(fp32_xml) + calibration_data = [np.zeros((1, 3, 224, 224), dtype=np.float32)] + validation_labels = [7] + + with ( + patch.object( + getitune_converter, + "_read_preprocessing_from_model", + return_value=([1, 3, 224, 224], "0 0 0", "1 1 1", True), + ), + patch.object( + getitune_converter, + "create_calibration_dataset", + return_value=(calibration_data, validation_labels), + ) as mock_dataset, + patch.object(getitune_converter, "quantize_model") as mock_quantize, + ): + getitune_converter._quantize_exported_model(config) + + mock_dataset.assert_called_once_with( + input_shape=[1, 3, 224, 224], + mean_values="0 0 0", + scale_values="1 1 1", + reverse_input_channels=True, + subset_size=500, + return_labels=True, + dataset_path=ANY, + dataset_type="imagenet-1k", + ) + mock_quantize.assert_called_once_with( + model_path=fp32_xml, + calibration_data=calibration_data, + model_config=config, + preset="mixed", + validation_data=calibration_data, + validation_labels=validation_labels, + validation_samples=None, + metric=ANY, + accuracy_results=ANY, + ) + + def test_quantize_exported_model_skips_accuracy_for_non_classification( + self, + getitune_converter, + sample_getitune_config, + ): + """Accuracy is not measured for non-classification model types.""" + config = { + **sample_getitune_config, + "getitune_task": "DETECTION", + "model_type": "SSD", + "dataset_type": "coco-detection", + "labels": "IMAGENET1K_V1", + } + output_folder = getitune_converter.output_dir / "dino_v2_cls-fp16-ov" + fp32_xml = output_folder / "dino_v2_cls_fp32.xml" + _write_openvino_model(fp32_xml) + calibration_data = [np.zeros((1, 3, 224, 224), dtype=np.float32)] + + with ( + patch.object( + getitune_converter, + "_read_preprocessing_from_model", + return_value=([1, 3, 224, 224], "0 0 0", "1 1 1", True), + ), + patch.object( + getitune_converter, + "create_calibration_dataset", + return_value=(calibration_data, []), + ) as mock_dataset, + patch.object(getitune_converter, "quantize_model") as mock_quantize, + ): + getitune_converter._quantize_exported_model(config) + + assert mock_dataset.call_args.kwargs["return_labels"] is False + mock_quantize.assert_called_once_with( + model_path=fp32_xml, + calibration_data=calibration_data, + model_config=config, + preset="mixed", + validation_data=None, + validation_labels=None, + validation_samples=None, + metric=ANY, + accuracy_results=ANY, + ) + + def test_quantize_exported_model_honours_measure_accuracy_false( + self, + tmp_output_dir, + tmp_cache_dir, + training_extensions_dir, + mock_dataset_registry, + sample_getitune_config, + ): + """measure_accuracy=False on the converter disables accuracy measurement entirely.""" + converter = GetituneConverter( + output_dir=tmp_output_dir, + cache_dir=tmp_cache_dir, + verbose=True, + dataset_registry=mock_dataset_registry, + training_extensions_dir=training_extensions_dir, + measure_accuracy=False, + ) + config = {**sample_getitune_config, "labels": "IMAGENET1K_V1"} + output_folder = converter.output_dir / "dino_v2_cls-fp16-ov" + fp32_xml = output_folder / "dino_v2_cls_fp32.xml" + _write_openvino_model(fp32_xml) + calibration_data = [np.zeros((1, 3, 224, 224), dtype=np.float32)] + + with ( + patch.object( + converter, + "_read_preprocessing_from_model", + return_value=([1, 3, 224, 224], "0 0 0", "1 1 1", True), + ), + patch.object( + converter, + "create_calibration_dataset", + return_value=(calibration_data, []), + ) as mock_dataset, + patch.object(converter, "quantize_model"), + ): + result = converter._quantize_exported_model(config) + + assert mock_dataset.call_args.kwargs["return_labels"] is False + assert result.metric_name is None + + +class TestApplyConfigLabels: + """Tests for GetituneConverter._apply_config_labels.""" + + def test_no_labels_in_config_is_a_noop(self, getitune_converter, sample_getitune_config, tmp_path): + """Without a labels config, no OpenVINO work is attempted.""" + with patch.object(getitune_converter, "get_labels") as mock_get_labels: + getitune_converter._apply_config_labels(sample_getitune_config, tmp_path / "model.xml") + mock_get_labels.assert_not_called() + + def test_unresolved_labels_logs_warning(self, getitune_converter, sample_getitune_config, tmp_path, caplog): + """An unknown label set logs a warning and skips rewriting.""" + config = {**sample_getitune_config, "labels": "UNKNOWN_SET"} + with ( + patch.object(getitune_converter, "get_labels", return_value=None), + caplog.at_level(logging.WARNING), + ): + getitune_converter._apply_config_labels(config, tmp_path / "model.xml") + assert "Could not load labels for: UNKNOWN_SET" in caplog.text + + def test_rewrites_labels_into_existing_models(self, getitune_converter, sample_getitune_config, tmp_path): + """Labels are written to rt_info and the models are re-saved. + + The save must go to a temporary path that is then moved over the + original, never writing back to the file currently being read (which + OpenVINO memory-maps and would corrupt). + """ + config = {**sample_getitune_config, "labels": "IMAGENET1K_V1"} + fp16_xml = tmp_path / "dino_v2_cls.xml" + fp32_xml = tmp_path / "dino_v2_cls_fp32.xml" + _write_openvino_model(fp16_xml) + _write_openvino_model(fp32_xml) + missing_xml = tmp_path / "missing.xml" + + fake_ov = types.ModuleType("openvino") + core = MagicMock() + model = MagicMock() + core.read_model.return_value = model + setattr(fake_ov, "Core", MagicMock(return_value=core)) + + def fake_save_model(_model, path, compress_to_fp16): + # Never overwrite the source model directly. + assert path != fp16_xml + assert path != fp32_xml + _write_openvino_model(Path(path)) + + save_model = MagicMock(side_effect=fake_save_model) + setattr(fake_ov, "save_model", save_model) + + with ( + patch.object(getitune_converter, "get_labels", return_value="cat dog"), + patch.dict("sys.modules", {"openvino": fake_ov}), + ): + getitune_converter._apply_config_labels(config, fp16_xml, fp32_xml, missing_xml) + + model.set_rt_info.assert_called_with("cat dog", ["model_info", "labels"]) + assert save_model.call_count == 2 + assert save_model.call_args_list[0].kwargs["compress_to_fp16"] is True + assert save_model.call_args_list[1].kwargs["compress_to_fp16"] is False + # Final models exist and the temporary files were moved away. + assert fp16_xml.exists() + assert fp32_xml.exists() + assert not (tmp_path / "dino_v2_cls_labeled_tmp.xml").exists() + assert not (tmp_path / "dino_v2_cls_fp32_labeled_tmp.xml").exists() + + def test_handles_openvino_runtime_error(self, getitune_converter, sample_getitune_config, tmp_path, caplog): + """A failure while rewriting labels is logged and swallowed.""" + config = {**sample_getitune_config, "labels": "IMAGENET1K_V1"} + model_xml = tmp_path / "dino_v2_cls.xml" + _write_openvino_model(model_xml) + + fake_ov = types.ModuleType("openvino") + core = MagicMock() + core.read_model.side_effect = RuntimeError("read failed") + setattr(fake_ov, "Core", MagicMock(return_value=core)) + + with ( + patch.object(getitune_converter, "get_labels", return_value="cat dog"), + patch.dict("sys.modules", {"openvino": fake_ov}), + caplog.at_level(logging.WARNING), + ): + getitune_converter._apply_config_labels(config, model_xml) + + assert "Could not apply labels to exported model" in caplog.text + + +class TestReadPreprocessingFromModel: + """Tests for GetituneConverter._read_preprocessing_from_model.""" + + def test_reads_all_params_from_rt_info(self, tmp_path): + """_read_preprocessing_from_model extracts params from model rt_info.""" + model_path = tmp_path / "model.xml" + model_path.write_text("") + + mock_model = MagicMock() + mock_model.input.return_value = MagicMock(shape=[1, 3, 640, 640]) + + def fake_get_rt_info(path): + values = { + "mean_values": "123.675 116.28 103.53", + "scale_values": "58.395 57.12 57.375", + "reverse_input_channels": "True", + } + key = path[1] + if key in values: + result = MagicMock() + result.astype.return_value = values[key] + return result + msg = "Cannot get runtime attribute. Path to runtime attribute is incorrect." + raise RuntimeError(msg) + + mock_model.get_rt_info.side_effect = fake_get_rt_info + + mock_ov = MagicMock() + mock_ov.Core.return_value.read_model.return_value = mock_model + + input_shape, mean_values, scale_values, reverse = GetituneConverter._read_preprocessing_from_model( + mock_ov, + model_path, + ) + + assert input_shape == [1, 3, 640, 640] + assert mean_values == "123.675 116.28 103.53" + assert scale_values == "58.395 57.12 57.375" + assert reverse is True + + def test_falls_back_to_defaults_when_rt_info_missing(self, tmp_path): + """_read_preprocessing_from_model uses defaults when rt_info keys are absent.""" + model_path = tmp_path / "model.xml" + model_path.write_text("") + + mock_model = MagicMock() + mock_model.input.return_value = MagicMock(shape=[1, 3, 224, 224]) + mock_model.get_rt_info.side_effect = RuntimeError( + "Cannot get runtime attribute. Path to runtime attribute is incorrect.", + ) + + mock_ov = MagicMock() + mock_ov.Core.return_value.read_model.return_value = mock_model + + input_shape, mean_values, scale_values, reverse = GetituneConverter._read_preprocessing_from_model( + mock_ov, + model_path, + ) + + assert input_shape == [1, 3, 224, 224] + assert mean_values == "0 0 0" + assert scale_values == "1 1 1" + assert reverse is True diff --git a/model_converter/tests/unit/test_registry.py b/model_converter/tests/unit/test_registry.py new file mode 100644 index 000000000..a5e477b47 --- /dev/null +++ b/model_converter/tests/unit/test_registry.py @@ -0,0 +1,579 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Tests for converter registry and individual converter classes.""" + +import json +import logging +import sys +import types +from unittest.mock import ANY, MagicMock, call, patch + +import pytest +import torch.nn as nn +from model_converter.converters import ( + CONVERTER_REGISTRY, + BaseConverter, + GetituneConverter, + PyTorchConverter, + TimmConverter, + TorchvisionConverter, + YoloConverter, +) +from model_converter.reporting import AccuracyResults + + +class TestConverterRegistry: + """Tests for the converter registry.""" + + def test_contains_expected_keys(self): + """CONVERTER_REGISTRY exposes all supported converter names.""" + assert set(CONVERTER_REGISTRY) == {"torchvision", "timm", "yolo", "getitune"} + + def test_maps_expected_classes(self): + """CONVERTER_REGISTRY maps each name to the expected class.""" + assert { + "torchvision": TorchvisionConverter, + "timm": TimmConverter, + "yolo": YoloConverter, + "getitune": GetituneConverter, + } == CONVERTER_REGISTRY + + @pytest.mark.parametrize( + ("name", "expected_cls"), + [ + ("BaseConverter", BaseConverter), + ("PyTorchConverter", PyTorchConverter), + ("TorchvisionConverter", TorchvisionConverter), + ("TimmConverter", TimmConverter), + ("YoloConverter", YoloConverter), + ("GetituneConverter", GetituneConverter), + ], + ) + def test_package_exports_converter_classes(self, name, expected_cls): + """model_converter.converters re-exports the converter classes.""" + import model_converter.converters as converters + + assert getattr(converters, name) is expected_cls + + +class TestTorchvisionConverter: + """Tests for TorchvisionConverter-specific behavior.""" + + def test_process_model_config_downloads_weights_and_loads_model( + self, + converter, + sample_model_config, + tmp_path, + mock_torch_model, + ): + """process_model_config downloads weights and builds a model instance.""" + weights_path = tmp_path / "weights.pth" + checkpoint = {"state_dict": {"layer.weight": "weights"}} + model_class = MagicMock() + fp16_path = converter.output_dir / "test_model-fp16-ov" / "test_model.xml" + fp32_path = converter.output_dir / "test_model-fp16-ov" / "test_model_fp32.xml" + + with ( + patch.object(converter._url_downloader, "download", return_value=weights_path) as mock_download, + patch.object(converter, "load_model_class", return_value=model_class) as mock_load_class, + patch.object(converter, "load_checkpoint", return_value=checkpoint) as mock_load_checkpoint, + patch.object(converter, "create_model", return_value=mock_torch_model) as mock_create_model, + patch.object(converter, "_build_metadata", return_value={("model_info", "model_type"): "Classification"}), + patch.object(converter, "export_to_openvino", return_value=(fp16_path, fp32_path)) as mock_export, + ): + result = converter.process_model_config(sample_model_config) + + assert result is True + mock_download.assert_called_once_with(url=sample_model_config["weights_url"]) + mock_load_class.assert_called_once_with(sample_model_config["model_class_name"]) + mock_load_checkpoint.assert_called_once_with(weights_path) + mock_create_model.assert_called_once_with(model_class, checkpoint, None) + assert mock_export.call_args.kwargs["model"] is mock_torch_model + assert mock_export.call_args.kwargs["output_path"] == converter.output_dir / "test_model" + + def test_process_model_config_skips_when_outputs_are_cached(self, converter, sample_model_config): + """process_model_config skips work when both output models already exist.""" + fp16_dir = converter.output_dir / "test_model-fp16-ov" + int8_dir = converter.output_dir / "test_model-int8-ov" + fp16_dir.mkdir(parents=True) + int8_dir.mkdir(parents=True) + (fp16_dir / "test_model.xml").write_text("") + (int8_dir / "test_model.xml").write_text("") + + with patch.object(converter._url_downloader, "download") as mock_download: + result = converter.process_model_config(sample_model_config) + + assert result is True + mock_download.assert_not_called() + + +class TestTimmConverter: + """Tests for TimmConverter-specific behavior.""" + + def test_load_huggingface_model_uses_timm_create_model(self, timm_converter): + """load_huggingface_model builds the expected timm hub reference.""" + mock_model = MagicMock() + mock_model.eval.return_value = mock_model + timm_module = types.ModuleType("timm") + timm_module.create_model = MagicMock(return_value=mock_model) + + with patch.dict(sys.modules, {"timm": timm_module}): + result = timm_converter.load_huggingface_model( + repo_id="timm/resnet50.a1_in1k", + revision="abc123", + model_library="timm", + model_params={"num_classes": 10}, + ) + + assert result is mock_model + timm_module.create_model.assert_called_once_with( + "hf-hub:timm/resnet50.a1_in1k@abc123", + pretrained=True, + cache_dir=timm_converter.cache_dir, + num_classes=10, + ) + mock_model.eval.assert_called_once_with() + + def test_process_model_config_uses_loader_and_export_pipeline( + self, + timm_converter, + sample_timm_config, + mock_torch_model, + ): + """process_model_config runs the timm conversion workflow.""" + fp16_path = timm_converter.output_dir / "test_timm_model-fp16-ov" / "test_timm_model.xml" + fp32_path = timm_converter.output_dir / "test_timm_model-fp16-ov" / "test_timm_model_fp32.xml" + + with ( + patch.object(timm_converter, "load_huggingface_model", return_value=mock_torch_model) as mock_load, + patch.object( + timm_converter, + "_build_metadata", + return_value={("model_info", "model_type"): "Classification"}, + ) as mock_metadata, + patch.object(timm_converter, "export_to_openvino", return_value=(fp16_path, fp32_path)) as mock_export, + ): + result = timm_converter.process_model_config(sample_timm_config) + + assert result is True + mock_load.assert_called_once_with( + repo_id=sample_timm_config["huggingface_repo"], + revision=sample_timm_config["huggingface_revision"], + model_library="timm", + model_params=None, + ) + mock_metadata.assert_called_once_with(sample_timm_config) + assert mock_export.call_args.kwargs["model"] is mock_torch_model + assert mock_export.call_args.kwargs["output_path"] == timm_converter.output_dir / "test_timm_model" + + def test_process_model_config_skips_when_outputs_exist(self, timm_converter, sample_timm_config): + """process_model_config skips work when both FP16 and INT8 models already exist.""" + model_short_name = sample_timm_config["model_short_name"] + fp16_dir = timm_converter.output_dir / f"{model_short_name}-fp16-ov" + int8_dir = timm_converter.output_dir / f"{model_short_name}-int8-ov" + fp16_dir.mkdir(parents=True) + int8_dir.mkdir(parents=True) + (fp16_dir / f"{model_short_name}.xml").write_text("") + (int8_dir / f"{model_short_name}.xml").write_text("") + + with patch.object(timm_converter, "load_huggingface_model") as mock_load: + result = timm_converter.process_model_config(sample_timm_config) + + assert result is True + mock_load.assert_not_called() + + def test_process_model_config_returns_false_when_license_is_missing( + self, + timm_converter, + sample_timm_config, + caplog, + ): + """process_model_config returns False when license is not defined.""" + config = dict(sample_timm_config) + config.pop("license") + + with caplog.at_level(logging.ERROR): + result = timm_converter.process_model_config(config) + + assert result is False + assert "must define 'license'" in caplog.text + + def test_process_model_config_returns_false_when_license_link_is_missing( + self, + timm_converter, + sample_timm_config, + caplog, + ): + """process_model_config returns False when license_link is not defined.""" + config = dict(sample_timm_config) + config.pop("license_link") + + with caplog.at_level(logging.ERROR): + result = timm_converter.process_model_config(config) + + assert result is False + assert "must define 'license_link'" in caplog.text + + def test_process_model_config_runs_quantization_when_dataset_available( + self, + tmp_path, + sample_timm_config, + mock_torch_model, + ): + """process_model_config runs quantization when dataset path exists.""" + dataset_dir = tmp_path / "dataset" + dataset_dir.mkdir() + timm_converter = TimmConverter( + output_dir=tmp_path / "output", + cache_dir=tmp_path / "cache", + dataset_registry=dataset_dir, + verbose=True, + ) + fp16_path = timm_converter.output_dir / "test_timm_model-fp16-ov" / "test_timm_model.xml" + fp32_path = timm_converter.output_dir / "test_timm_model-fp16-ov" / "test_timm_model_fp32.xml" + + with ( + patch.object(timm_converter, "load_huggingface_model", return_value=mock_torch_model), + patch.object( + timm_converter, + "_build_metadata", + return_value={("model_info", "model_type"): "Classification"}, + ), + patch.object(timm_converter, "export_to_openvino", return_value=(fp16_path, fp32_path)), + patch.object(timm_converter, "_quantize_and_cleanup", return_value=AccuracyResults()) as mock_quantize, + ): + result = timm_converter.process_model_config(sample_timm_config) + + assert result is True + mock_quantize.assert_called_once_with( + sample_timm_config, + fp32_path, + model_type=sample_timm_config.get("model_type", ""), + input_shape=sample_timm_config["input_shape"], + mean_values=sample_timm_config["mean_values"], + scale_values=sample_timm_config["scale_values"], + reverse_input_channels=sample_timm_config["reverse_input_channels"], + torch_model=ANY, + ) + + +class TestYoloConverter: + """Tests for YoloConverter-specific behavior.""" + + def test_process_model_config_exports_fp16_and_int8(self, tmp_path, monkeypatch): + """process_model_config exports both YOLO variants and repackages them.""" + converter = YoloConverter(output_dir=tmp_path / "output", cache_dir=tmp_path / "cache") + config = {"model_short_name": "YOLO11n", "yolo_version": "yolo11n"} + mock_model = MagicMock() + yolo_module = types.ModuleType("ultralytics") + yolo_module.YOLO = MagicMock(return_value=mock_model) + export_calls: list[dict[str, object]] = [] + + def export_side_effect(**kwargs): + export_calls.append(kwargs) + suffix = "_int8_openvino_model" if kwargs.get("int8") else "_openvino_model" + export_dir = converter.cache_dir / f"yolo11n{suffix}" + export_dir.mkdir(parents=True) + (export_dir / "yolo11n.xml").write_text("") + + mock_model.export.side_effect = export_side_effect + monkeypatch.chdir(tmp_path) + + with ( + patch.dict(sys.modules, {"ultralytics": yolo_module}), + patch.object(converter, "_update_model_type_in_xml") as mock_update, + patch.object(converter, "_copy_yolo_readme") as mock_copy_readme, + ): + result = converter.process_model_config(config) + + assert result is True + yolo_module.YOLO.assert_called_once_with(str(converter.cache_dir / "yolo11n.pt")) + assert export_calls == [ + {"format": "openvino", "half": True}, + {"format": "openvino", "int8": True, "data": "coco128.yaml"}, + ] + assert (converter.output_dir / "YOLO11n-fp16-ov" / "yolo11n.xml").exists() + assert (converter.output_dir / "YOLO11n-int8-ov" / "yolo11n.xml").exists() + assert mock_update.call_count == 2 + mock_copy_readme.assert_has_calls( + [ + call("README-yolo-fp16.md", converter.output_dir / "YOLO11n-fp16-ov", "n"), + call("README-yolo-int8.md", converter.output_dir / "YOLO11n-int8-ov", "n"), + ], + ) + + def test_process_model_config_returns_false_on_export_failure(self, tmp_path, monkeypatch): + """process_model_config handles YOLO export errors gracefully.""" + converter = YoloConverter(output_dir=tmp_path / "output", cache_dir=tmp_path / "cache") + yolo_module = types.ModuleType("ultralytics") + failing_model = MagicMock() + failing_model.export.side_effect = RuntimeError("export failed") + yolo_module.YOLO = MagicMock(return_value=failing_model) + monkeypatch.chdir(tmp_path) + + with patch.dict(sys.modules, {"ultralytics": yolo_module}): + result = converter.process_model_config({"model_short_name": "YOLO11n", "yolo_version": "yolo11n"}) + + assert result is False + + def test_process_model_config_skips_when_outputs_exist(self, tmp_path): + """process_model_config skips work when both FP16 and INT8 YOLO models already exist.""" + converter = YoloConverter(output_dir=tmp_path / "output", cache_dir=tmp_path / "cache") + config = {"model_short_name": "YOLO11n", "yolo_version": "yolo11n"} + + fp16_dir = converter.output_dir / "YOLO11n-fp16-ov" + int8_dir = converter.output_dir / "YOLO11n-int8-ov" + fp16_dir.mkdir(parents=True) + int8_dir.mkdir(parents=True) + (fp16_dir / "yolo11n.xml").write_text("") + (int8_dir / "yolo11n.xml").write_text("") + + result = converter.process_model_config(config) + + assert result is True + + def test_process_model_config_replaces_existing_int8_folder(self, tmp_path, monkeypatch): + """process_model_config removes pre-existing INT8 folder before renaming new output.""" + converter = YoloConverter(output_dir=tmp_path / "output", cache_dir=tmp_path / "cache") + config = {"model_short_name": "YOLO11n", "yolo_version": "yolo11n"} + mock_model = MagicMock() + yolo_module = types.ModuleType("ultralytics") + yolo_module.YOLO = MagicMock(return_value=mock_model) + + # Pre-create the INT8 folder to trigger the rmtree branch + int8_dir = converter.output_dir / "YOLO11n-int8-ov" + int8_dir.mkdir(parents=True) + (int8_dir / "old_file.txt").write_text("old") + + def export_side_effect(**kwargs): + suffix = "_int8_openvino_model" if kwargs.get("int8") else "_openvino_model" + export_dir = converter.cache_dir / f"yolo11n{suffix}" + export_dir.mkdir(parents=True) + (export_dir / "yolo11n.xml").write_text("") + + mock_model.export.side_effect = export_side_effect + monkeypatch.chdir(tmp_path) + + with ( + patch.dict(sys.modules, {"ultralytics": yolo_module}), + patch.object(converter, "_update_model_type_in_xml"), + patch.object(converter, "_copy_yolo_readme"), + ): + result = converter.process_model_config(config) + + assert result is True + assert (converter.output_dir / "YOLO11n-int8-ov" / "yolo11n.xml").exists() + assert not (converter.output_dir / "YOLO11n-int8-ov" / "old_file.txt").exists() + + def test_copy_yolo_readme_replaces_size_placeholder(self, tmp_path, monkeypatch): + """_copy_yolo_readme reads template and replaces <>.""" + import model_converter.converters.yolo as yolo_module + + converter = YoloConverter(output_dir=tmp_path / "output", cache_dir=tmp_path / "cache") + template_dir = tmp_path / "templates" + template_dir.mkdir() + (template_dir / "README-yolo-fp16.md").write_text("# YOLO <> model") + monkeypatch.setattr(yolo_module, "__file__", str(template_dir.parent / "converters" / "yolo.py")) + + dest_dir = tmp_path / "YOLO11n-fp16-ov" + dest_dir.mkdir() + + converter._copy_yolo_readme("README-yolo-fp16.md", dest_dir, "n") + + assert (dest_dir / "README.md").read_text() == "# YOLO n model" + + def test_copy_yolo_readme_warns_when_template_missing(self, tmp_path, monkeypatch, caplog): + """_copy_yolo_readme logs a warning when template file doesn't exist.""" + import model_converter.converters.yolo as yolo_module + + converter = YoloConverter(output_dir=tmp_path / "output", cache_dir=tmp_path / "cache") + empty_templates = tmp_path / "templates" + empty_templates.mkdir() + monkeypatch.setattr(yolo_module, "__file__", str(empty_templates.parent / "converters" / "yolo.py")) + + dest_dir = tmp_path / "YOLO11n-fp16-ov" + dest_dir.mkdir() + + converter._copy_yolo_readme("README-yolo-fp16.md", dest_dir, "n") + + assert "YOLO README template not found" in caplog.text + assert not (dest_dir / "README.md").exists() + + def test_update_model_type_in_xml_modifies_element_and_writes_config(self, tmp_path): + """_update_model_type_in_xml updates model_type and writes config.json.""" + converter = YoloConverter(output_dir=tmp_path / "output", cache_dir=tmp_path / "cache") + xml_path = tmp_path / "model.xml" + xml_path.write_text( + '\n' + "\n" + " \n" + " \n" + ' \n' + ' \n' + " \n" + " \n" + "\n", + ) + + converter._update_model_type_in_xml(xml_path, "YOLO11") + + config_json = json.loads((tmp_path / "config.json").read_text()) + assert config_json["model_type"] == "YOLO11" + assert config_json["labels"] == "person car" + + def test_update_model_type_in_xml_handles_parse_error(self, tmp_path, caplog): + """_update_model_type_in_xml logs a warning when XML is malformed.""" + converter = YoloConverter(output_dir=tmp_path / "output", cache_dir=tmp_path / "cache") + xml_path = tmp_path / "bad.xml" + xml_path.write_text("not valid xml <<<>>>") + + converter._update_model_type_in_xml(xml_path, "YOLO11") + + assert "Failed to update" in caplog.text + + +class TestPyTorchConverterSharedLogic: + """Tests for logic shared by PyTorch-based converters.""" + + def test_export_to_openvino_converts_and_saves_models( + self, + converter, + sample_model_config, + mock_torch_model, + mock_ov_model, + ): + """export_to_openvino converts the model and saves FP32 and FP16 artifacts.""" + dummy_input = object() + openvino_module = types.ModuleType("openvino") + openvino_module.convert_model = MagicMock(return_value=mock_ov_model) + openvino_module.save_model = MagicMock() + + with ( + patch.dict(sys.modules, {"openvino": openvino_module}), + 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.object(converter, "copy_readme") as mock_copy_readme, + ): + fp16_path, fp32_path = converter.export_to_openvino( + model=mock_torch_model, + input_shape=[1, 3, 224, 224], + output_path=converter.output_dir / "test_model", + model_config=sample_model_config, + input_names=["input"], + output_names=["result"], + metadata={("model_info", "model_type"): "Classification"}, + ) + + assert fp16_path == converter.output_dir / "test_model-fp16-ov" / "test_model.xml" + assert fp32_path == converter.output_dir / "test_model-fp16-ov" / "test_model_fp32.xml" + mock_prepare.assert_called_once_with(mock_torch_model, sample_model_config) + mock_example_input.assert_called_once_with([1, 3, 224, 224], sample_model_config) + openvino_module.convert_model.assert_called_once_with(mock_torch_model, example_input=dummy_input) + mock_ov_model.reshape.assert_called_once_with({"input": [1, 3, 224, 224]}) + mock_postprocess.assert_called_once_with( + mock_ov_model, + input_names=["input"], + output_names=["result"], + metadata={("model_info", "model_type"): "Classification"}, + ) + openvino_module.save_model.assert_has_calls( + [ + call(mock_ov_model, fp32_path, compress_to_fp16=False), + call(mock_ov_model, fp16_path, compress_to_fp16=True), + ], + ) + mock_copy.assert_called_once() + mock_copy_readme.assert_called_once_with(sample_model_config, fp16_path.parent, variant="fp16") + assert (fp16_path.parent / "config.json").exists() + + def test_build_metadata_includes_core_fields_and_labels(self, converter, sample_model_config): + """_build_metadata adds required metadata fields and resolved labels.""" + with patch.object(converter, "get_labels", return_value="tabby_cat golden_retriever"): + metadata = converter._build_metadata( + { + **sample_model_config, + "confidence_threshold": 0.5, + "iou_threshold": 0.45, + }, + ) + + assert metadata["model_info", "model_type"] == "Classification" + assert metadata["model_info", "model_short_name"] == "test_model" + assert metadata["model_info", "reverse_input_channels"] == "True" + assert metadata["model_info", "mean_values"] == "123.675 116.28 103.53" + assert metadata["model_info", "scale_values"] == "58.395 57.12 57.375" + assert metadata["model_info", "confidence_threshold"] == "0.5" + assert metadata["model_info", "iou_threshold"] == "0.45" + assert metadata["model_info", "labels"] == "tabby_cat golden_retriever" + + def test_get_labels_resolves_imagenet1k(self, converter): + """get_labels resolves torchvision ImageNet-1K labels.""" + torchvision_module = types.ModuleType("torchvision") + torchvision_module.__path__ = [] + torchvision_models_module = types.ModuleType("torchvision.models") + torchvision_models_module.__path__ = [] + torchvision_meta_module = types.ModuleType("torchvision.models._meta") + torchvision_meta_module._IMAGENET_CATEGORIES = ["tabby cat", "golden retriever"] + + with patch.dict( + sys.modules, + { + "torchvision": torchvision_module, + "torchvision.models": torchvision_models_module, + "torchvision.models._meta": torchvision_meta_module, + }, + ): + assert converter.get_labels("IMAGENET1K_V1") == "tabby_cat golden_retriever" + + def test_get_labels_resolves_imagenet21k(self, converter): + """get_labels resolves timm ImageNet-21K labels.""" + timm_module = types.ModuleType("timm") + timm_module.__path__ = [] + timm_data_module = types.ModuleType("timm.data") + image_net_info = MagicMock() + image_net_info.label_descriptions.return_value = ["tabby, tabby cat", "golden retriever, dog"] + timm_data_module.ImageNetInfo = MagicMock(return_value=image_net_info) + + with patch.dict(sys.modules, {"timm": timm_module, "timm.data": timm_data_module}): + assert converter.get_labels("IMAGENET21K") == "tabby golden_retriever" + + def test_get_labels_resolves_coco(self, converter): + """get_labels resolves torchvision COCO labels.""" + torchvision_module = types.ModuleType("torchvision") + torchvision_module.__path__ = [] + torchvision_models_module = types.ModuleType("torchvision.models") + torchvision_models_module.__path__ = [] + torchvision_detection_module = types.ModuleType("torchvision.models.detection") + mock_weights = MagicMock() + mock_weights.COCO_V1.meta = {"categories": ["person", "traffic light"]} + torchvision_detection_module.MaskRCNN_ResNet50_FPN_Weights = mock_weights + + with patch.dict( + sys.modules, + { + "torchvision": torchvision_module, + "torchvision.models": torchvision_models_module, + "torchvision.models.detection": torchvision_detection_module, + }, + ): + assert converter.get_labels("COCO_V1") == "person traffic_light" + + def test_create_model_loads_state_dict_into_instance(self, converter): + """create_model instantiates a class and loads its state_dict.""" + model_class = MagicMock() + model_instance = MagicMock(spec=nn.Module) + model_instance.eval.return_value = model_instance + model_class.return_value = model_instance + checkpoint = {"state_dict": {"layer.weight": "weights"}} + + result = converter.create_model(model_class, checkpoint, model_params={"num_classes": 2}) + + assert result is model_instance + model_class.assert_called_once_with(num_classes=2) + model_instance.load_state_dict.assert_called_once_with(checkpoint["state_dict"], strict=False) + model_instance.eval.assert_called_once_with() diff --git a/model_converter/tests/unit/test_reporting.py b/model_converter/tests/unit/test_reporting.py new file mode 100644 index 000000000..846fe04ec --- /dev/null +++ b/model_converter/tests/unit/test_reporting.py @@ -0,0 +1,285 @@ +# +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +"""Tests for the conversion summary reporting module.""" + +from typing import Any + +from model_converter.reporting import ( + STATUS_ACCURACY_DROP, + STATUS_FAILED_CONVERSION, + STATUS_FAILED_QUANTIZATION, + STATUS_OK, + STATUS_OK_NO_ACCURACY, + STATUS_SKIPPED, + AccuracyResults, + ConversionResult, + determine_status, + format_console_table, + format_markdown_report, + original_url_for_config, + upsert_result, + write_markdown_report, +) + + +def _result(**kwargs) -> ConversionResult: + base: dict[str, Any] = { + "model_short_name": "m", + "model_full_name": "Model", + "model_type": "Classification", + "model_library": "timm", + } + base.update(kwargs) + return ConversionResult(**base) + + +class TestAccuracyResults: + """Tests for the AccuracyResults dataclass defaults.""" + + def test_defaults(self): + acc = AccuracyResults() + assert acc.original_accuracy is None + assert acc.fp32_accuracy is None + assert acc.fp16_accuracy is None + assert acc.int8_accuracy is None + assert acc.int8_succeeded is False + assert acc.measured is False + assert acc.metric_name is None + + def test_metric_name_can_be_set(self): + acc = AccuracyResults(metric_name="mIoU") + assert acc.metric_name == "mIoU" + + +class TestOriginalUrlForConfig: + """Tests for original_url_for_config.""" + + def test_weights_url_takes_precedence(self): + config = {"weights_url": "https://example.com/w.pth", "huggingface_repo": "org/model"} + assert original_url_for_config(config) == "https://example.com/w.pth" + + def test_huggingface_repo(self): + assert original_url_for_config({"huggingface_repo": "org/model"}) == "https://huggingface.co/org/model" + + def test_none_when_unavailable(self): + assert original_url_for_config({}) is None + + +class TestDetermineStatus: + """Tests for every branch of determine_status.""" + + def test_skipped(self): + status, detail = determine_status(_result(), converted=False, quantized=False, skipped=True) + assert status == STATUS_SKIPPED + assert detail + + def test_failed_conversion(self): + status, _ = determine_status(_result(), converted=False, quantized=False, skipped=False) + assert status == STATUS_FAILED_CONVERSION + + def test_failed_quantization(self): + status, _ = determine_status(_result(), converted=True, quantized=False, skipped=False) + assert status == STATUS_FAILED_QUANTIZATION + + def test_no_accuracy_data(self): + status, _ = determine_status( + _result(fp32_accuracy=None, original_accuracy=None), + converted=True, + quantized=True, + skipped=False, + ) + assert status == STATUS_OK_NO_ACCURACY + + def test_original_as_baseline_when_no_fp32(self): + """When fp32_accuracy is absent, original_accuracy serves as baseline → OK.""" + result = _result(original_accuracy=0.37, fp16_accuracy=0.37, int8_accuracy=0.365) + status, detail = determine_status(result, converted=True, quantized=True, skipped=False) + assert status == STATUS_OK + assert detail == "" + + def test_fp16_drop_flagged_against_original_baseline(self): + """FP16 drop > 5% relative to original_accuracy is flagged when fp32 is absent.""" + result = _result(original_accuracy=0.37, fp16_accuracy=0.30, int8_accuracy=0.30) + status, detail = determine_status(result, converted=True, quantized=True, skipped=False) + assert status == STATUS_ACCURACY_DROP + assert "FP16 drop" in detail + + def test_int8_drop_flagged_against_original_baseline(self): + """INT8 drop > 5% relative to original_accuracy is flagged when fp32 is absent.""" + result = _result(original_accuracy=0.37, fp16_accuracy=0.37, int8_accuracy=0.30) + status, detail = determine_status(result, converted=True, quantized=True, skipped=False) + assert status == STATUS_ACCURACY_DROP + assert "INT8 drop" in detail + + def test_fp32_drop_not_checked_when_fp32_is_baseline(self): + """When fp32 is used as baseline (no original), FP32-drop check is skipped.""" + result = _result(fp32_accuracy=0.90, fp16_accuracy=0.89, int8_accuracy=0.88) + status, _ = determine_status(result, converted=True, quantized=True, skipped=False) + assert status == STATUS_OK + + def test_ok_within_threshold(self): + result = _result(fp32_accuracy=0.90, fp16_accuracy=0.89, int8_accuracy=0.88) + status, detail = determine_status(result, converted=True, quantized=True, skipped=False) + assert status == STATUS_OK + assert detail == "" + + def test_boundary_exactly_5_percent_is_ok(self): + result = _result(fp32_accuracy=0.90, int8_accuracy=0.85) + status, _ = determine_status(result, converted=True, quantized=True, skipped=False) + assert status == STATUS_OK + + def test_fp16_drop_flagged(self): + result = _result(fp32_accuracy=0.90, fp16_accuracy=0.80, int8_accuracy=0.89) + status, detail = determine_status(result, converted=True, quantized=True, skipped=False) + assert status == STATUS_ACCURACY_DROP + assert "FP16 drop" in detail + + def test_int8_drop_flagged(self): + result = _result(fp32_accuracy=0.90, fp16_accuracy=0.89, int8_accuracy=0.70) + status, detail = determine_status(result, converted=True, quantized=True, skipped=False) + assert status == STATUS_ACCURACY_DROP + assert "INT8 drop" in detail + + def test_fp32_conversion_drop_flagged(self): + result = _result(original_accuracy=0.90, fp32_accuracy=0.80, fp16_accuracy=0.80, int8_accuracy=0.79) + status, detail = determine_status(result, converted=True, quantized=True, skipped=False) + assert status == STATUS_ACCURACY_DROP + assert "FP32 drop" in detail + + def test_original_within_threshold_is_ok(self): + result = _result(original_accuracy=0.90, fp32_accuracy=0.89, fp16_accuracy=0.89, int8_accuracy=0.88) + status, detail = determine_status(result, converted=True, quantized=True, skipped=False) + assert status == STATUS_OK + assert detail == "" + + +class TestRendering: + """Tests for console and markdown rendering.""" + + def test_console_table_includes_values_and_na(self): + results = [ + _result( + original_accuracy=0.9123, + fp32_accuracy=0.9012, + fp16_accuracy=0.9, + int8_accuracy=0.88, + status=STATUS_OK, + ), + _result(model_full_name="NoAcc", original_url=None, status=STATUS_OK_NO_ACCURACY), + ] + table = format_console_table(results) + assert "Conversion Summary Report" in table + assert "Original Accuracy" in table + assert "91.23%" in table + assert "90.12%" in table + assert "N/A" in table + assert "Status" in table + + def test_markdown_report_escapes_pipes_and_renders_rows(self): + results = [_result(model_full_name="A|B", fp32_accuracy=0.5, status=STATUS_OK)] + md = format_markdown_report(results) + assert md.startswith("# Conversion Summary Report") + assert "A\\|B" in md + assert "50.00%" in md + assert md.endswith("\n") + + def test_markdown_report_empty(self): + md = format_markdown_report([]) + assert "No models were processed." in md + + def test_console_table_uses_short_name_when_full_name_missing(self): + result = ConversionResult( + model_short_name="short", + model_full_name="", + model_type="", + model_library="", + ) + table = format_console_table([result]) + assert "short" in table + + def test_write_markdown_report_creates_file(self, tmp_path): + path = tmp_path / "nested" / "report.md" + write_markdown_report([_result(status=STATUS_OK)], path) + assert path.exists() + assert "# Conversion Summary Report" in path.read_text() + + +class TestUpsertResult: + """Tests for upsert_result.""" + + def test_creates_files_on_first_write(self, tmp_path): + """upsert_result creates both the Markdown and JSON sidecar.""" + path = tmp_path / "report.md" + upsert_result(_result(model_short_name="m1", model_full_name="Model One", status=STATUS_OK), path) + assert path.exists() + assert path.with_suffix(".json").exists() + assert "# Conversion Summary Report" in path.read_text() + + def test_appends_new_entry(self, tmp_path): + """upsert_result appends a new entry when the model is not yet tracked.""" + path = tmp_path / "report.md" + upsert_result(_result(model_short_name="m1", model_full_name="Model One", status=STATUS_OK), path) + upsert_result(_result(model_short_name="m2", model_full_name="Model Two", status=STATUS_OK), path) + md = path.read_text() + assert "Model One" in md + assert "Model Two" in md + + def test_replaces_existing_entry(self, tmp_path): + """upsert_result replaces an existing entry for the same model_short_name.""" + path = tmp_path / "report.md" + upsert_result( + _result(model_short_name="m1", model_full_name="Model One", status=STATUS_FAILED_CONVERSION), + path, + ) + upsert_result(_result(model_short_name="m1", model_full_name="Model One", status=STATUS_OK), path) + md = path.read_text() + assert STATUS_OK in md + assert STATUS_FAILED_CONVERSION not in md + # Only one row for Model One + assert md.count("Model One") == 1 + + def test_preserves_other_entries_on_replace(self, tmp_path): + """upsert_result keeps all other entries when replacing one.""" + path = tmp_path / "report.md" + upsert_result(_result(model_short_name="m1", model_full_name="Model One", status=STATUS_OK), path) + upsert_result(_result(model_short_name="m2", model_full_name="Model Two", status=STATUS_OK), path) + upsert_result( + _result(model_short_name="m1", model_full_name="Model One", status=STATUS_FAILED_CONVERSION), + path, + ) + md = path.read_text() + assert "Model Two" in md + assert STATUS_FAILED_CONVERSION in md + + def test_creates_parent_directories(self, tmp_path): + """upsert_result creates missing parent directories.""" + path = tmp_path / "nested" / "deep" / "report.md" + upsert_result(_result(status=STATUS_OK), path) + assert path.exists() + + def test_json_sidecar_is_valid(self, tmp_path): + """upsert_result writes valid JSON to the sidecar file.""" + import json + + path = tmp_path / "report.md" + upsert_result(_result(model_short_name="m1", status=STATUS_OK), path) + data = json.loads(path.with_suffix(".json").read_text()) + assert isinstance(data, list) + assert data[0]["model_short_name"] == "m1" + + def test_handles_corrupt_json_sidecar(self, tmp_path): + """upsert_result gracefully handles a corrupt JSON sidecar by treating it as empty.""" + import json + + path = tmp_path / "report.md" + json_path = path.with_suffix(".json") + json_path.write_text("{ not valid json }") + + upsert_result(_result(model_short_name="m1", status=STATUS_OK), path) + + data = json.loads(json_path.read_text()) + assert len(data) == 1 + assert data[0]["model_short_name"] == "m1" diff --git a/model_converter/tests/unit/test_yolo.py b/model_converter/tests/unit/test_yolo.py index b6d48ccc5..8e14620d6 100644 --- a/model_converter/tests/unit/test_yolo.py +++ b/model_converter/tests/unit/test_yolo.py @@ -3,170 +3,643 @@ # SPDX-License-Identifier: Apache-2.0 # -"""Tests for model_converter.yolo module.""" +"""Tests for YoloConverter — including COCO mAP accuracy measurement.""" import json from pathlib import Path from unittest.mock import MagicMock, patch +import cv2 +import numpy as np +import pytest +from model_converter.converters.yolo import _COCO80_TO_COCO91, YoloConverter +from model_converter.datasets.base import CalibrationSample +from model_converter.metrics.coco_detection import COCO80_TO_COCO91 +from model_converter.reporting import AccuracyResults -class TestCopyReadmeTemplate: - """Tests for copy_readme_template function.""" +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- - def test_copies_and_replaces_placeholder(self, tmp_path): - """copy_readme_template replaces <> in template.""" - from model_converter.yolo.yolo import copy_readme_template - # Create a template in the actual templates dir (we need to mock it) - template_content = "# YOLO11<> Model\nSize: <>" - dest_dir = tmp_path / "output" - dest_dir.mkdir() +def _make_yolo_converter(tmp_output_dir, tmp_cache_dir, *, measure_accuracy=True, registry=None): + return YoloConverter( + output_dir=tmp_output_dir, + cache_dir=tmp_cache_dir, + verbose=True, + dataset_registry=registry, + measure_accuracy=measure_accuracy, + ) + + +def _make_coco_dataset(root: Path) -> Path: + """Create a minimal COCO-style dataset directory.""" + images_dir = root / "images" + annotations_dir = root / "annotations" + images_dir.mkdir(parents=True) + annotations_dir.mkdir(parents=True) + + img = np.zeros((64, 64, 3), dtype=np.uint8) + img_path = images_dir / "000000000001.jpg" + cv2.imwrite(str(img_path), img) + + annotation = { + "images": [{"id": 1, "file_name": "000000000001.jpg"}], + "annotations": [ + { + "id": 1, + "image_id": 1, + "category_id": 1, + "bbox": [10, 10, 20, 20], + "area": 400, + "iscrowd": 0, + }, + ], + "categories": [{"id": 1, "name": "person"}], + } + (annotations_dir / "instances_val2017.json").write_text(json.dumps(annotation)) + return root + + +def _make_coco_sample(image_path: Path) -> CalibrationSample: + return CalibrationSample(image_path=image_path, label=0, image_id=1) + + +# --------------------------------------------------------------------------- +# Module-level constant +# --------------------------------------------------------------------------- + + +class TestCoco80ToCoco91: + def test_length_is_80(self): + assert len(_COCO80_TO_COCO91) == 80 + + def test_first_and_last_values(self): + assert _COCO80_TO_COCO91[0] == 1 + assert _COCO80_TO_COCO91[-1] == 90 + + def test_no_duplicates(self): + assert len(set(_COCO80_TO_COCO91)) == 80 + + def test_alias_matches_canonical(self): + assert _COCO80_TO_COCO91 == COCO80_TO_COCO91 + + def test_class_11_maps_to_13_not_12(self): + """Stop sign (class 11) maps to COCO category 13, not 12 (missing category).""" + assert _COCO80_TO_COCO91[11] == 13 + + +# --------------------------------------------------------------------------- +# _measure_original_accuracy +# --------------------------------------------------------------------------- + + +class TestMeasureOriginalAccuracy: + """Unit tests for YoloConverter._measure_original_accuracy.""" + + @pytest.fixture + def converter(self, tmp_path): + return _make_yolo_converter(tmp_path / "out", tmp_path / "cache") + + def test_returns_map_on_success(self, converter, tmp_path): + """Returns a float mAP when inference and evaluation succeed.""" + coco_root = _make_coco_dataset(tmp_path / "coco") + annotation_file = coco_root / "annotations" / "instances_val2017.json" + img_path = coco_root / "images" / "000000000001.jpg" + sample = _make_coco_sample(img_path) + + fake_result = MagicMock() + fake_result.boxes.xyxy.cpu.return_value.tolist.return_value = [[5.0, 5.0, 25.0, 25.0]] + fake_result.boxes.cls.cpu.return_value.tolist.return_value = [0.0] + fake_result.boxes.conf.cpu.return_value.tolist.return_value = [0.9] + + fake_model = MagicMock(return_value=[fake_result]) + + with patch("model_converter.converters.yolo.CocoDetectionMAP") as mock_metric_cls: + mock_metric = MagicMock() + mock_metric.compute.return_value = 0.42 + mock_metric_cls.return_value = mock_metric + + with patch("ultralytics.YOLO", return_value=fake_model): + result = converter._measure_original_accuracy( + pt_model_path=tmp_path / "cache" / "yolo11n.pt", + samples=[sample], + annotation_file=annotation_file, + ) + + assert result == pytest.approx(0.42) + mock_metric.update.assert_called_once() + preds = mock_metric.update.call_args[1]["predictions"] + assert len(preds) == 1 + assert preds[0]["image_id"] == 1 + assert preds[0]["category_id"] == _COCO80_TO_COCO91[0] # class 0 → COCO cat 1 + + def test_returns_none_when_ultralytics_import_fails(self, converter, tmp_path): + """Returns None gracefully when Ultralytics is not importable.""" + coco_root = _make_coco_dataset(tmp_path / "coco") + annotation_file = coco_root / "annotations" / "instances_val2017.json" + sample = _make_coco_sample(coco_root / "images" / "000000000001.jpg") + + with patch("ultralytics.YOLO", side_effect=ImportError("no ultralytics")): + result = converter._measure_original_accuracy( + pt_model_path=tmp_path / "cache" / "yolo11n.pt", + samples=[sample], + annotation_file=annotation_file, + ) + + assert result is None + + def test_returns_none_when_pt_file_not_found(self, converter, tmp_path): + """Returns None gracefully when the .pt weights file is absent.""" + coco_root = _make_coco_dataset(tmp_path / "coco") + annotation_file = coco_root / "annotations" / "instances_val2017.json" + sample = _make_coco_sample(coco_root / "images" / "000000000001.jpg") + + with patch("ultralytics.YOLO", side_effect=FileNotFoundError("not found")): + result = converter._measure_original_accuracy( + pt_model_path=tmp_path / "cache" / "missing.pt", + samples=[sample], + annotation_file=annotation_file, + ) + + assert result is None + + def test_skips_samples_without_image_id(self, converter, tmp_path): + """Samples with image_id=None are silently skipped.""" + coco_root = _make_coco_dataset(tmp_path / "coco") + annotation_file = coco_root / "annotations" / "instances_val2017.json" + img_path = coco_root / "images" / "000000000001.jpg" + sample_no_id = CalibrationSample(image_path=img_path, label=0, image_id=None) + + fake_model = MagicMock() + + with patch("model_converter.converters.yolo.CocoDetectionMAP") as mock_metric_cls: + mock_metric = MagicMock() + mock_metric.compute.return_value = 0.0 + mock_metric_cls.return_value = mock_metric + + with patch("ultralytics.YOLO", return_value=fake_model): + converter._measure_original_accuracy( + pt_model_path=tmp_path / "cache" / "yolo11n.pt", + samples=[sample_no_id], + annotation_file=annotation_file, + ) + + fake_model.assert_not_called() + + def test_skips_unreadable_images(self, converter, tmp_path): + """Samples whose image file cannot be read are silently skipped.""" + coco_root = _make_coco_dataset(tmp_path / "coco") + annotation_file = coco_root / "annotations" / "instances_val2017.json" + sample = CalibrationSample( + image_path=tmp_path / "does_not_exist.jpg", + label=0, + image_id=1, + ) + + fake_model = MagicMock() + + with patch("model_converter.converters.yolo.CocoDetectionMAP") as mock_metric_cls: + mock_metric = MagicMock() + mock_metric.compute.return_value = 0.0 + mock_metric_cls.return_value = mock_metric + + with patch("ultralytics.YOLO", return_value=fake_model): + converter._measure_original_accuracy( + pt_model_path=tmp_path / "cache" / "yolo11n.pt", + samples=[sample], + annotation_file=annotation_file, + ) + + fake_model.assert_not_called() + + def test_skips_failed_inference_samples(self, converter, tmp_path): + """Per-sample inference errors are swallowed; remaining samples still processed.""" + coco_root = _make_coco_dataset(tmp_path / "coco") + annotation_file = coco_root / "annotations" / "instances_val2017.json" + img_path = coco_root / "images" / "000000000001.jpg" + sample = _make_coco_sample(img_path) + + fake_model = MagicMock(side_effect=RuntimeError("inference exploded")) + + with patch("model_converter.converters.yolo.CocoDetectionMAP") as mock_metric_cls: + mock_metric = MagicMock() + mock_metric.compute.return_value = 0.0 + mock_metric_cls.return_value = mock_metric + + with patch("ultralytics.YOLO", return_value=fake_model): + result = converter._measure_original_accuracy( + pt_model_path=tmp_path / "cache" / "yolo11n.pt", + samples=[sample], + annotation_file=annotation_file, + ) + + assert result == pytest.approx(0.0) + mock_metric.update.assert_called_once_with(predictions=[]) + + def test_uses_fallback_category_id_for_out_of_range_cls(self, converter, tmp_path): + """Class indices >= 80 fall back to idx+1 instead of the lookup table.""" + coco_root = _make_coco_dataset(tmp_path / "coco") + annotation_file = coco_root / "annotations" / "instances_val2017.json" + img_path = coco_root / "images" / "000000000001.jpg" + sample = _make_coco_sample(img_path) + + fake_result = MagicMock() + fake_result.boxes.xyxy.cpu.return_value.tolist.return_value = [[0.0, 0.0, 10.0, 10.0]] + fake_result.boxes.cls.cpu.return_value.tolist.return_value = [99.0] # out of range + fake_result.boxes.conf.cpu.return_value.tolist.return_value = [0.5] + fake_model = MagicMock(return_value=[fake_result]) + + with patch("model_converter.converters.yolo.CocoDetectionMAP") as mock_metric_cls: + mock_metric = MagicMock() + mock_metric.compute.return_value = 0.0 + mock_metric_cls.return_value = mock_metric + + with patch("ultralytics.YOLO", return_value=fake_model): + converter._measure_original_accuracy( + pt_model_path=tmp_path / "cache" / "yolo11n.pt", + samples=[sample], + annotation_file=annotation_file, + ) + + preds = mock_metric.update.call_args[1]["predictions"] + assert preds[0]["category_id"] == 100 # 99 + 1 fallback + + +# --------------------------------------------------------------------------- +# _measure_yolo_accuracy +# --------------------------------------------------------------------------- + + +class TestMeasureYoloAccuracy: + """Unit tests for YoloConverter._measure_yolo_accuracy.""" + + @pytest.fixture + def coco_dataset(self, tmp_path): + return _make_coco_dataset(tmp_path / "coco") + + @pytest.fixture + def mock_registry(self, coco_dataset): + mock = MagicMock() + mock.resolve_from_config.return_value = coco_dataset + return mock + + @pytest.fixture + def converter(self, tmp_path, mock_registry): + return _make_yolo_converter(tmp_path / "out", tmp_path / "cache", registry=mock_registry) + + @pytest.fixture + def yolo_config(self): + return { + "model_short_name": "YOLO11n", + "model_full_name": "YOLO11n", + "model_library": "yolo", + "yolo_version": "yolo11n", + "model_type": "YOLO11", + "dataset_type": "coco-detection", + "license": "agpl-3.0", + "license_link": "https://spdx.org/licenses/AGPL-3.0-only.html", + } + + def test_returns_accuracy_results_on_success(self, converter, yolo_config, coco_dataset, tmp_path): + """Full happy path: all three mAP values are populated.""" + fp16_folder = tmp_path / "out" / "YOLO11n-fp16-ov" + int8_folder = tmp_path / "out" / "YOLO11n-int8-ov" + fp16_folder.mkdir(parents=True) + int8_folder.mkdir(parents=True) + (fp16_folder / "yolo11n.xml").write_text("") + (int8_folder / "yolo11n.xml").write_text("") - with patch.object(Path, "read_text", return_value=template_content): - copy_readme_template("README-yolo-fp16.md", dest_dir, "n") + with ( + patch.object(converter, "_collect_validation_samples") as mock_samples, + patch.object(converter, "_metric_for_config") as mock_metric_fn, + patch.object(converter, "_measure_metric") as mock_measure, + patch.object(converter, "_measure_original_accuracy") as mock_original, + ): + mock_samples.return_value = [_make_coco_sample(coco_dataset / "images" / "000000000001.jpg")] + mock_metric = MagicMock() + mock_metric.name = "mAP" + mock_metric_fn.return_value = mock_metric + mock_original.return_value = 0.50 + mock_measure.side_effect = [0.48, 0.45] # fp16, int8 + + result = converter._measure_yolo_accuracy(yolo_config, "yolo11n", fp16_folder, int8_folder) + + assert result is not None + assert result.measured is True + assert result.original_accuracy == pytest.approx(0.50) + assert result.fp16_accuracy == pytest.approx(0.48) + assert result.int8_accuracy == pytest.approx(0.45) + assert result.metric_name == "mAP" + + def test_returns_none_when_dataset_path_missing(self, tmp_path, yolo_config): + """Returns None when dataset registry resolves to a non-existent path.""" + mock_registry = MagicMock() + mock_registry.resolve_from_config.return_value = tmp_path / "nonexistent" + conv = _make_yolo_converter(tmp_path / "out", tmp_path / "cache", registry=mock_registry) + + result = conv._measure_yolo_accuracy( + yolo_config, + "yolo11n", + tmp_path / "fp16", + tmp_path / "int8", + ) + + assert result is None + + def test_returns_none_when_dataset_path_is_none(self, tmp_path, yolo_config): + """Returns None when dataset registry returns None.""" + mock_registry = MagicMock() + mock_registry.resolve_from_config.return_value = None + conv = _make_yolo_converter(tmp_path / "out", tmp_path / "cache", registry=mock_registry) + + result = conv._measure_yolo_accuracy( + yolo_config, + "yolo11n", + tmp_path / "fp16", + tmp_path / "int8", + ) + + assert result is None + + def test_returns_none_when_dataset_type_not_coco( + self, + tmp_path, + coco_dataset, + mock_registry, + ): + """Returns None for non-COCO dataset types (no annotation file entry).""" + conv = _make_yolo_converter(tmp_path / "out", tmp_path / "cache", registry=mock_registry) + config = { + "model_short_name": "YOLO11n", + "model_type": "YOLO11", + "dataset_type": "imagenet-1k", # not in _COCO_ANNOTATION_FILES + } + + result = conv._measure_yolo_accuracy(config, "yolo11n", tmp_path / "fp16", tmp_path / "int8") + + assert result is None + + def test_returns_none_when_annotation_file_missing(self, tmp_path, yolo_config): + """Returns None when the COCO annotation JSON file is absent.""" + coco_root = tmp_path / "coco" + (coco_root / "images").mkdir(parents=True) + # No annotations/ directory created + mock_registry = MagicMock() + mock_registry.resolve_from_config.return_value = coco_root + conv = _make_yolo_converter(tmp_path / "out", tmp_path / "cache", registry=mock_registry) + + result = conv._measure_yolo_accuracy( + yolo_config, + "yolo11n", + tmp_path / "fp16", + tmp_path / "int8", + ) + + assert result is None + + def test_returns_none_when_no_validation_samples( + self, + converter, + yolo_config, + tmp_path, + ): + """Returns None when _collect_validation_samples returns an empty list.""" + with patch.object(converter, "_collect_validation_samples", return_value=[]): + result = converter._measure_yolo_accuracy( + yolo_config, + "yolo11n", + tmp_path / "fp16", + tmp_path / "int8", + ) + + assert result is None + + def test_returns_none_when_metric_is_none(self, converter, yolo_config, coco_dataset, tmp_path): + """Returns None when _metric_for_config cannot produce a metric.""" + sample = _make_coco_sample(coco_dataset / "images" / "000000000001.jpg") + with ( + patch.object(converter, "_collect_validation_samples", return_value=[sample]), + patch.object(converter, "_metric_for_config", return_value=None), + ): + result = converter._measure_yolo_accuracy( + yolo_config, + "yolo11n", + tmp_path / "fp16", + tmp_path / "int8", + ) + + assert result is None + + def test_skips_fp16_when_xml_missing(self, converter, yolo_config, coco_dataset, tmp_path): + """FP16 accuracy is not measured when the FP16 XML file does not exist.""" + fp16_folder = tmp_path / "fp16" # no XML inside + int8_folder = tmp_path / "int8" + fp16_folder.mkdir() + int8_folder.mkdir() + (int8_folder / "yolo11n.xml").write_text("") + + sample = _make_coco_sample(coco_dataset / "images" / "000000000001.jpg") + mock_metric = MagicMock() + mock_metric.name = "mAP" - readme = dest_dir / "README.md" - assert readme.exists() - content = readme.read_text() - assert "YOLO11n Model" in content - assert "Size: n" in content - assert "<>" not in content + with ( + patch.object(converter, "_collect_validation_samples", return_value=[sample]), + patch.object(converter, "_metric_for_config", return_value=mock_metric), + patch.object(converter, "_measure_original_accuracy", return_value=0.5), + patch.object(converter, "_measure_metric", return_value=0.44) as mock_measure, + ): + result = converter._measure_yolo_accuracy(yolo_config, "yolo11n", fp16_folder, int8_folder) + assert result is not None + assert result.fp16_accuracy is None + # _measure_metric called once for INT8 only + assert mock_measure.call_count == 1 -class TestUpdateModelTypeInXml: - """Tests for update_model_type_in_xml function.""" + def test_skips_int8_when_xml_missing(self, converter, yolo_config, coco_dataset, tmp_path): + """INT8 accuracy is not measured when the INT8 XML file does not exist.""" + fp16_folder = tmp_path / "fp16" + int8_folder = tmp_path / "int8" # no XML inside + fp16_folder.mkdir() + int8_folder.mkdir() + (fp16_folder / "yolo11n.xml").write_text("") - def test_updates_model_type(self, tmp_path): - """update_model_type_in_xml updates model_type value in XML.""" - from model_converter.yolo.yolo import update_model_type_in_xml + sample = _make_coco_sample(coco_dataset / "images" / "000000000001.jpg") + mock_metric = MagicMock() + mock_metric.name = "mAP" - # Create a valid OpenVINO XML structure - xml_content = """ - - - - - - - -""" - xml_path = tmp_path / "model.xml" - xml_path.write_text(xml_content) + with ( + patch.object(converter, "_collect_validation_samples", return_value=[sample]), + patch.object(converter, "_metric_for_config", return_value=mock_metric), + patch.object(converter, "_measure_original_accuracy", return_value=0.5), + patch.object(converter, "_measure_metric", return_value=0.48) as mock_measure, + ): + result = converter._measure_yolo_accuracy(yolo_config, "yolo11n", fp16_folder, int8_folder) - update_model_type_in_xml(xml_path, "YOLO11") + assert result is not None + assert result.int8_accuracy is None + # _measure_metric called once for FP16 only + assert mock_measure.call_count == 1 - # Verify XML was updated - from defusedxml.ElementTree import parse - tree = parse(xml_path) - root = tree.getroot() - model_type_elem = root.find(".//rt_info/model_info/model_type") - assert model_type_elem is not None - assert model_type_elem.get("value") == "YOLO11" +# --------------------------------------------------------------------------- +# process_model_config — accuracy integration +# --------------------------------------------------------------------------- - # Verify config.json was created - config_json = tmp_path / "config.json" - assert config_json.exists() - config = json.loads(config_json.read_text()) - assert config["model_type"] == "YOLO11" - assert config["labels"] == "person car" - def test_handles_parse_error(self, tmp_path, capsys): - """update_model_type_in_xml handles invalid XML gracefully.""" - from model_converter.yolo.yolo import update_model_type_in_xml +class TestProcessModelConfigAccuracy: + """Integration-style tests verifying accuracy is wired into process_model_config.""" - xml_path = tmp_path / "bad.xml" - xml_path.write_text("not valid xml <<<>>>") + @pytest.fixture + def yolo_config(self): + return { + "model_short_name": "YOLO11n", + "model_full_name": "YOLO11n", + "model_library": "yolo", + "yolo_version": "yolo11n", + "model_type": "YOLO11", + "dataset_type": "coco-detection", + "license": "agpl-3.0", + "license_link": "https://spdx.org/licenses/AGPL-3.0-only.html", + } - # Should not raise - update_model_type_in_xml(xml_path, "YOLO11") + def _stub_ultralytics_export(self, cache_dir: Path, model_short_name: str, yolo_version: str) -> MagicMock: + """Return a fake YOLO model that fakes both FP16 and INT8 exports.""" + fp16_export = cache_dir / f"{yolo_version}_openvino_model" + int8_export = cache_dir / f"{yolo_version}_int8_openvino_model" + fp16_export.mkdir(parents=True) + int8_export.mkdir(parents=True) + (fp16_export / f"{yolo_version}.xml").write_text("") + (int8_export / f"{yolo_version}.xml").write_text("") - captured = capsys.readouterr() - assert "Failed to update" in captured.out + fake_model = MagicMock() - def test_handles_file_not_found(self, tmp_path, capsys): - """update_model_type_in_xml handles missing file gracefully.""" - from model_converter.yolo.yolo import update_model_type_in_xml + def fake_export(format, **kwargs): + pass - xml_path = tmp_path / "nonexistent.xml" + fake_model.export.side_effect = fake_export + return fake_model - update_model_type_in_xml(xml_path, "YOLO11") + def test_accuracy_measured_when_flag_true_and_dataset_available( + self, + tmp_path, + yolo_config, + ): + """process_model_config calls _measure_yolo_accuracy when measure_accuracy=True.""" + cache_dir = tmp_path / "cache" + cache_dir.mkdir() + output_dir = tmp_path / "out" + output_dir.mkdir() + + conv = _make_yolo_converter(output_dir, cache_dir) + fake_model = self._stub_ultralytics_export(cache_dir, "YOLO11n", "yolo11n") + fake_accuracy = AccuracyResults() + fake_accuracy.original_accuracy = 0.5 + fake_accuracy.fp16_accuracy = 0.48 + fake_accuracy.int8_accuracy = 0.45 + fake_accuracy.measured = True + + with ( + patch("ultralytics.YOLO", return_value=fake_model), + patch.object(conv, "_update_model_type_in_xml"), + patch.object(conv, "_copy_yolo_readme"), + patch.object(conv, "_measure_yolo_accuracy", return_value=fake_accuracy) as mock_acc, + ): + success = conv.process_model_config(yolo_config) + + assert success is True + mock_acc.assert_called_once() + assert len(conv.results) == 1 + result = conv.results[0] + assert result.original_accuracy == pytest.approx(0.5) + assert result.fp16_accuracy == pytest.approx(0.48) + assert result.int8_accuracy == pytest.approx(0.45) + + def test_accuracy_skipped_when_measure_accuracy_false(self, tmp_path, yolo_config): + """process_model_config skips accuracy when measure_accuracy=False.""" + cache_dir = tmp_path / "cache" + cache_dir.mkdir() + output_dir = tmp_path / "out" + output_dir.mkdir() + + conv = _make_yolo_converter(output_dir, cache_dir, measure_accuracy=False) + fake_model = self._stub_ultralytics_export(cache_dir, "YOLO11n", "yolo11n") + + with ( + patch("ultralytics.YOLO", return_value=fake_model), + patch.object(conv, "_update_model_type_in_xml"), + patch.object(conv, "_copy_yolo_readme"), + patch.object(conv, "_measure_yolo_accuracy") as mock_acc, + ): + success = conv.process_model_config(yolo_config) - captured = capsys.readouterr() - assert "Failed to update" in captured.out + assert success is True + mock_acc.assert_not_called() + result = conv.results[0] + assert result.original_accuracy is None + def test_accuracy_skipped_when_int8_not_produced(self, tmp_path, yolo_config): + """process_model_config skips accuracy when INT8 XML is missing.""" + cache_dir = tmp_path / "cache" + cache_dir.mkdir() + output_dir = tmp_path / "out" + output_dir.mkdir() -class TestConvertYoloModels: - """Tests for convert_yolo_models function.""" + conv = _make_yolo_converter(output_dir, cache_dir) - @patch("model_converter.yolo.yolo.copy_readme_template") - @patch("model_converter.yolo.yolo.update_model_type_in_xml") - @patch("model_converter.yolo.yolo.YOLO") - @patch("shutil.rmtree") - def test_convert_single_model(self, mock_rmtree, mock_yolo_class, mock_update_xml, mock_copy_readme, tmp_path): - """convert_yolo_models converts a single YOLO variant.""" - from model_converter.yolo.yolo import convert_yolo_models + # Only produce FP16 export, not INT8 + fp16_export = cache_dir / "yolo11n_openvino_model" + fp16_export.mkdir(parents=True) + (fp16_export / "yolo11n.xml").write_text("") - mock_model = MagicMock() - mock_yolo_class.return_value = mock_model + fake_model = MagicMock() + fake_model.export.return_value = None - # Mock that the output folders exist after export with ( - patch("model_converter.yolo.yolo.Path"), - patch.object(Path, "exists", return_value=False), + patch("ultralytics.YOLO", return_value=fake_model), + patch.object(conv, "_update_model_type_in_xml"), + patch.object(conv, "_copy_yolo_readme"), + patch.object(conv, "_measure_yolo_accuracy") as mock_acc, ): - convert_yolo_models(["yolo11n"]) + success = conv.process_model_config(yolo_config) - mock_yolo_class.assert_called_once_with("yolo11n.pt") - assert mock_model.export.call_count == 2 + assert success is True + mock_acc.assert_not_called() - @patch("model_converter.yolo.yolo.copy_readme_template") - @patch("model_converter.yolo.yolo.update_model_type_in_xml") - @patch("model_converter.yolo.yolo.YOLO") - def test_convert_with_existing_output( - self, - mock_yolo_class, - mock_update_xml, - mock_copy_readme, - tmp_path, - monkeypatch, - ): - """convert_yolo_models handles existing output directories.""" - from model_converter.yolo.yolo import convert_yolo_models + def test_skipped_models_do_not_measure_accuracy(self, tmp_path, yolo_config): + """process_model_config skips accuracy for already-converted models.""" + cache_dir = tmp_path / "cache" + output_dir = tmp_path / "out" + fp16_folder = output_dir / "YOLO11n-fp16-ov" + int8_folder = output_dir / "YOLO11n-int8-ov" + fp16_folder.mkdir(parents=True) + int8_folder.mkdir(parents=True) + (fp16_folder / "yolo11n.xml").write_text("") + (int8_folder / "yolo11n.xml").write_text("") - monkeypatch.chdir(tmp_path) + conv = _make_yolo_converter(output_dir, cache_dir) - mock_model = MagicMock() - mock_yolo_class.return_value = mock_model + with patch.object(conv, "_measure_yolo_accuracy") as mock_acc: + success = conv.process_model_config(yolo_config) - # Create the output dirs that YOLO export would create - fp16_dir = tmp_path / "yolo11n_openvino_model" - fp16_dir.mkdir() - (fp16_dir / "yolo11n.xml").write_text("") + assert success is True + mock_acc.assert_not_called() - int8_dir = tmp_path / "yolo11n_int8_openvino_model" - int8_dir.mkdir() - (int8_dir / "yolo11n.xml").write_text("") + def test_failed_conversion_does_not_measure_accuracy(self, tmp_path, yolo_config): + """process_model_config records failure without measuring accuracy.""" + cache_dir = tmp_path / "cache" + cache_dir.mkdir() + output_dir = tmp_path / "out" + output_dir.mkdir() - # Create target dirs that already exist (to test rmtree path) - (tmp_path / "YOLO11n-fp16-ov").mkdir() - (tmp_path / "YOLO11n-int8-ov").mkdir() + conv = _make_yolo_converter(output_dir, cache_dir) with ( - patch("model_converter.yolo.yolo.update_model_type_in_xml"), - patch("model_converter.yolo.yolo.copy_readme_template"), + patch("ultralytics.YOLO", side_effect=RuntimeError("export failed")), + patch.object(conv, "_measure_yolo_accuracy") as mock_acc, ): - convert_yolo_models(["yolo11n"]) - - -class TestYoloMain: - """Tests for yolo main function.""" - - @patch("model_converter.yolo.yolo.convert_yolo_models") - def test_main_returns_zero(self, mock_convert): - """main() calls convert_yolo_models and returns 0.""" - from model_converter.yolo.yolo import main + success = conv.process_model_config(yolo_config) - result = main() - assert result == 0 - mock_convert.assert_called_once_with() + assert success is False + mock_acc.assert_not_called() + assert conv.results[0].original_accuracy is None diff --git a/model_converter/uv.lock b/model_converter/uv.lock index fe1249d1d..05b03c091 100644 --- a/model_converter/uv.lock +++ b/model_converter/uv.lock @@ -730,6 +730,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0a/dd/8050c947d435c8d4bc94e3252f4d8bb8a76cfb424f043a8680be637a57f1/kiwisolver-1.5.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:59cd8683f575d96df5bb48f6add94afc055012c29e28124fcae2b63661b9efb1", size = 73558, upload-time = "2026-03-09T13:15:52.112Z" }, ] +[[package]] +name = "lightning-utilities" +version = "0.15.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "packaging" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f1/45/7fa8f56b17dc0f0a41ec70dd307ecd6787254483549843bef4c30ab5adce/lightning_utilities-0.15.3.tar.gz", hash = "sha256:792ae0204c79f6859721ac7f386c237a33b0ed06ba775009cb894e010a842033", size = 33553, upload-time = "2026-02-22T14:48:53.348Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/25/f4/ead6e0e37209b07c9baa3e984ccdb0348ca370b77cea3aaea8ddbb097e00/lightning_utilities-0.15.3-py3-none-any.whl", hash = "sha256:6c55f1bee70084a1cbeaa41ada96e4b3a0fea5909e844dd335bd80f5a73c5f91", size = 31906, upload-time = "2026-02-22T14:48:52.488Z" }, +] + [[package]] name = "markdown-it-py" version = "4.2.0" @@ -1397,8 +1410,11 @@ dependencies = [ { name = "onnx" }, { name = "opencv-python-headless" }, { name = "openvino" }, + { name = "openvino-model-api" }, + { name = "pycocotools" }, { name = "timm" }, { name = "torch" }, + { name = "torchmetrics" }, { name = "torchvision" }, { name = "transformers" }, { name = "ultralytics" }, @@ -1406,7 +1422,6 @@ dependencies = [ [package.dev-dependencies] tests = [ - { name = "openvino-model-api" }, { name = "pytest" }, { name = "pytest-cov" }, { name = "pytest-mock" }, @@ -1421,8 +1436,11 @@ requires-dist = [ { name = "onnx" }, { name = "opencv-python-headless", specifier = "<4.14" }, { name = "openvino", specifier = ">=2025.3" }, + { name = "openvino-model-api", editable = "../model_api" }, + { name = "pycocotools" }, { name = "timm" }, { name = "torch" }, + { name = "torchmetrics" }, { name = "torchvision" }, { name = "transformers" }, { name = "ultralytics" }, @@ -1431,7 +1449,6 @@ requires-dist = [ [package.metadata.requires-dev] docs = [] tests = [ - { name = "openvino-model-api", editable = "../model_api" }, { name = "pytest" }, { name = "pytest-cov" }, { name = "pytest-mock" }, @@ -1622,6 +1639,45 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8c/c7/7bb2e321574b10df20cbde462a94e2b71d05f9bbda251ef27d104668306a/psutil-7.2.2-cp37-abi3-win_arm64.whl", hash = "sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee", size = 134617, upload-time = "2026-01-28T18:15:36.514Z" }, ] +[[package]] +name = "pycocotools" +version = "2.0.11" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a2/df/32354b5dda963ffdfc8f75c9acf8828ef7890723a4ed57bb3ff2dc1d6f7e/pycocotools-2.0.11.tar.gz", hash = "sha256:34254d76da85576fcaf5c1f3aa9aae16b8cb15418334ba4283b800796bd1993d", size = 25381, upload-time = "2025-12-15T22:31:46.148Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/3f/41ce3fce61b7721158f21b61727eb054805babc0088cfa48506935b80a36/pycocotools-2.0.11-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:81bdceebb4c64e9265213e2d733808a12f9c18dfb14457323cc6b9af07fa0e61", size = 158947, upload-time = "2025-12-15T22:31:03.291Z" }, + { url = "https://files.pythonhosted.org/packages/e2/9b/a739705b246445bd1376394bf9d1ec2dd292b16740e92f203461b2bb12ed/pycocotools-2.0.11-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a1c05f91ccc658dfe01325267209c4b435da1722c93eeb5749fabc1d087b6882", size = 485174, upload-time = "2025-12-15T22:31:04.395Z" }, + { url = "https://files.pythonhosted.org/packages/34/70/7a12752784e57d8034a76c245c618a2f88a9d2463862b990f314aea7e5d6/pycocotools-2.0.11-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:18ba75ff58cedb33a85ce2c18f1452f1fe20c9dd59925eec5300b2bf6205dbe1", size = 493172, upload-time = "2025-12-15T22:31:05.504Z" }, + { url = "https://files.pythonhosted.org/packages/5c/fc/d703599ac728209dba08aea8d4bee884d5adabfcd9041abed1658d863747/pycocotools-2.0.11-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:693417797f0377fd094eb815c0a1e7d1c3c0251b71e3b3779fce3b3cf24793c5", size = 480506, upload-time = "2025-12-15T22:31:06.77Z" }, + { url = "https://files.pythonhosted.org/packages/81/d9/e1cfc320bbb2cd58c3b4398c3821cbe75d93c16ed3135ac9e774a18a02d3/pycocotools-2.0.11-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:b6a07071c441d0f5e480a8f287106191582e40289d4e242dfe684e0c8a751088", size = 497595, upload-time = "2025-12-15T22:31:08.277Z" }, + { url = "https://files.pythonhosted.org/packages/a2/23/d17f6111c2a6ae8631d4fa90202bea05844da715d61431fbc34d276462d5/pycocotools-2.0.11-cp311-cp311-win_amd64.whl", hash = "sha256:8e159232adae3aef6b4e2d37b008bff107b26e9ed3b48e70ea6482302834bd34", size = 80519, upload-time = "2025-12-15T22:31:09.613Z" }, + { url = "https://files.pythonhosted.org/packages/00/4c/76b00b31a724c3f5ccdab0f85e578afb2ca38d33be0a0e98f1770cafd958/pycocotools-2.0.11-cp311-cp311-win_arm64.whl", hash = "sha256:4fc9889e819452b9c142036e1eabac8a13a8bd552d8beba299a57e0da6bfa1ec", size = 69304, upload-time = "2025-12-15T22:31:10.592Z" }, + { url = "https://files.pythonhosted.org/packages/87/12/2f2292332456e4e4aba1dec0e3de8f1fc40fb2f4fdb0ca1cb17db9861682/pycocotools-2.0.11-cp312-abi3-macosx_10_13_universal2.whl", hash = "sha256:a2e9634bc7cadfb01c88e0b98589aaf0bd12983c7927bde93f19c0103e5441f4", size = 147795, upload-time = "2025-12-15T22:31:11.519Z" }, + { url = "https://files.pythonhosted.org/packages/63/3c/68d7ea376aada9046e7ea2d7d0dad0d27e1ae8b4b3c26a28346689390ab2/pycocotools-2.0.11-cp312-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7fd4121766cc057133534679c0ec3f9023dbd96e9b31cf95c86a069ebdac2b65", size = 398434, upload-time = "2025-12-15T22:31:12.558Z" }, + { url = "https://files.pythonhosted.org/packages/23/59/dc81895beff4e1207a829d40d442ea87cefaac9f6499151965f05c479619/pycocotools-2.0.11-cp312-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a82d1c9ed83f75da0b3f244f2a3cf559351a283307bd9b79a4ee2b93ab3231dd", size = 411685, upload-time = "2025-12-15T22:31:13.995Z" }, + { url = "https://files.pythonhosted.org/packages/0b/0b/5a8a7de300862a2eb5e2ecd3cb015126231379206cd3ebba8f025388d770/pycocotools-2.0.11-cp312-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:89e853425018e2c2920ee0f2112cf7c140a1dcf5f4f49abd9c2da112c3e0f4b3", size = 390500, upload-time = "2025-12-15T22:31:15.138Z" }, + { url = "https://files.pythonhosted.org/packages/63/b5/519bb68647f06feea03d5f355c33c05800aeae4e57b9482b2859eb00752e/pycocotools-2.0.11-cp312-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:87af87b8d06d5b852a885a319d9362dca3bed9f8bbcc3feb6513acb1f88ea242", size = 409790, upload-time = "2025-12-15T22:31:16.326Z" }, + { url = "https://files.pythonhosted.org/packages/83/b4/f6708404ff494706b80e714b919f76dc4ec9845a4007affd6d6b0843f928/pycocotools-2.0.11-cp312-abi3-win_amd64.whl", hash = "sha256:ffe806ce535f5996445188f9a35643791dc54beabc61bd81e2b03367356d604f", size = 77570, upload-time = "2025-12-15T22:31:17.703Z" }, + { url = "https://files.pythonhosted.org/packages/6e/63/778cd0ddc9d4a78915ac0a72b56d7fb204f7c3fabdad067d67ea0089762e/pycocotools-2.0.11-cp312-abi3-win_arm64.whl", hash = "sha256:c230f5e7b14bd19085217b4f40bba81bf14a182b150b8e9fab1c15d504ade343", size = 64564, upload-time = "2025-12-15T22:31:18.652Z" }, + { url = "https://files.pythonhosted.org/packages/5d/78/31c81e99d596a20c137d8a2e7a25f39a88f88fada5e0b253fce7323ecf0d/pycocotools-2.0.11-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:fd72b9734e6084b217c1fc3945bfd4ec05bdc75a44e4f0c461a91442bb804973", size = 168931, upload-time = "2025-12-15T22:31:19.845Z" }, + { url = "https://files.pythonhosted.org/packages/5f/63/fdd488e4cd0fdc6f93134f2cd68b1fce441d41566e86236bf6156961ef9b/pycocotools-2.0.11-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f7eb43b79448476b094240450420b7425d06e297880144b8ea6f01e9b4340e43", size = 484856, upload-time = "2025-12-15T22:31:21.231Z" }, + { url = "https://files.pythonhosted.org/packages/a1/fc/c83648a8fb7ea3b8e2ce2e761b469807e6cadb81577bf1af31c4f2ef0d87/pycocotools-2.0.11-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c3546b93b39943347c4f5b0694b5824105cbe2174098a416bcad4acd9c21e957", size = 480994, upload-time = "2025-12-15T22:31:22.426Z" }, + { url = "https://files.pythonhosted.org/packages/b6/2d/35e1122c0d007288aa9545be9549cbc7a4987b2c22f21d75045260a8b5b8/pycocotools-2.0.11-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:efd1694b2075f2f10c5828f10f6e6c4e44368841fd07dae385c3aa015c8e25f9", size = 467956, upload-time = "2025-12-15T22:31:23.754Z" }, + { url = "https://files.pythonhosted.org/packages/e4/ff/30cfe8142470da3e45abe43a9842449ca0180d993320559890e2be19e4a5/pycocotools-2.0.11-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:368244f30eb8d6cae7003aa2c0831fbdf0153664a32859ec7fbceea52bfb6878", size = 474658, upload-time = "2025-12-15T22:31:24.883Z" }, + { url = "https://files.pythonhosted.org/packages/bc/62/254ca92604106c7a5af3258e589e465e681fe0166f9b10f97d8ca70934d6/pycocotools-2.0.11-cp313-cp313t-win_amd64.whl", hash = "sha256:ac8aa17263e6489aa521f9fa91e959dfe0ea3a5519fde2cbf547312cdce7559e", size = 89681, upload-time = "2025-12-15T22:31:26.025Z" }, + { url = "https://files.pythonhosted.org/packages/6e/f0/c019314dc122ad5e6281de420adc105abe9b59d00008f72ef3ad32b1e328/pycocotools-2.0.11-cp313-cp313t-win_arm64.whl", hash = "sha256:04480330df5013f6edd94891a0ee8294274185f1b5093d1b0f23d51778f0c0e9", size = 70520, upload-time = "2025-12-15T22:31:26.999Z" }, + { url = "https://files.pythonhosted.org/packages/66/2b/58b35c88f2086c043ff1c87bd8e7bf36f94e84f7b01a5e00b6f5fabb92a7/pycocotools-2.0.11-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:a6b13baf6bfcf881b6d6ac6e23c776f87a68304cd86e53d1d6b9afa31e363c4e", size = 169883, upload-time = "2025-12-15T22:31:28.233Z" }, + { url = "https://files.pythonhosted.org/packages/24/c0/b970eefb78746c8b4f8b3fa1b49d9f3ec4c5429ef3c5d4bbcc55abebe478/pycocotools-2.0.11-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:78bae4a9de9d34c4759754a848dfb3306f9ef1c2fcb12164ffbd3d013d008321", size = 486894, upload-time = "2025-12-15T22:31:29.283Z" }, + { url = "https://files.pythonhosted.org/packages/5b/f7/db7436820a1948d96fa9764b6026103e808840979be01246049f2c1e7f94/pycocotools-2.0.11-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:83d896f4310379849dfcfa7893afb0ff21f4f3cdb04ab3f61b05dd98953dd0ad", size = 483249, upload-time = "2025-12-15T22:31:31.687Z" }, + { url = "https://files.pythonhosted.org/packages/1e/a6/a14a12c9f50c41998fdc0d31fd3755bcbce124bac9abb1d6b99d1853cafd/pycocotools-2.0.11-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:eebd723503a2eb2c8b285f56ea3be1d9f3875cd7c40d945358a428db94f14015", size = 469070, upload-time = "2025-12-15T22:31:32.821Z" }, + { url = "https://files.pythonhosted.org/packages/46/de/aa4f65ece3da8e89310a1be00cad0700170fd13f41a3aaae2712291269d5/pycocotools-2.0.11-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:bd7a1e19ef56a828a94bace673372071d334a9232cd32ae3cd48845a04d45c4f", size = 475589, upload-time = "2025-12-15T22:31:34.188Z" }, + { url = "https://files.pythonhosted.org/packages/44/6f/04a30df03ae6236b369b361df0c50531d173d03678978806aa2182e02d1e/pycocotools-2.0.11-cp314-cp314t-win_amd64.whl", hash = "sha256:63026e11a56211058d0e84e8263f74cbccd5e786fac18d83fd221ecb9819fcc7", size = 93863, upload-time = "2025-12-15T22:31:35.38Z" }, + { url = "https://files.pythonhosted.org/packages/da/05/8942b640d6307a21c3ede188e8c56f07bedf246fac0e501437dbda72a350/pycocotools-2.0.11-cp314-cp314t-win_arm64.whl", hash = "sha256:8cedb8ccb97ffe9ed2c8c259234fa69f4f1e8665afe3a02caf93f6ef2952c07f", size = 72038, upload-time = "2025-12-15T22:31:36.768Z" }, +] + [[package]] name = "pydot" version = "3.0.4" @@ -2237,6 +2293,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/88/21/afadd25ecd81b3cea1e11c73cf1ab41a983a50271548c3ec7ec3b9efc3e9/torch-2.12.0-cp314-cp314t-win_amd64.whl", hash = "sha256:5f96b63f8287f66a005dd1b5a6abba2920f11156c5e5c4d815f3e2050fd1aa16", size = 123231092, upload-time = "2026-05-13T14:51:18.854Z" }, ] +[[package]] +name = "torchmetrics" +version = "1.9.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "lightning-utilities" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "torch" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/81/34/39b8b749333db56c0585d7a11fa62a283c087bb1dfc897d69fb8cedbefb1/torchmetrics-1.9.0.tar.gz", hash = "sha256:a488609948600df52d3db4fcdab02e62aab2a85ef34da67037dc3e65b8512faa", size = 581765, upload-time = "2026-03-09T17:41:22.443Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/a2/c7f6ebf546f8f644edf0f999aa98ece106986a77a7b922316bf6414ff825/torchmetrics-1.9.0-py3-none-any.whl", hash = "sha256:bfdcbff3dd1d96b3374bb2496eb39f23c4b28b8a845b6a18c313688e0d2d9ca1", size = 983384, upload-time = "2026-03-09T17:41:19.756Z" }, +] + [[package]] name = "torchvision" version = "0.27.0"