Refactor models conversions

docs/deepstream_validation.md

@@ -69,7 +69,7 @@ The models are **validated and ready for performance benchmarking**.
 - **ONNX predictions:** `/workspace/assets/predictions/onnx/`  
 - **Configs used:** `configs/resnet18.txt`, `configs/mobilenet.txt`  
 - **Engines:** `/workspace/models/deepstream/`  
-- **Script used:** `src/validate_deepstream_vs_onnx.py`  
+- **Script used:** `src/model_conversion/validate_trt_vs_onnx.py`  
 - **Execution environments:**
   - `run_dev.sh` (x86_64) → ONNX inference  
   - `run_dev_jetson.sh` (ARM Jetson) → DeepStream inference  

docs/export_to_onnx.md

@@ -8,10 +8,11 @@ All conversions should be executed **inside the dev container** to ensure depend
 ### Commands
 Export ResNet18:
 ```bash
-python scripts/export_to_onnx.py resnet18
+python src/model_conversion/pytorch_to_onnx.py resnet18
+
 ```
 
 Export MobileNetV2:
 ```bash
-python scripts/export_to_onnx.py mobilenet
+python src/model_conversion/pytorch_to_onnx.py mobilenet
 ```

src/deepstream_pipeline.py

@@ -1,12 +1,16 @@
-
-import gi
-
-gi.require_version("Gst", "1.0")
 import os
+import sys
 from datetime import datetime
 
+import gi
 from gi.repository import GLib, Gst
-from engine_helper import ensure_engine_exists
+
+# Add workspace root to sys.path
+sys.path.append("/workspace")
+
+from src.model_conversion.onnx_to_trt import build_engine_if_missing
+
+gi.require_version("Gst", "1.0")
 
 
 rtsp_port = os.environ.get("RTSP_PORT", "8554")
@@ -133,7 +137,7 @@ def on_message(bus, msg):
 
 
 def main():
-    ensure_engine_exists(CONFIG_FILE)
+    build_engine_if_missing(CONFIG_FILE)
     # OS-level setup
     Gst.init(None)
     os.makedirs(OUTPUT_DIR, exist_ok=True)

src/engine_helper.py → src/model_conversion/onnx_to_trt.py

@@ -10,7 +10,7 @@ def file_sha256(path: str) -> str:
         return hashlib.sha256(f.read()).hexdigest()
 
 
