Jason Wu
commited on
Commit
•
7803ddd
1
Parent(s):
2a5021e
initial commit
Browse files- app.py +39 -0
- res/class_map_enrico.json +1 -0
app.py
ADDED
@@ -0,0 +1,39 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import gradio as gr
|
3 |
+
import json
|
4 |
+
from torchvision import transforms
|
5 |
+
import torch.nn.functional as F
|
6 |
+
|
7 |
+
TORCHSCRIPT_PATH = "res/screenclassification-resnet-noisystudent+web350k.torchscript"
|
8 |
+
LABELS_PATH = "res/class_map_enrico.json"
|
9 |
+
IMG_SIZE = 128
|
10 |
+
|
11 |
+
model = torch.jit.load(TORCHSCRIPT_PATH)
|
12 |
+
|
13 |
+
with open(LABELS_PATH, "r") as f:
|
14 |
+
label2Idx = json.load(f)["label2Idx"]
|
15 |
+
|
16 |
+
img_transforms = transforms.Compose([
|
17 |
+
transforms.Resize(IMG_SIZE),
|
18 |
+
transforms.ToTensor(),
|
19 |
+
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
|
20 |
+
])
|
21 |
+
|
22 |
+
def predict(img):
|
23 |
+
img_input = img_transforms(img).unsqueeze(0)
|
24 |
+
predictions = F.softmax(model(img_input), dim=-1)[0]
|
25 |
+
confidences = {}
|
26 |
+
for label in label2Idx:
|
27 |
+
confidences[label] = float(predictions[int(label2Idx[label])])
|
28 |
+
|
29 |
+
return confidences
|
30 |
+
|
31 |
+
example_imgs = [
|
32 |
+
"examples/example.jpg",
|
33 |
+
"examples/example_pair1.jpg",
|
34 |
+
"examples/example_pair2.jpg"
|
35 |
+
]
|
36 |
+
|
37 |
+
interface = gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs=gr.Label(num_top_classes=3), examples=example_imgs)
|
38 |
+
|
39 |
+
interface.launch()
|
res/class_map_enrico.json
ADDED
@@ -0,0 +1 @@
|
|
|
|
|
1 |
+
{"label2Idx": {"bare": 0, "calculator": 1, "camera": 2, "chat": 3, "editor": 4, "form": 5, "gallery": 6, "list": 7, "login": 8, "maps": 9, "mediaplayer": 10, "menu": 11, "modal": 12, "news": 13, "other": 14, "profile": 15, "search": 16, "settings": 17, "terms": 18, "tutorial": 19}, "idx2Label": {"0": "bare", "1": "calculator", "2": "camera", "3": "chat", "4": "editor", "5": "form", "6": "gallery", "7": "list", "8": "login", "9": "maps", "10": "mediaplayer", "11": "menu", "12": "modal", "13": "news", "14": "other", "15": "profile", "16": "search", "17": "settings", "18": "terms", "19": "tutorial"}}
|