Update app.py

app.py

@@ -34,7 +34,7 @@ cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rc
 cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # set threshold for this model
 # Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well
 cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
-#cfg.MODEL.DEVICE= 'cpu'
+cfg.MODEL.DEVICE= 'cpu'
 predictor = DefaultPredictor(cfg)
 '''
 os.system(wget http://images.cocodataset.org/val2017/000000439715.jpg -q -O input.jpg)
@@ -66,10 +66,13 @@ carTransforms = transforms.Compose([transforms.Resize(224),
 
 def classifyCar(im):
   im = Image.fromarray(im.astype('uint8'), 'RGB')
-  with torch.no_grad():
-    outputs = predictor(im)
-    v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
-    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
+  try:
+    with torch.no_grad():
+      outputs = predictor(im)
+      v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
+      out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
+  except Exception as err:
+    return im, {err: float(0.5)}
   im2 = carTransforms(im).unsqueeze(0)  # transform and add batch dimension
   with torch.no_grad():
     scores = torch.nn.functional.softmax(DesignModernityModel(im2)[0])