Jason Wu commited on
Commit
7803ddd
1 Parent(s): 2a5021e

initial commit

Browse files
Files changed (2) hide show
  1. app.py +39 -0
  2. 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"}}