Update app.py

app.py

@@ -5,70 +5,91 @@ except:
     os.system('pip install git+https://github.com/facebookresearch/detectron2.git')
 
 import cv2
-
-from matplotlib.pyplot import axis
 import gradio as gr
 import requests
 import numpy as np
-from torch import nn
-import requests
-
-import torch
-
-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
 
-url1 = 'https://images.unsplash.com/photo-1503463168353-9d883c7f5255?q=80&w=1748&auto=format&fit=crop&ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D'
-r = requests.get(url1, allow_redirects=True)
-open("balloon.jpg", 'wb').write(r.content)
-url2 = 'https://images.unsplash.com/photo-1599158150601-1417ebbaafdd?q=80&w=1636&auto=format&fit=crop&ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D'
-r = requests.get(url2, allow_redirects=True)
-open("football.jpg", 'wb').write(r.content)
-
-model_name='COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml'
-
-# model = model_zoo.get(model_name, trained=True)
+# Predefined Detectron2 models
+models = [
+    {
+        "name": "Instance Segmentation",
+        "config_file": "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml",
+    },
+    {
+        "name": "Panoptic Segmentation",
+        "config_file": "COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml",
+    },
+]
+
+def setup_model(config_file):
+    cfg = get_cfg()
+    cfg.merge_from_file(model_zoo.get_config_file(config_file))
+    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(config_file)
+    if not torch.cuda.is_available():
+        cfg.MODEL.DEVICE = "cpu"
+    return cfg
+
+for model in models:
+    model["cfg"] = setup_model(model["config_file"])
+    model["metadata"] = MetadataCatalog.get(model["cfg"].DATASETS.TRAIN[0])
+
+def inference(image_url, image, min_score, model_name):
+    model = next((m for m in models if m["name"] == model_name), None)
+    if not model:
+        raise ValueError("Model not found")
+
+    if image_url:
+        r = requests.get(image_url)
+        if r:
+            im = np.frombuffer(r.content, dtype="uint8")
+            im = cv2.imdecode(im, cv2.IMREAD_COLOR)
+    else:
+        # Model expects BGR!
+        im = image[:,:,::-1]
+
+    model["cfg"].MODEL.ROI_HEADS.SCORE_THRESH_TEST = min_score
+    predictor = DefaultPredictor(model["cfg"])
+    outputs = predictor(im)
+
+    v = Visualizer(im, model["metadata"], scale=1.2)
+    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
 
-cfg = get_cfg()
-# 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(model_name))
-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 w ell
-cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(model_name)
+    return out.get_image()
 
-if not torch.cuda.is_available():
-    cfg.MODEL.DEVICE='cpu'
+title = "# Segmentation Model Demo"
+description = """
+This demo introduces an interactive playground for pretrained Detectron2 model.
+Currently, two models are supported that were trained on COCO dataset:
+* [Instance Segmentation](https://github.com/facebookresearch/detectron2/blob/main/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml): Identifies, outlines individual object instances.
+* [Panoptic Segmentation](https://github.com/facebookresearch/detectron2/blob/main/configs/COCO-PanopticSegmentation/panoptic_fpn_R_50_1x.yaml): Unifies instance and semantic segmentation.
+"""
+footer = "Made by eyepop.ai with ❤️."
+
+with gr.Blocks() as demo:
+    gr.Markdown(title)
+    gr.Markdown(description)
+    
+    with gr.Tab("From URL"):
+        url_input = gr.Textbox(label="Image URL", placeholder="https://images.unsplash.com/photo-1701226362119-cc86312846af?q=80&w=1587&auto=format&fit=crop&ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D")
+    
+    with gr.Tab("From Image"):
+        image_input = gr.Image(type="numpy", label="Input Image")
 
-predictor = DefaultPredictor(cfg)
+    min_score = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Minimum score")
 
+    model_name = gr.Radio(choices=[model["name"] for model in models], value=models[0]["name"], label="Select Detectron2 model")
 
-def inference(image):
-    img = np.array(image.resize((1024,1024)))
-    outputs = predictor(img)
+    output_image = gr.Image(type="pil", label="Output")
 
-    v = Visualizer(img, MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
-    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
+    inference_button = gr.Button("Submit")
     
-    return out.get_image()
-
+    inference_button.click(fn=inference, inputs=[url_input, image_input, min_score, model_name], outputs=output_image)
 
+    gr.Markdown(footer)
 
-title = "Detectron2-MaskRCNN X101"
-description = "demo for Detectron2. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below.\
-    </br><b>Model: COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml</b>"
-article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2012.07177'>Simple Copy-Paste is a Strong Data Augmentation Method for Instance Segmentation</a> | <a href='https://github.com/facebookresearch/detectron2/blob/main/MODEL_ZOO.md'>Detectron model ZOO</a></p>"
+demo.launch()
 
-gr.Interface(
-    inference, 
-    [gr.inputs.Image(type="pil", label="Input")], 
-    gr.outputs.Image(type="numpy", label="Output"),
-    title=title,
-    description=description,
-    article=article,
-    examples=[
-            ["balloon.jpg"],
-            ["football.jpg"]
-        ]).launch()