Update app.py

app.py

@@ -1,14 +1,10 @@
 import gradio as gr
-import spaces
 import numpy as np
-import torch
-import torch.nn.functional as F
 import onnxruntime
 import cv2
-from PIL import Image
 
 # Declare ONNX session as a global variable
-MODEL_PATH = "weights/Glint360K_R200_TopoFR_9784.onnx"
+MODEL_PATH = "Glint360K_R200_TopoFR_9784.onnx"
 session = onnxruntime.InferenceSession(MODEL_PATH)
 
 def pil_to_cv2(pil_image):
@@ -28,8 +24,9 @@ def process_image(pil_img):
 
     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
     img = np.transpose(img, (2, 0, 1))
-    img = torch.from_numpy(img).unsqueeze(0).float()
-    img.div_(255).sub_(0.5).div_(0.5)
+    img = img.astype(np.float32)
+    img = np.expand_dims(img, axis=0)
+    img = (img / 255.0 - 0.5) / 0.5
     return img
 
 def calculate_similarity(img1, img2):
@@ -40,19 +37,19 @@ def calculate_similarity(img1, img2):
     # Extract features using ONNX model
     def get_features(img_tensor):
         input_name = session.get_inputs()[0].name
-        features = session.run(None, {input_name: img_tensor.numpy()})[0]
-        return torch.from_numpy(features)
+        features = session.run(None, {input_name: img_tensor})[0]
+        return features
     
     # Extract features for each image
     feat1 = get_features(img1_tensor)
     feat2 = get_features(img2_tensor)
     
-    # Normalize features
-    feat1 = F.normalize(feat1, p=2, dim=1)
-    feat2 = F.normalize(feat2, p=2, dim=1)
+    # Normalize features (L2 normalization)
+    feat1 = feat1 / np.linalg.norm(feat1, axis=1, keepdims=True)
+    feat2 = feat2 / np.linalg.norm(feat2, axis=1, keepdims=True)
     
     # Calculate cosine similarity
-    cosine_similarity = torch.sum(feat1 * feat2, dim=1).item()
+    cosine_similarity = np.sum(feat1 * feat2, axis=1).item()
     return f"Cosine Similarity: {cosine_similarity:.4f}"
 
 # Create Gradio interface with custom layout
@@ -75,4 +72,4 @@ with gr.Blocks() as iface:
     )
 
 # Launch the interface
-iface.launch(debug=True)
+iface.launch()