Update app.py

app.py

@@ -8,8 +8,8 @@ import torch.nn.functional as F
 import torchvision.transforms as T
 
 # Get pre-trained model
-#weights = torchvision.models.ResNet18_Weights.IMAGENET1K_V1
-model = torchvision.models.resnet18()# weights=weights)
+weights = torchvision.models.ResNet18_Weights.IMAGENET1K_V1
+model = torchvision.models.resnet18(weights=weights)
 
 # Set model to evaluation mode
 model.eval()
@@ -20,10 +20,13 @@ labels = r.text.split("\n")
 
 # Define prediction function
 def predict(img):
+    '''
+    img: PIL image to be predicted
+    confidences: python dictionary containing confidences
+    '''
     # Transform image to pytorch tensor of shape [1, 3, 224, 224]
     img = T.PILToTensor()(img).unsqueeze(0)
     img = T.Resize(size=(224, 224))(img)
-
     
     # Use model without gradients to reduce computation
     with torch.no_grad():
@@ -38,4 +41,4 @@ gr.Interface(fn=predict,
              inputs=gr.Image(type="pil"),
              outputs=gr.Label(num_top_classes=10),
              theme="default",
-            )
+            ).launch()