-def ensure_engine_exists(config_path: str) -> str:
+def build_engine_if_missing(config_path: str) -> str:
     """
     Ensure that a TensorRT engine file exists based on a DeepStream config file.
     If not found, create it using trtexec.

src/model_conversion/pytorch_to_onnx.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+import argparse
+from pathlib import Path
+import subprocess
+
+import torch
+
+from src.inference import MobileNetInference, ResNetInference
+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)
+
+    # Detect parameter dtype (fp16/fp32) and match input accordingly
+    param_dtype = next(
+        (p.dtype for p in pytorch_model.model.parameters() if p.is_floating_point()),
+        torch.float32,
+    )
+
+    # Fixed input size: 256x256, 3 channels, batch=1
+    dummy = torch.zeros(1, 3, INPUT_SIZE[0], INPUT_SIZE[1], dtype=param_dtype)
+
+    # ONNX path/name
+    onnx_dir = MODELS_DIR_PATH / "onnx"
+    ensure_clean_directory(onnx_dir)
+    onnx_path = onnx_dir / f"{model_name}.onnx"
+    dynamic_axes = {"input": {0: "batch"}, "output": {0: "batch"}}
+
+    with torch.inference_mode():
+        torch.onnx.export(
+            pytorch_model.model,
+            dummy,
+            onnx_path.as_posix(),
+            input_names=["input"],
+            output_names=["output"],
+            dynamic_axes=dynamic_axes,
+            opset_version=17,
+            do_constant_folding=True,
+        )
+
+    subprocess.run(["python3", "src/model_conversion/validate_onnx_vs_pytorch.py", model_name], check=True)
+
+    print(f"Exported: {onnx_path}")
+
+
+if __name__ == "__main__":
+    p = argparse.ArgumentParser("Export PyTorch model to ONNX (fixed 256x256 input)")
+    p.add_argument(
+        "model_name",
+        type=str,
+        choices=["resnet18", "mobilenet"],
+        help="Name of the model to export (must exist in models directory)",
+    )
+    args = p.parse_args()
+
+    main(args.model_name)

src/model_conversion/validate_onnx_vs_pytorch.py

@@ -0,0 +1,116 @@
+#!/usr/bin/env python3
+import argparse
+import os
+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.onnx_model import OnnxClassifierInferenceBase
+
+NUM_CLASSES = 83
+INPUT_SIZE = (256, 256)
+
+MODELS_DIR_PATH = Path("models")
+IMAGES_DIR_PATH = Path("assets/test_images")
+DEVICE = "cpu"
+INPUT_SHAPE = (1, 3, *INPUT_SIZE)
+
+
+def assert_the_same_shapes(
+    model_name: str,
+    onnx_clf: OnnxClassifierInferenceBase,
+):
+    onnx_inputs = [input.shape for input in onnx_clf.session.get_inputs()]
+    onnx_outputs = [output.shape for output in onnx_clf.session.get_outputs()]
+
+    assert onnx_inputs == [["batch", 3, *INPUT_SIZE]], "ONNX input shape mismatch"
+    assert onnx_outputs == [
+        ["batch", NUM_CLASSES]
+    ], f"Number of outputs mismatch: {onnx_outputs} is not [['batch', {NUM_CLASSES}]]"
+
+
+def assert_numerical_accuracy(
+    model_name: str,
+    torch_clf: ClassifierInferenceBase,
+    onnx_clf: OnnxClassifierInferenceBase,
+    rtol: float = 1e-3,
+    atol: float = 1e-5,
+):
+    with torch.inference_mode():
+        param_dtype = next((p.dtype for p in torch_clf.model.parameters() if p.is_floating_point()), torch.float32)
+        torch_input = torch.randn(*INPUT_SHAPE, dtype=param_dtype)
+        torch_logits = torch_clf.model(torch_input).to(torch.float32)
+        torch_np = torch_logits.detach().cpu().numpy()
+
+        onnx_logits = onnx_clf._forward(torch_input.to(torch.float32))
+        onnx_np = onnx_logits.detach().cpu().numpy()
+
+    assert np.allclose(torch_np, onnx_np, rtol=rtol, atol=atol)
+
+
+def assert_same_predictions(
+    model_name: str,
+    torch_clf: ClassifierInferenceBase,
+    onnx_clf: OnnxClassifierInferenceBase,
+):
+    torch_results = {}
+    onnx_results = {}
+    img_extensions = {".jpg", ".jpeg", ".png", ".bmp"}
+
+    for filename in os.listdir(IMAGES_DIR_PATH):
+        file_ext = os.path.splitext(filename)[1].lower()
+        if file_ext in img_extensions:
+            image_path = os.path.join(IMAGES_DIR_PATH, filename)
+            image = Image.open(image_path)
+
+            torch_pred = torch_clf.infer(image)
+            onnx_pred = onnx_clf.infer(image)
+
+            torch_results[filename] = torch_pred[0]["class_id"]
+            onnx_results[filename] = onnx_pred[0]["class_id"]
+
+    assert torch_results == onnx_results
+
+
+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}")
+    onnx_model = OnnxClassifierInferenceBase(
+        device=DEVICE,
+        weights_path=MODELS_DIR_PATH / "onnx" / f"{model_name}.onnx",
+        class_mapping=ID_TO_NAME,
+        topk=1,
+    )
+
+    assert_the_same_shapes(model_name, onnx_model)
+    assert_numerical_accuracy(model_name, torch_model, onnx_model)
+    assert_same_predictions(model_name, torch_model, onnx_model)
+
+
+if __name__ == "__main__":
+    p = argparse.ArgumentParser("Validate ONNX model against PyTorch")
+    p.add_argument(
+        "model_name",
+        type=str,
+        choices=["resnet18", "mobilenet"],
+        help="Name of the model to validate",
+    )
+    args = p.parse_args()
+    main(args.model_name)

src/test_onnx_pipeline.py

@@ -8,11 +8,15 @@
 
 gi.require_version("Gst", "1.0")
 from gi.repository import Gst  # noqa: E402
-from engine_helper import ensure_engine_exists
+
+# Add workspace root to sys.path
+sys.path.append("/workspace")
+
+from src.model_conversion.onnx_to_trt import build_engine_if_missing
 
 def main(config_file):
     # Create deepstream_engines directory with full permissions
-    ensure_engine_exists(config_file)
+    build_engine_if_missing(config_file)
 
     Gst.init(None)
 

src/test_onnx_pipeline_real.py

@@ -7,10 +7,14 @@
 
 gi.require_version("Gst", "1.0")
 from gi.repository import GLib, Gst  # noqa: E402
-from engine_helper import ensure_engine_exists
+
+# Add workspace root to sys.path
+sys.path.append("/workspace")
+
+from src.model_conversion.onnx_to_trt import build_engine_if_missing
 
 def main(config_file, input_path, is_video=True):
-    ensure_engine_exists(config_file)
+    build_engine_if_missing(config_file)
 
     Gst.init(None)
     pipeline = Gst.Pipeline()