Update app.py

app.py

@@ -14,7 +14,7 @@ os.system('pip install git+https://github.com/facebookresearch/detectron2.git')
 os.system('pip install pyyaml==5.1')
 
 import detectron2
-'''
+
 from detectron2.utils.logger import setup_logger # ????
 
 from detectron2 import model_zoo
@@ -32,7 +32,7 @@ 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")
 predictor = DefaultPredictor(cfg)
-
+'''
 os.system(wget http://images.cocodataset.org/val2017/000000439715.jpg -q -O input.jpg)
 im = cv2.imread("./input.jpg")
 cv2_imshow(im)
@@ -64,13 +64,16 @@ def classifyCar(im):
   im = Image.fromarray(im.astype('uint8'), 'RGB')
   im = carTransforms(im).unsqueeze(0)  # transform and add batch dimension
   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"))
     scores = torch.nn.functional.softmax(DesignModernityModel(im)[0])
-  return {LABELS[i]: float(scores[i]) for i in range(n_labels)}
+  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
 
 # create interface for model
-interface = gr.Interface(classifyCar, inputs='image', outputs='label', cache_examples=False, title='VW Up or Fiat 500')
+interface = gr.Interface(classifyCar, inputs='image', outputs=['image','label'], cache_examples=False, title='VW Up or Fiat 500')
 interface.launch()