Update app.py

app.py

@@ -40,23 +40,10 @@ def convert_img_to_tensor(img):
 
 
 def predict(img) -> Dict:
-    """Transforms and performs a prediction on img and returns prediction and time taken.
+    """Transforms and performs a prediction on img and returns prediction.
     """
-    # Start the timer
-    # print(type(img))
-    # img = img['composite']
-    # if type(img) == numpy.ndarray:
-    #   img = PIL.Image.fromarray(img)
     img = convert_img_to_tensor(img)
-    # padding =  transforms.CenterCrop([224, 224])
-    # img = padding(img)
-    # img = torch.from_numpy(img)
-    # print(img.shape)
-    # Transform the target image and add a batch dimension
-    # # img = torch.tensor(img, dtype=torch.float32).unsqueeze(0).unsqueeze(0) / 255
     img = mobileNetV3_transform(img).unsqueeze(0)
-    # img = image_transform(img).unsqueeze(0)
-
 
     # Put model into evaluation mode and turn on inference mode
     mobileNetV3.eval()
@@ -64,7 +51,6 @@ def predict(img) -> Dict:
         # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
         y_pred = mobileNetV3(img)
         pred_probs = torch.softmax(y_pred, dim=1)
-        # y_pred_class = torch.argmax(torch.softmax(y_pred, dim=1), dim=1)
 
     # Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
     pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
@@ -73,6 +59,7 @@ def predict(img) -> Dict:
 
     # Return the prediction dictionary and prediction time
     return pred_labels_and_probs
+
 link_to_code = "https://github.com/henrythdu/MobileNetV3-QuickDraw-Animals"
 title = "MobileNetV3 - Quick Draw - Animals 🔢"
 description = "An MobileNetV3 feature extractor computer vision model to classify doodling of animals from Quick Draw dataset. The goal is to have a small model without significantly compromising on the predictive performance"