Integrate Ultralytics Backend (OD + IS) (#6648)

Co-authored-by: Romanov, Evgeny <evgeny.romanov@intel.com>

Author: kprokofi

application/backend/app/execution/common/geti_config_converter.py

@@ -66,7 +66,7 @@ def convert_metrics(metrics: dict) -> dict[str, float]:
     """
     result: dict[str, float] = {}
     for k, v in metrics.items():
-        name = k.split("/")[1] if "/" in k else k
+        name = k.split("/", 1)[1] if "/" in k else k
         if isinstance(v, torch.Tensor):
             if v.numel() == 1:
                 result[name] = v.item()

application/backend/app/execution/common/getitune_converters.py

@@ -7,6 +7,7 @@
 from pathlib import Path
 from uuid import UUID
 
+import yaml
 from datumaro.experimental.fields import Subset
 from getitune.backend.openvino.engine import OVEngine
 from getitune.config.data import SamplerConfig, SubsetConfig
@@ -56,8 +57,12 @@ def __init__(self, quantization_deps: QuantizationDependencies) -> None:
         self._db_session_factory = quantization_deps.db_session_factory
 
     @step("Validate Model", 5)
-    def validate_model(self, params: QuantizationJobParams) -> ModelRevision:
-        """Verify the source model is valid for quantization."""
+    def validate_model(self, params: QuantizationJobParams) -> tuple[ModelRevision, ModelVariant]:
+        """Verify the source model is valid for quantization.
+
+        Returns:
+            Tuple of (model, openvino_fp16_variant).
+        """
         with self._db_session_factory() as db:
             self._model_service.set_db_session(db)
             model = self._model_service.get_model(project_id=params.project_id, model_id=params.model_id)
@@ -82,8 +87,18 @@ def validate_model(self, params: QuantizationJobParams) -> ModelRevision:
         if not ov_model_xml_path.exists():
             raise FileNotFoundError(f"OpenVINO model files not found at {ov_model_xml_path}")
 
+        # Validate metadata.yaml structure if it exists (optional — Lightning models don't produce it)
+        metadata_path = ov_model_xml_path.parent / "metadata.yaml"
+        if metadata_path.exists():
+            with open(metadata_path) as f:
+                metadata = yaml.safe_load(f)
+            if not isinstance(metadata, dict):
+                raise ValueError(f"metadata.yaml is malformed (not a dict) at {metadata_path}")
+            if "args" not in metadata or not isinstance(metadata["args"], dict):
+                raise ValueError(f"metadata.yaml is missing 'args' dict at {metadata_path}")
+
         logger.info("Model {} validated for quantization (architecture={})", model.id, model.architecture)
-        return model
+        return model, openvino_variant
 
     @step("Prepare Calibration Dataset", 20)
     def prepare_calibration_dataset(
@@ -259,6 +274,7 @@ def store_artifacts(
         params: QuantizationJobParams,
         quantized_model_path: Path,
         model_variant_id: UUID,
+        source_variant_id: UUID,
     ) -> None:
         """Store quantized model files.
 
@@ -274,11 +290,25 @@ def store_artifacts(
         bin_path = quantized_model_path.with_suffix(".bin")
         if bin_path.exists():
             shutil.copyfile(bin_path, variant_dir / "model.bin")
+
+        # Copy metadata.yaml from the source FP16 variant and update for INT8
+        source_variant_dir = model_dir / "variants" / str(source_variant_id)
+        source_metadata_path = source_variant_dir / "metadata.yaml"
+        if source_metadata_path.exists():
+            with open(source_metadata_path) as f:
+                metadata = yaml.safe_load(f)
+            metadata["args"]["int8"] = True
+            metadata["args"]["half"] = False
+            dest_metadata_path = variant_dir / "metadata.yaml"
+            with open(dest_metadata_path, "w") as f:
+                yaml.dump(metadata, f, default_flow_style=False, sort_keys=False)
+            logger.info("Wrote INT8 metadata.yaml to {}", dest_metadata_path)
+
         logger.info("Stored quantized model variant at {}", variant_dir)
 
     def execute(self, params: QuantizationJobParams) -> None:
         """Execute the full quantization pipeline."""
-        model = self.validate_model(params=params)
+        model, openvino_fp16_variant = self.validate_model(params=params)
 
         with self._db_session_factory() as db:
             self._project_service.set_db_session(db)
@@ -323,6 +353,7 @@ def execute(self, params: QuantizationJobParams) -> None:
             params=params,
             quantized_model_path=quantized_model_path,
             model_variant_id=variant.id,
+            source_variant_id=openvino_fp16_variant.id,
         )
 
     def _base_model_path(self, project_id: UUID, model_id: UUID) -> Path:

application/backend/app/execution/quantization/getitune_quantizer.py

@@ -7,22 +7,24 @@
 from dataclasses import dataclass
 from datetime import UTC, datetime
 from pathlib import Path
-from typing import cast
+from typing import Any, cast
 from uuid import UUID
 
 import polars as pl
 import yaml
 from datumaro.experimental import Dataset
 from datumaro.experimental.fields import Subset
 from getitune import TaskType
-from getitune.backend.lightning.engine import LightningEngine
 from getitune.backend.lightning.models.base import DataInputParams, LightningModel
 from getitune.backend.openvino.engine import OVEngine
+from getitune.backend.ultralytics.models.base import UltralyticsModel
 from getitune.config.data import SamplerConfig, SubsetConfig
 from getitune.data.dataset.base import VisionDataset
 from getitune.data.entity.utils import detect_storage_dtype
 from getitune.data.factory import TransformLibFactory
 from getitune.data.module import DataModule
+from getitune.engine import create_engine
+from getitune.engine.engine import Engine
 from getitune.types.device import DeviceType as GetiTuneDeviceType
 from getitune.types.export import ExportFormat
 from getitune.types.precision import Precision
@@ -171,7 +173,12 @@ def prepare_weights(self, training_params: TrainingJobParams) -> Path:
                     "Can't start training - the parent revision has no variants (it may have failed). "
                     "Review the previous revision and retry."
                 )
