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)
@@ -46,7 +46,7 @@ outputs = predictor(im)
 print(outputs["instances"].pred_classes)
 print(outputs["instances"].pred_boxes)
 '''
-# -- load Mask R-CNN model for segmentation
+# -- load design modernity model for classification
 DesignModernityModel = torch.load("DesignModernityModel.pt")
 
 #INPUT_FEATURES = DesignModernityModel.fc.in_features
@@ -70,9 +70,9 @@ def classifyCar(im):
     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"))
-  im = carTransforms(im).unsqueeze(0)  # transform and add batch dimension
+  im2 = carTransforms(im).unsqueeze(0)  # transform and add batch dimension
   with torch.no_grad():
-    scores = torch.nn.functional.softmax(DesignModernityModel(im)[0])
+    scores = torch.nn.functional.softmax(DesignModernityModel(im2)[0])
   return Image.fromarray(np.uint8(out.get_image())).convert('RGB'), {LABELS[i]: float(scores[i]) for i in range(n_labels)}
 
 #examples = [[example_img.jpg], [example_img2.jpg]]  # must be uploaded in repo