pjdevelop commited on
Commit
e444fbb
1 Parent(s): 6403686

commit files to HF hub

Browse files
Files changed (2) hide show
  1. app.py +26 -0
  2. requirements.txt +8 -0
app.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import gradio as gr
3
+ import torch
4
+ import numpy as np
5
+ from PIL import Image
6
+ from torchvision import transforms as T
7
+ import joblib
8
+
9
+ # Load models
10
+ dinov2_vits14 = torch.load('dinov2_vits14.pth', map_location=torch.device('cpu'))
11
+ clf = joblib.load('svm_model.joblib')
12
+
13
+ # Transform for input image
14
+ transform_image = T.Compose([T.ToTensor(), T.Resize(244), T.CenterCrop(224), T.Normalize([0.5], [0.5])])
15
+
16
+ def predict(image):
17
+ image = Image.fromarray(image)
18
+ transformed_img = transform_image(image)[:3].unsqueeze(0)
19
+ with torch.no_grad():
20
+ embedding = dinov2_vits14(transformed_img)
21
+ prediction = clf.predict(np.array(embedding[0].cpu()).reshape(1, -1))
22
+ return prediction[0]
23
+
24
+ iface = gr.Interface(fn=predict, inputs="image", outputs="text")
25
+ iface.launch()
26
+
requirements.txt ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+
2
+ gradio
3
+ torch
4
+ torchvision
5
+ joblib
6
+ Pillow
7
+ numpy
8
+