-            parent_pytorch_variant = next(v for v in parent_variants if v.format == ModelFormat.PYTORCH)
+            parent_pytorch_variant = next((v for v in parent_variants if v.format == ModelFormat.PYTORCH), None)
+            if parent_pytorch_variant is None:
+                raise ExecutionErr(
+                    "Can't start training - the parent revision has no PyTorch variant. "
+                    "Review the previous revision and retry."
+                )
             weights_path = self.__build_model_weights_path(
                 self._data_dir, project_id, parent_model_revision_id, parent_pytorch_variant.id
             )
@@ -384,20 +391,17 @@ def train_model(
         weights_path: Path,
         model_id: UUID,
         device: DeviceInfo,
-        has_parent_revision: bool,
-    ) -> tuple[Path, LightningEngine]:
-        """Execute model training."""
-        # TODO use weights path to initialize model from pre-downloaded weights
-        #  after resolving https://github.com/open-edge-platform/training_extensions/issues/5100
-        logger.warning(
-            "Argument 'weights_path' (value='{}') is not used in model training yet; "
-            "the weights location will be determined internally by getitune",
-            weights_path,
-        )
-
-        # Build the DataModule
+    ) -> tuple[Path, Engine]:
+        """Execute model training.
+
+        Instantiates the model from *training_config* with jsonargparse and
+        creates the engine through ``create_engine`` from the library.
+        Lightning callbacks are built when a ``callbacks`` section is present
+        in the config; backends that do not support them simply ignore the
+        parameter.
+        """
         logger.info("Preparing the DataModule for training (model_id={})", model_id)
