tut 12

app.py

@@ -1,8 +1,31 @@
 import gradio as gr
 import torch
+from torchvision import datasets, models, transforms
+from PIL import Image
 
-def greet(name):
-    return "Hello " + name + "!!"
+LABELS = ['VW Up!', 'Fiat 500']
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
+model = models.resnet18(pretrained=True)
+num_ftrs = model.fc.in_features
+
+model.fc = torch.nn.Linear(num_ftrs, 2)
+
+state_dict = torch.load('up500Model.pt', map_location='cpu')
+model.load_state_dict(state_dict, strict=False)
+model.eval()
+
+imgTransforms = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
+
+def predict(im):
+    inp = Image.fromarray(inp.astype('uint8'), 'RGB')
+    inp = imgTransforms(inp).unsqueeze(0)
+    with torch.no_grad():
+        prediction = torch.nn.functional.softmax(model(inp)[0], dim=0) 
+    return {LABELS[i]: float(prediction[i]) for i in range(2)}
+
+iface = gr.Interface(predict, inputs=gr.inputs.Image(), outputs="label")
 iface.launch()

requirements.txt

@@ -1 +1,2 @@
-torch
+torch
+torchvision