Update main.py

main.py

@@ -38,8 +38,31 @@ print(outputs["instances"].pred_classes)
 print(outputs["instances"].pred_boxes)
 
 # -- load Mask R-CNN model for segmentation
-MaskRCNN = torch.load("MaskRCNN.pt")
+DesignModernityModel = torch.load("DesignModernityModel.pt")
 
+#INPUT_FEATURES = DesignModernityModel.fc.in_features
+#linear = nn.linear(INPUT_FEATURES, 5)
 
+DesignModernityModel.eval() # set state of the model to inference
+
+LABELS = ['2003-2006', 'VW Up!']
+
+carTransforms = transforms.Compose([
+			transforms.RandomResizedCrop(224),
+			...
+])
+
+def classifyCar(im):
+  im = Image.fromarray(im.astype('uint8'), 'RGB')
+  im = carTransforms(im).unsqueeze(0)  # transform and add batch dimension
+  with torch.no_grad():
+    scores = torch.nn.functional.softmax(model(im)[0])
+  return {LABELS[i]: float(scores[i]) for i in range(2)}
+
+examples = [[example_img.jpg], [example_img2.jpg]]  # must be uploaded in repo
+
+# create interface for model
+interface = gr.Interface(classifyCar, inputs='Image', outputs='label', cache_examples=False, title='VW Up or Fiat 500', example=examples)
+interface.launch()