-        getitune_datamodule = DataModule.from_vision_datasets(
+        datamodule = DataModule.from_vision_datasets(
             train_dataset=dataset_info.getitune_training_dataset,
             val_dataset=dataset_info.getitune_validation_dataset,
             test_dataset=dataset_info.getitune_testing_dataset,
@@ -406,36 +410,36 @@ def train_model(
             test_subset=dataset_info.getitune_testing_subset_config,
         )
 
-        # Create the LightningModel according to the training configuration
-        logger.info("Instantiating the LightningModel for training (model_id={})", model_id)
+        logger.info("Instantiating model for training (model_id={})", model_id)
         model_cfg = training_config["model"]
-        model_cfg["init_args"]["label_info"] = getitune_datamodule.label_info.label_names
+        model_cfg["init_args"]["label_info"] = datamodule.label_info.label_names
         model_cfg["init_args"]["data_input_params"] = DataInputParams(
-            input_size=cast(tuple[int, int], getitune_datamodule.input_size),
-            mean=getitune_datamodule.input_mean if getitune_datamodule.input_mean is not None else (0.0, 0.0, 0.0),
-            std=getitune_datamodule.input_std if getitune_datamodule.input_std is not None else (1.0, 1.0, 1.0),
-            intensity_config=getitune_datamodule.input_intensity_config,
+            input_size=cast(tuple[int, int], datamodule.input_size),
+            mean=datamodule.input_mean if datamodule.input_mean is not None else (0.0, 0.0, 0.0),
+            std=datamodule.input_std if datamodule.input_std is not None else (1.0, 1.0, 1.0),
+            intensity_config=datamodule.input_intensity_config,
         ).as_dict()
-        model_parser = ArgumentParser()
-        model_parser.add_subclass_arguments(LightningModel, "model", required=False, fail_untyped=False)
-        getitune_model: LightningModel = model_parser.instantiate_classes(Namespace(model=model_cfg)).get("model")
-        if hasattr(getitune_model, "tile_config"):
-            getitune_model.tile_config = getitune_datamodule.tile_config
-
-        # Set up the LightningEngine
-        logger.info("Initializing the LightningEngine for training (model_id={})", model_id)
+        logger.info("Initializing engine for training (model_id={})", model_id)
         getitune_device_type = (
             GetiTuneDeviceType.gpu if device.type is DeviceType.CUDA else GetiTuneDeviceType(device.type)
         )
-        getitune_engine = LightningEngine(
-            model=getitune_model,
-            data=getitune_datamodule,
-            checkpoint=weights_path if has_parent_revision else None,
-            work_dir=self._data_dir / f"getitune-workspace-{model_id}",
-            device=getitune_device_type,
-        )
+        engine_kwargs: dict[str, Any] = {
+            "work_dir": self._data_dir / f"getitune-workspace-{model_id}",
+            "device": getitune_device_type,
+            "checkpoint": weights_path,
+        }
+
+        model_parser = ArgumentParser()
+        class_path = model_cfg.get("class_path", "")
+        model_type = UltralyticsModel if "ultralytics" in class_path else LightningModel
+        model_parser.add_argument("--model", type=model_type)
+        getitune_model = model_parser.instantiate_classes(Namespace(model=model_cfg)).get("model")
+
+        if hasattr(getitune_model, "tile_config"):
+            getitune_model.tile_config = datamodule.tile_config
+
+        getitune_engine = create_engine(model=getitune_model, data=datamodule, **engine_kwargs)
 
-        # Set up the callbacks
         callbacks_cfg = training_config.get("callbacks", [])
         for cb_cfg in callbacks_cfg:
             if "init_args" in cb_cfg and "dirpath" in cb_cfg["init_args"]:
@@ -446,25 +450,35 @@ def train_model(
         callbacks_list = parser.instantiate_classes(parsed_callbacks_cfg).get("callbacks", [])
         callbacks_list.append(TrainingProgressCallback(self.update_progress, min_p=10, max_p=80))
 
-        # Start training
-        logger.info("Starting the training loop (model_id={})", model_id)
-        train_kwargs = {
+        logger.info("Starting training loop (model_id={})", model_id)
+        train_kwargs: dict[str, Any] = {
             "max_epochs": training_config["max_epochs"],
             "callbacks": callbacks_list,
         }
-        if device.type is not DeviceType.CPU and device.index:
+        if device.type is not DeviceType.CPU and device.index is not None:
             train_kwargs["devices"] = [device.index]
         if "precision" in training_config:
             train_kwargs["precision"] = training_config["precision"]
+        # Forward backend-specific training args (e.g. patience for Ultralytics).
+        if "training" in training_config:
+            train_kwargs.update(training_config["training"])
         getitune_engine.train(**train_kwargs)  # pyrefly: ignore[bad-argument-type]
-        trained_model_path = Path(getitune_engine.work_dir) / "best_checkpoint.ckpt"
+
+        trained_model_path = getitune_engine.best_checkpoint
+        if trained_model_path is None:
+            raise RuntimeError(
+                f"Training completed but no best checkpoint was produced "
+                f"by the engine ({type(getitune_engine).__name__})."
+            )
+        if not trained_model_path.exists():
+            raise FileNotFoundError(f"Trained checkpoint not found at {trained_model_path}")
         logger.info("Model training completed. Trained model saved at {}", trained_model_path)
         return trained_model_path, getitune_engine
 
     @step("Evaluate Model", 95)
     def evaluate_model(
         self,
-        getitune_engine: LightningEngine,
+        getitune_engine: Engine,
         task: Task,
         model_revision_id: UUID,
         model_variants: list[ModelVariantDescriptor],
@@ -533,7 +547,7 @@ def _save_evaluation_result(
             )
 
     @step("Export Model")
-    def export_model(self, getitune_engine: LightningEngine, model_checkpoint_path: Path) -> ExportedModels:
+    def export_model(self, getitune_engine: Engine, model_checkpoint_path: Path) -> ExportedModels:
         """Export the trained model to desired OpenVINO and ONNX formats"""
         logger.info("Exporting the model to OpenVINO format (FP16 precision)...")
         exported_ov_model_path = getitune_engine.export(
@@ -578,10 +592,9 @@ def store_model_artifacts(
         variants_dir = model_dir / "variants"
         variants_dir.mkdir(parents=True, exist_ok=True)
 
-        # Copy PyTorch checkpoint
         pytorch_variant_dir = variants_dir / str(created_variants[ModelFormat.PYTORCH])
         pytorch_variant_dir.mkdir(parents=True, exist_ok=True)
-        shutil.copyfile(trained_model_path, pytorch_variant_dir / "model.ckpt")
+        shutil.copyfile(trained_model_path, pytorch_variant_dir / "model.pt")
         logger.info("Stored PyTorch variant at {}", pytorch_variant_dir)
 
         # Copy OpenVINO IR files
@@ -595,6 +608,9 @@ def store_model_artifacts(
             exported_model_paths.openvino_model_path.with_suffix(".bin"),
             openvino_variant_dir / "model.bin",
         )
+        ov_metadata = exported_model_paths.openvino_model_path.parent / "metadata.yaml"
+        if ov_metadata.exists():
+            shutil.copyfile(ov_metadata, openvino_variant_dir / "metadata.yaml")
         logger.info("Stored OpenVINO variant at {}", openvino_variant_dir)
 
         # Copy ONNX file
@@ -604,14 +620,16 @@ def store_model_artifacts(
             exported_model_paths.onnx_model_path.with_suffix(".onnx"),
             onnx_variant_dir / "model.onnx",
         )
+        onnx_metadata = exported_model_paths.onnx_model_path.parent / "metadata.yaml"
+        if onnx_metadata.exists():
+            shutil.copyfile(onnx_metadata, onnx_variant_dir / "metadata.yaml")
         logger.info("Stored ONNX variant at {}", onnx_variant_dir)
 
         # Store the metrics
         metrics_source_path = getitune_work_dir / "csv"
         metrics_dest_path = model_dir / "metrics"
         if metrics_source_path.exists():
             shutil.move(metrics_source_path, metrics_dest_path)
-            logger.info("Stored training metrics at {}", metrics_dest_path)
 
     def create_model_variants(self, model_revision_id: UUID) -> dict[ModelFormat, UUID]:
         """Create variant records in the database for all exported formats.
@@ -680,7 +698,6 @@ def execute(self, params: TrainingJobParams) -> None:
                 weights_path=weights_path,
                 model_id=params.model_id,
                 device=params.device,
-                has_parent_revision=params.parent_model_revision_id is not None,
             )
             exported_model_paths = self.export_model(
                 getitune_engine=getitune_engine, model_checkpoint_path=trained_model_path
@@ -849,9 +866,10 @@ def __base_model_path(data_dir: Path, project_id: UUID, model_id: UUID) -> Path:
 
     @classmethod
     def __build_model_weights_path(cls, data_dir: Path, project_id: UUID, model_id: UUID, model_variant: UUID) -> Path:
-        """Get the path to the PyTorch checkpoint from a model's variants directory."""
+        """Get the path to the stored PyTorch checkpoint."""
         model_dir = cls.__base_model_path(data_dir, project_id, model_id)
-        return model_dir / "variants" / str(model_variant) / "model.ckpt"
+        variant_dir = model_dir / "variants" / str(model_variant)
+        return variant_dir / "model.pt"
 
     @classmethod
     def __build_model_config_path(cls, data_dir: Path, project_id: UUID, model_id: UUID) -> Path:

application/backend/app/execution/training/getitune_trainer.py

@@ -371,19 +371,20 @@ class Mosaic(BaseAugmentationParameter):
         default=0.5,
         title="Scale",
         description=(
-            "Scale factor for random perspective crop applied after mosaic assembly. "
-            "The random scale is sampled from [1-scale, 1+scale]. "
-            "Higher values produce more zoom variation."
+            "Scale factor for the random perspective crop applied after mosaic assembly. "
+            "The random scale is sampled from [1 - scale, 1 + scale]. "
+            "Higher values produce more aggressive zoom variation."
         ),
     )
     translate: float = Field(
         ge=0.0,
-        le=0.5,
+        le=1.0,
         default=0.1,
         title="Translate",
         description=(
-            "Translation fraction for random perspective crop. "
-            "The crop center shifts by ±translate * image_size pixels."
+            "Translate fraction for the random perspective crop. "
+            "The crop centre shifts by up to +/- translate * image_size pixels. "
+            "Higher values allow more positional variation."
         ),
     )
 

application/backend/app/models/training_configuration/augmentation.py

@@ -52,8 +52,7 @@ def _find_project_label(
     label_name: str | None,
     predicted_class_index: int | None = None,
 ) -> Label | None:
-    """
-    Find a project label matching the predicted label using a three-stage fallback strategy.
+    """Find a project label matching the predicted label using a three-stage fallback strategy.
 
     Resolution order:
     1. **Exact name match** - performs an O(1) lookup of ``label_name`` in ``name_map``.

application/backend/app/services/data_collect/prediction_converter.py

@@ -318,7 +318,7 @@ def main() -> None:
 
 dependencies = [
     "openvino~=2026.1.0",
-    "openvino-model-api[onnx]~=0.4.3",
+    "openvino-model-api[onnx]~=0.4.5",
     "opencv-python-headless~=4.13.0",
     "numpy>=2.0",
     "pillow~=12.0",

application/backend/app/services/demo_files_service.py

@@ -116,7 +116,7 @@ def set_inference_model(
                         resource_type=ResourceType.MODEL,
                         resource_id=f"{model_id} variant {model_variant.id}",
                     )
-                model_xml_path, _ = paths
+                model_xml_path = paths[0]
 
                 if self._loaded_model is not None:
                     ModelLoader.unload(self._loaded_model)