Upload feature_extractor.py

src/modules/feature_extractor.py

@@ -1,28 +1,60 @@
-
 import torchvision.models.feature_extraction
 import torchvision
 import os
 import torch
+import onnx
+import onnxruntime
+import numpy as np
 
 from .config_extractor import MODEL_CONFIG
 
 os.environ["KMP_DUPLICATE_LIB_OK"] = "True"
 
+
 class FeatureExtractor:
     """Class for extracting features from images using a pre-trained model"""
 
-    def __init__(self, base_model):
+    def __init__(self, base_model, onnx_path=None):
         # set the base model
         self.base_model = base_model
         # get the number of features
         self.feat_dims = MODEL_CONFIG[base_model]["feat_dims"]
-        # initialize the image transformations
-        self.model, self.transforms = self.init_model(base_model)
-        self.model.eval()  # set the model to evaluation mode
-        # move the model to MPS
-        self.device = torch.device("cpu")
-        print("device is", self.device)
-        self.model.to(self.device)
+        # get the feature layer name
+        self.feat_layer = MODEL_CONFIG[base_model]["feat_layer"]
+
+        # Set default ONNX path if not provided
+        if onnx_path is None:
+            onnx_path = f"model/{base_model}_feature_extractor.onnx"
+
+        self.onnx_path = onnx_path
+        self.onnx_session = None
+
+        # Initialize transforms (needed for both ONNX and PyTorch)
+        _, self.transforms = self.init_model(base_model)
+
+        # Check if ONNX model exists
+        if os.path.exists(onnx_path):
+            print(f"Loading existing ONNX model from {onnx_path}")
+            self.onnx_session = onnxruntime.InferenceSession(onnx_path)
+        else:
+            print(
+                f"ONNX model not found at {onnx_path}. Initializing PyTorch model and converting to ONNX..."
+            )
+            # Initialize PyTorch model
+            self.model, _ = self.init_model(base_model)
+            self.model.eval()
+            self.device = torch.device("cpu")
+            self.model.to(self.device)
+
+            # Create directory if it doesn't exist
+            os.makedirs(os.path.dirname(onnx_path), exist_ok=True)
+
+            # Convert to ONNX
+            self.convert_to_onnx(onnx_path)
+
+            # Load the newly created ONNX model
+            self.onnx_session = onnxruntime.InferenceSession(onnx_path)
+            print(f"Successfully created and loaded ONNX model from {onnx_path}")
 
     def init_model(self, base_model):
         """Initialize the model for feature extraction
@@ -60,8 +92,49 @@ class FeatureExtractor:
         x = self.transforms(img)
         # add batch dimension
         x = x.unsqueeze(0)
-        # move the tensor to MPS
-        x = x.to(self.device)
-        # output now has the features corresponding to input x
-        output = self.model(x)[MODEL_CONFIG[self.base_model]["feat_layer"]]
-        return output
+
+        # Convert to numpy for ONNX Runtime
+        x_numpy = x.numpy()
+        # Run inference with ONNX Runtime
+        print("Running inference with ONNX Runtime")
+        output = self.onnx_session.run(
+            None,
+            {'input': x_numpy}
+        )[0]
+        # Convert back to torch tensor
+        output = torch.from_numpy(output)
+
+        return output
+
+    def convert_to_onnx(self, save_path):
+        """Convert the model to ONNX format and save it
+
+        Args:
+            save_path: str, the path to save the ONNX model
+
+        Returns:
+            None
+        """
+        # Create a dummy input tensor
+        dummy_input = torch.randn(1, 3, 224, 224, device=self.device)
+
+        # Export the model
+        torch.onnx.export(
+            self.model,
+            dummy_input,
+            save_path,
+            export_params=True,
+            opset_version=14,
+            do_constant_folding=True,
+            input_names=['input'],
+            output_names=['output'],
+            dynamic_axes={
+                'input': {0: 'batch_size'},
+                'output': {0: 'batch_size'}
+            }
+        )
+
+        # Verify the exported model
+        onnx_model = onnx.load(save_path)
+        onnx.checker.check_model(onnx_model)
+        print(f"ONNX model saved to {save_path}")