Refactor pytorch logic

scripts/export_to_onnx.py

@@ -1,37 +1,27 @@
-#!/usr/bin/env python3
+import os
+import sys
 import argparse
 from pathlib import Path
 
 import torch
 
-from src.inference import MobileNetInference, ResNetInference
+# Add workspace root to sys.path
+sys.path.append("/workspace")
+
+from src.inference.utils.inference_factory import InferenceFactory
 from src.path_utils import ensure_clean_directory
 
 NUM_CLASSES = 83
 INPUT_SIZE = (256, 256)
 DEVICE = "cpu"
 MODELS_DIR_PATH = Path("models")
 
-
-def get_model(model_name: str):
-    if model_name == "resnet18":
-        return ResNetInference(
-            device=DEVICE,
-            weights_path=MODELS_DIR_PATH / "pytorch" / f"{model_name}.pt",
-            num_classes=NUM_CLASSES,
-        )
-    elif model_name == "mobilenet":
-        return MobileNetInference(
-            device=DEVICE,
-            weights_path=MODELS_DIR_PATH / "pytorch" / f"{model_name}.pt",
-            num_classes=NUM_CLASSES,
-        )
-    else:
-        raise ValueError(f"Unsupported model: {model_name}")
-
-
 def main(model_name: str):
-    pytorch_model = get_model(model_name)
+    pytorch_model = InferenceFactory.create(
+        model_type=model_name,
+        model_path=MODELS_DIR_PATH / "pytorch" / f"{model_name}.pt",
+        device=DEVICE,
+    )
 
     # Detect parameter dtype (fp16/fp32) and match input accordingly
     param_dtype = next(

scripts/onnx_validation.py

@@ -1,15 +1,17 @@
-#!/usr/bin/env python3
-import argparse
 import os
+import argparse
 from pathlib import Path
 
 import numpy as np
 import torch
 from PIL import Image
 
-from src.inference import ID_TO_NAME, ClassifierInferenceBase, MobileNetInference, ResNetInference
+from src.optimal_class_mapping import MODEL_NAMES as ID_TO_NAME
+from src.inference.utils.inference_factory import InferenceFactory
+from src.inference.base.classifier_inference_base import ClassifierInferenceBase
 from src.onnx_model import OnnxClassifierInferenceBase
 
+
 NUM_CLASSES = 83
 INPUT_SIZE = (256, 256)
 
@@ -76,22 +78,13 @@ def assert_same_predictions(
 
 
 def main(model_name: str):
-    if model_name == "resnet18":
-        torch_model = ResNetInference(
-            device=DEVICE,
-            weights_path=MODELS_DIR_PATH / "pytorch" / f"{model_name}.pt",
-            num_classes=NUM_CLASSES,
-            strict=True,
-        )
-    elif model_name == "mobilenet":
-        torch_model = MobileNetInference(
-            device=DEVICE,
-            weights_path=MODELS_DIR_PATH / "pytorch" / f"{model_name}.pt",
-            num_classes=NUM_CLASSES,
-            strict=True,
-        )
-    else:
-        raise ValueError(f"Unsupported model: {model_name}")
+    torch_model = InferenceFactory.create(
+        model_type=model_name,
+        model_path=MODELS_DIR_PATH / "pytorch" / f"{model_name}.pt",
+        device=DEVICE,
+        class_mapping=ID_TO_NAME,
+    )
+
     onnx_model = OnnxClassifierInferenceBase(
         device=DEVICE,
         weights_path=MODELS_DIR_PATH / "onnx" / f"{model_name}.onnx",

src/evaluation/evaluator.py

@@ -1,14 +1,18 @@
+import os
+import sys
 import time
 from dataclasses import dataclass
 from pathlib import Path
 from typing import List, Optional, Tuple
 
+# Add workspace root to sys.path
+sys.path.append("/workspace")
+
 import torch
 from PIL import Image
 
+from src.inference.base.classifier_inference_base import ClassifierInferenceBase
 from metrics.metrics_api import ClassificationReport, compute_metrics
-from src.inference import ClassifierInferenceBase as InferenceModel
-
 
 @dataclass
 class EvaluationConfig:
@@ -37,7 +41,7 @@ def load_dataset(self) -> Tuple[List[Path], List[int]]:
         print(f"✅ Loaded {len(image_paths)} images from {len(set(labels))} classes")
         return image_paths, labels
 
-    def evaluate_model(self, model: InferenceModel) -> ClassificationReport:
+    def evaluate_model(self, model: ClassifierInferenceBase) -> ClassificationReport:
         print("🔄 Starting evaluation...")
         image_paths, true_labels = self.load_dataset()
         predictions = []

src/evaluation/run_evaluation.py

@@ -1,18 +1,23 @@
+import os
+import sys
 import argparse
 import json
 from pathlib import Path
 
 import pandas as pd
 
+# Add workspace root to sys.path
+sys.path.append("/workspace")
+
+from src.optimal_class_mapping import MODEL_NAMES as class_mapping, map_prediction
+from src.inference.utils.inference_factory import InferenceFactory
 from src.evaluation.evaluator import EvaluationConfig, ModelEvaluator
-from src.inference import MobileNetInference, ResNetInference
 from src.onnx_model import OnnxClassifierInferenceBase as OnnxModel
 from dataset.optimal_class_mapping import map_prediction
 from src.path_utils import ensure_clean_directory
 
-MODELS_DIR_PATH = Path("models")
-NUM_CLASSES = 83
 
+MODELS_DIR_PATH = Path("models")
 
 class MappedModelWrapper:
     """Wrapper that adds mapping between 83 model classes to 76 dataset classes"""
@@ -34,21 +39,12 @@ def infer(self, image):
 
 
 def create_model(model_name: str, model_type: str, device: str):
-    weights_path = MODELS_DIR_PATH / model_type / f"{model_name}.pt"
+    weights_path = MODELS_DIR_PATH / model_type / f"{model_name}"
 
     if model_type == "pytorch":
-        if model_name == "resnet18":
-            base_model = ResNetInference(
-                device=device, weights_path=weights_path, num_classes=NUM_CLASSES
-            )
-        elif model_name == "mobilenet":
-            base_model = MobileNetInference(
-                device=device, weights_path=weights_path, num_classes=NUM_CLASSES
-            )
-        else:
-            raise ValueError(f"Unknown model: {model_name}")
+        base_model = InferenceFactory.create(model_name, weights_path.with_suffix(".pt"), device, class_mapping)
     elif model_type == "onnx":
-        base_model = OnnxModel(device=device, weights_path=weights_path, topk=1)
+        base_model = OnnxModel(device=device, weights_path=weights_path.with_suffix(".onnx"), topk=1)
     else:
         raise ValueError(f"Unknown model_type: {model_type}")
 

src/evaluation/run_hierarchical_evaluation.py

@@ -1,14 +1,19 @@
+import os
+import sys
 import json
 import time
 from pathlib import Path
 
-from PIL import Image
+# Add workspace root to sys.path
+sys.path.append("/workspace")
 
+from PIL import Image
 from dataset.utilities.datasets import DATASETS
 from metrics.metrics_api import compute_metrics
+from dataset_preparation.utilities.datasets import DATASETS
 from src.evaluation.evaluator import EvaluationConfig, ModelEvaluator
-from src.inference import MobileNetInference, ResNetInference
-from dataset.optimal_class_mapping import map_prediction
+from src.inference.utils.inference_factory import InferenceFactory
+from src.optimal_class_mapping import MODEL_NAMES as class_mapping, map_prediction
 
 
 class MappedModelWrapper:
@@ -66,7 +71,7 @@ def main():
 
     # ResNet
     print("🔄 ResNet...")
-    resnet = ResNetInference(weights_path="models/pytorch/resnet18.pt", num_classes=83)
+    resnet = InferenceFactory.create("resnet", "models/pytorch/resnet18.pt", "cpu")
     wrapped_resnet = MappedModelWrapper(resnet)
 
     predictions, latencies = [], []
@@ -83,7 +88,7 @@ def main():
 
     # MobileNet
     print("🔄 MobileNet...")
-    mobilenet = MobileNetInference(weights_path="models/pytorch/mobilenet.pt", num_classes=83)
+    mobilenet = InferenceFactory.create("mobilenet", "models/pytorch/mobilenet.pt", "cpu")
     wrapped_mobilenet = MappedModelWrapper(mobilenet)
 
     predictions, latencies = [], []