Upload 3 files

app.py

@@ -39,19 +39,24 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# Initialize paths
 index_path = "./model/db_vit_b_16.index"
+onnx_path = "./model/vit_b_16_feature_extractor.onnx"
 
+# Check if index file exists
 if not os.path.exists(index_path):
     raise FileNotFoundError(f"Index file not found: {index_path}")
 
 try:
+    # Load FAISS index
     index = faiss.read_index(index_path)
-except RuntimeError as e:
-    raise RuntimeError(f"Error reading FAISS index: {e}")
-feature_extractor = FeatureExtractor(base_model="vit_b_16")
+    print(f"Successfully loaded FAISS index from {index_path}")
 
-if torch.backends.mps.is_built():
-    torch.set_default_device("mps")
+    # Initialize feature extractor with ONNX support
+    feature_extractor = FeatureExtractor(base_model="vit_b_16", onnx_path=onnx_path)
+    print("Successfully initialized feature extractor with ONNX support")
+except Exception as e:
+    raise RuntimeError(f"Error initializing models: {str(e)}")
 
 
 def base64_to_image(base64_str: str) -> Image.Image:
@@ -69,12 +74,6 @@ def image_to_base64(image: Image.Image) -> str:
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
 
-def image_to_base64(image: Image.Image) -> str:
-    buffered = BytesIO()
-    image.save(buffered, format="JPEG")
-    return base64.b64encode(buffered.getvalue()).decode("utf-8")
-
-
 def unzip_folder(zip_file_path, extract_to_path):
     if not os.path.exists(zip_file_path):
         raise FileNotFoundError(f"Zip file not found: {zip_file_path}")
@@ -105,20 +104,26 @@ class ImageSearchBody(BaseModel):
 @app.post("/search-image/")
 async def search_image(body: ImageSearchBody):
     try:
+        # Convert base64 to image
         image = base64_to_image(body.base64_image)
-        with torch.no_grad():
-            output = feature_extractor.extract_features(image)
-            output = output.view(output.size(0), -1)
-            output = output / output.norm(p=2, dim=1, keepdim=True)
-            D, I = index.search(output.cpu().numpy(), 1)
-        print(D, I)
+
+        # Extract features using ONNX model
+        output = feature_extractor.extract_features(image)
+
+        # Prepare features for FAISS search
+        output = output.view(output.size(0), -1)
+        output = output / output.norm(p=2, dim=1, keepdim=True)
+
+        # Search for similar images
+        D, I = index.search(output.cpu().numpy(), 1)
+
+        # Get the matched image
         image_list = sorted([f for f in os.listdir(extract_path) if is_image_file(f)])
-        print(image_list)
         image_name = image_list[int(I[0][0])]
-        matched_image_path = f"{extract_path}/{image_list[int(I[0][0])]}"
+        matched_image_path = f"{extract_path}/{image_name}"
         matched_image = Image.open(matched_image_path)
         matched_image_base64 = image_to_base64(matched_image)
-    
+
         return JSONResponse(
             content={
                 "image_base64": matched_image_base64,
@@ -129,11 +134,15 @@ async def search_image(body: ImageSearchBody):
         )
 
     except Exception as e:
-        return JSONResponse(content={"error": str(e)}, status_code=500)
+        print(f"Error in search_image: {str(e)}")
+        return JSONResponse(
+            content={"error": f"Error processing image: {str(e)}"}, status_code=500
+        )
 
 
 from src.firebase.firebase_provider import process_images
 
+
 class Body(BaseModel):
     base64_image: list[str] = Field(..., title="Base64 Image String")
     model_config = {

feature_extractor.py

@@ -0,0 +1,140 @@
+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, 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"]
+        # 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
+
+        Args:
+            base_model: str, the name of the base model
+
+        Returns:
+            model: torch.nn.Module, the feature extraction model
+            transforms: torchvision.transforms.Compose, the image transformations
+        """
+        if base_model not in MODEL_CONFIG:
+            raise ValueError(f"Invalid base model: {base_model}")
+
+        # get the model and weights
+        weights = MODEL_CONFIG[base_model]["weights"]
+        model = torchvision.models.feature_extraction.create_feature_extractor(
+            MODEL_CONFIG[base_model]["model"](weights=weights),
+            [MODEL_CONFIG[base_model]["feat_layer"]],
+        )
+        # get the image transformations
+        transforms = weights.transforms()
+        return model, transforms
+
+    def extract_features(self, img):
+        """Extract features from an image
+
+        Args:
+            img: PIL.Image, the input image
+
+        Returns:
+            output: torch.Tensor, the extracted features
+        """
+        # apply transformations
+        x = self.transforms(img)
+        # add batch dimension
+        x = x.unsqueeze(0)
+
+        # 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}")

requirements.txt

@@ -10,4 +10,6 @@ python-multipart
 firebase-admin
 python-dotenv
 aiofiles
-pytz
+pytz
+onnx
+onnxruntime