Update app.py

app.py

@@ -10,40 +10,50 @@ os.system("python setup.py build install")
 os.chdir("/home/user/app/Mask2Former/")
 import gradio as gr
 # check pytorch installation: 
-import torch, torchvision
-print(torch.__version__, torch.cuda.is_available())
-assert torch.__version__.startswith("1.9")   # please manually install torch 1.9 if Colab changes its default version
-# Some basic setup:
-# Setup detectron2 logger
 import detectron2
 from detectron2.utils.logger import setup_logger
+setup_logger()
+setup_logger(name="mask2former")
 
 # import some common libraries
 import numpy as np
-import os, json, cv2, random
+import cv2
+import torch
+from google.colab.patches import cv2_imshow
 
 # import some common detectron2 utilities
 from detectron2 import model_zoo
 from detectron2.engine import DefaultPredictor
 from detectron2.config import get_cfg
-from detectron2.utils.visualizer import Visualizer
-from detectron2.data import MetadataCatalog, DatasetCatalog
-from PIL import Image
+from detectron2.utils.visualizer import Visualizer, ColorMode
+from detectron2.data import MetadataCatalog
+from detectron2.projects.deeplab import add_deeplab_config
+coco_metadata = MetadataCatalog.get("coco_2017_val_panoptic")
+
+# import Mask2Former project
+from mask2former import add_maskformer2_config
 
 cfg = get_cfg()
 cfg.MODEL.DEVICE='cpu'
-# add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library
-cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
-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")
+add_deeplab_config(cfg)
+add_maskformer2_config(cfg)
+cfg.merge_from_file("configs/coco/panoptic-segmentation/swin/maskformer2_swin_large_IN21k_384_bs16_100ep.yaml")
+cfg.MODEL.WEIGHTS = 'https://dl.fbaipublicfiles.com/maskformer/mask2former/coco/panoptic/maskformer2_swin_large_IN21k_384_bs16_100ep/model_final_f07440.pkl'
+cfg.MODEL.MASK_FORMER.TEST.SEMANTIC_ON = True
+cfg.MODEL.MASK_FORMER.TEST.INSTANCE_ON = True
+cfg.MODEL.MASK_FORMER.TEST.PANOPTIC_ON = True
 predictor = DefaultPredictor(cfg)
+outputs = predictor(im)
+
 def inference(img):
-    im = cv2.imread(img.name)
-    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"))
-    return Image.fromarray(np.uint8(out.get_image())).convert('RGB')
+    im = cv2.imread(img)
+    v = Visualizer(im[:, :, ::-1], coco_metadata, scale=1.2, instance_mode=ColorMode.IMAGE_BW)
+    panoptic_result = v.draw_panoptic_seg(outputs["panoptic_seg"][0].to("cpu"), outputs["panoptic_seg"][1]).get_image()
+    v = Visualizer(im[:, :, ::-1], coco_metadata, scale=1.2, instance_mode=ColorMode.IMAGE_BW)
+    instance_result = v.draw_instance_predictions(outputs["instances"].to("cpu")).get_image()
+    v = Visualizer(im[:, :, ::-1], coco_metadata, scale=1.2, instance_mode=ColorMode.IMAGE_BW)
+    semantic_result = v.draw_sem_seg(outputs["sem_seg"].argmax(0).to("cpu")).get_image()    
+    return Image.fromarray(np.uint8(panoptic_result)).convert('RGB'),Image.fromarray(np.uint8(instance_result)).convert('RGB'),Image.fromarray(np.uint8(semantic_result)).convert('RGB')
     
 
 title = "Detectron 2"
@@ -52,7 +62,7 @@ article = "<p style='text-align: center'><a href='https://ai.facebook.com/blog/-
 
 examples = [['airplane.png']]
 
-gr.Interface(inference, inputs=gr.inputs.Image(type="file"), outputs=gr.outputs.Image(type="pil"),enable_queue=True, title=title,
+gr.Interface(inference, inputs=gr.inputs.Image(type="filepath"), outputs=[gr.outputs.Image(label="Panoptic segmentation",type="pil"),gr.outputs.Image(label="instance segmentation",type="pil"),gr.outputs.Image(label="semantic segmentation",type="pil")],enable_queue=True, title=title,
     description=description,
     article=article,
     examples=examples).launch()