developer0hye commited on
Commit
c29b42f
·
verified ·
1 Parent(s): e5b7056

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -14
app.py CHANGED
@@ -1,14 +1,10 @@
1
  import gradio as gr
2
- import spaces
3
  import numpy as np
4
- import torch
5
- import torch.nn.functional as F
6
  import onnxruntime
7
  import cv2
8
- from PIL import Image
9
 
10
  # Declare ONNX session as a global variable
11
- MODEL_PATH = "weights/Glint360K_R200_TopoFR_9784.onnx"
12
  session = onnxruntime.InferenceSession(MODEL_PATH)
13
 
14
  def pil_to_cv2(pil_image):
@@ -28,8 +24,9 @@ def process_image(pil_img):
28
 
29
  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
30
  img = np.transpose(img, (2, 0, 1))
31
- img = torch.from_numpy(img).unsqueeze(0).float()
32
- img.div_(255).sub_(0.5).div_(0.5)
 
33
  return img
34
 
35
  def calculate_similarity(img1, img2):
@@ -40,19 +37,19 @@ def calculate_similarity(img1, img2):
40
  # Extract features using ONNX model
41
  def get_features(img_tensor):
42
  input_name = session.get_inputs()[0].name
43
- features = session.run(None, {input_name: img_tensor.numpy()})[0]
44
- return torch.from_numpy(features)
45
 
46
  # Extract features for each image
47
  feat1 = get_features(img1_tensor)
48
  feat2 = get_features(img2_tensor)
49
 
50
- # Normalize features
51
- feat1 = F.normalize(feat1, p=2, dim=1)
52
- feat2 = F.normalize(feat2, p=2, dim=1)
53
 
54
  # Calculate cosine similarity
55
- cosine_similarity = torch.sum(feat1 * feat2, dim=1).item()
56
  return f"Cosine Similarity: {cosine_similarity:.4f}"
57
 
58
  # Create Gradio interface with custom layout
@@ -75,4 +72,4 @@ with gr.Blocks() as iface:
75
  )
76
 
77
  # Launch the interface
78
- iface.launch(debug=True)
 
1
  import gradio as gr
 
2
  import numpy as np
 
 
3
  import onnxruntime
4
  import cv2
 
5
 
6
  # Declare ONNX session as a global variable
7
+ MODEL_PATH = "Glint360K_R200_TopoFR_9784.onnx"
8
  session = onnxruntime.InferenceSession(MODEL_PATH)
9
 
10
  def pil_to_cv2(pil_image):
 
24
 
25
  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
26
  img = np.transpose(img, (2, 0, 1))
27
+ img = img.astype(np.float32)
28
+ img = np.expand_dims(img, axis=0)
29
+ img = (img / 255.0 - 0.5) / 0.5
30
  return img
31
 
32
  def calculate_similarity(img1, img2):
 
37
  # Extract features using ONNX model
38
  def get_features(img_tensor):
39
  input_name = session.get_inputs()[0].name
40
+ features = session.run(None, {input_name: img_tensor})[0]
41
+ return features
42
 
43
  # Extract features for each image
44
  feat1 = get_features(img1_tensor)
45
  feat2 = get_features(img2_tensor)
46
 
47
+ # Normalize features (L2 normalization)
48
+ feat1 = feat1 / np.linalg.norm(feat1, axis=1, keepdims=True)
49
+ feat2 = feat2 / np.linalg.norm(feat2, axis=1, keepdims=True)
50
 
51
  # Calculate cosine similarity
52
+ cosine_similarity = np.sum(feat1 * feat2, axis=1).item()
53
  return f"Cosine Similarity: {cosine_similarity:.4f}"
54
 
55
  # Create Gradio interface with custom layout
 
72
  )
73
 
74
  # Launch the interface
75
+ iface.launch()