Update app.py

app.py

@@ -78,17 +78,22 @@ def classifyCar(im):
     
     car_class_true = outputs["instances"].pred_classes == 2
     boxes = list(outputs["instances"].pred_boxes[car_class_true])
-    max_idx = torch.tensor([(x[2] - x[0])*(x[3] - x[1]) for x in boxes]).argmax().item()
     
-    im2 = Image.fromarray(np.uint8(im)).convert('RGB').crop(boxes[max_idx].to(torch.int64).numpy())
-    
-    carResize = transforms.Compose([transforms.Resize((224, 224))])
-    im2 = carResize(im2)
-    
-    with torch.no_grad():
-      scores = torch.nn.functional.softmax(DesignModernityModel(carTransforms(im2).unsqueeze(0))[0])
-
-    label = "success"
+    if len(boxes) != 0:
+      max_idx = torch.tensor([(x[2] - x[0])*(x[3] - x[1]) for x in boxes]).argmax().item()
+      
+      im2 = Image.fromarray(np.uint8(im)).convert('RGB').crop(boxes[max_idx].to(torch.int64).numpy())
+      
+      carResize = transforms.Compose([transforms.Resize((224, 224))])
+      im2 = carResize(im2)
+      
+      with torch.no_grad():
+        scores = torch.nn.functional.softmax(DesignModernityModel(carTransforms(im2).unsqueeze(0))[0])
+      label = {LABELS[i]: float(scores[i]) for i in range(n_labels)}
+      
+    else:
+        im2 = Image.fromarray(np.uint8(im)).convert('RGB')
+        label = "No car detected"
   except:
     label = "fail2"
     #im2 = carTransforms(im).unsqueeze(0)  # transform and add batch dimension
@@ -96,7 +101,7 @@ def classifyCar(im):
     #  scores = torch.nn.functional.softmax(DesignModernityModel(im2)[0])
       #{LABELS[i]: float(scores[i]) for i in range(n_labels)}
       #Image.fromarray(np.uint8(out.get_image())).convert('RGB')
-  return im2, {LABELS[i]: float(scores[i]) for i in range(n_labels)}
+  return im2, label
 
 #examples = [[example_img.jpg], [example_img2.jpg]]  # must be uploaded in repo