Refactor ONNX Inference Logic

[r83575]

Refactored ONNX-related inference code to align with the new project structure.
No logic or functional changes — verified identical results pre/post refactor.

Changes

• Moved:
  • src/onnx_model.py → src/inference/base/classifier_inference_base_onnx.py
  • scripts/onnx_predict_images.py → src/inference/utils/onnx_predict_images.py
• Updated imports in:
  • scripts/onnx_validation.py
  • src/evaluation/run_evaluation.py
  • src/inference/utils/onnx_predict_images.py

Verification

• Outputs identical to pre-refactor version
• Imports and paths updated

Closes #120

[Sarah5567]

I think it would be nice to add the creation of the onnx_model to InferenceFactory, and update all instances where this object is created to use the factory.

[Sarah5567]

I think you should move this file into the onnx directory and rename it to onnx_inference.py (similar to mobilenet_inference.py and resnet_inference.py).

[r83575]

According to the task description, the ONNX file should stay under src/inference/base/ - moving it to onnx/ is out of scope for this ticket.

[lyuzinmaxim]

I like the fact that now the InferenceFactory supports onnx - that makes sense. However, I think the implementation could be more transparent.

Currently, the main method is def create(model_type: str, model_path: str, device: str, class_mapping=None, **kwargs):, and we have many if/elif/else statements. In general, the best practice is to use enums for that.

I asked chatgpt to refactor current factory using pydantic's BaseModel and Enums, and I actually like this version - it's transparent, clean and also scalable (imagine, adding 3rd backend like TRT).

from enum import Enum
from pydantic import BaseModel, Field, validator
from typing import Optional, Type, Union
from pathlib import Path

from src.inference.pytorch.resnet_inference import ResNetInference
from src.inference.pytorch.mobilenet_inference import MobileNetInference
from src.inference.base.classifier_inference_base_onnx import OnnxClassifierInferenceBase


class ModelArch(str, Enum):
    RESNET = "resnet"
    MOBILENET = "mobilenet"


class Backend(str, Enum):
    PYTORCH = "pytorch"
    ONNX = "onnx"


class InferenceConfig(BaseModel):
    arch: ModelArch = Field(..., description="Model architecture (resnet or mobilenet).")
    backend: Backend = Field(..., description="Backend (pytorch or onnx).")
    model_path: Union[str, Path] = Field(..., description="Model checkpoint or ONNX path.")
    device: str = Field(..., description="Device string (e.g. 'cuda', 'cpu').")
    class_mapping: Optional[dict] = None
    extra_args: dict = Field(default_factory=dict)

    @validator("backend", pre=True, always=True)
    def infer_backend(cls, v, values):
        """If backend not given, infer ONNX from file extension."""
        if v:
            return v
        model_path = Path(values.get("model_path", ""))
        if model_path.suffix == ".onnx":
            return Backend.ONNX
        return Backend.PYTORCH


class InferenceFactory:
    """Factory to create model inference instances from architecture/backend pairs."""

    _PYTORCH_CLASSES : dict[ModelArch, Type] = {
        ModelArch.RESNET: ResNetInference,
        ModelArch.MOBILENET: MobileNetInference,
    }

    _ONNX_CLASSES : Type = OnnxClassifierInferenceBase

    @staticmethod
    def create(config: InferenceConfig):
        if config.backend == Backend.ONNX:
            model_cls = InferenceFactory._ONNX_IMPL
        elif config.backend == Backend.PYTORCH:
            model_cls = InferenceFactory._PYTORCH_IMPLS.get(config.arch)
            if model_cls is None:
                raise ValueError(f"Unsupported PyTorch architecture: {config.arch}")
        else:
            raise ValueError(f"Unsupported backend: {config.backend}")

        return model_cls(
            device=config.device,
            weights_path=config.model_path,
            class_mapping=config.class_mapping,
            **config.extra_args,
        )

[lyuzinmaxim]

Please define the enums separately, e.g.

    arch = ModelArch.RESNET if "resnet" in model_type.lower() else ModelArch.MOBILENET
    backend = Backend.ONNX if model_path.endswith(".onnx") else Backend.PYTORCH

    config = InferenceConfig(
        arch=arch,
        backend=backend,
        model_path=Path(model_path),
        device=device,
        class_mapping=class_mapping,
    )