Update app.py

app.py

@@ -21,86 +21,54 @@ from detectron2.config import get_cfg
 from detectron2.utils.visualizer import Visualizer
 from detectron2.data import MetadataCatalog
 
-# Can add more models if needed for model selection purposes. 
-models = [
-{
-    "name": "Version 1 (2-class)",
-    "model_path": "https://huggingface.co/spaces/eyepop-ai/segmentation/blob/main/model_final.pth",
-    "classes": ["cell", "object"],
-    "cfg": None,
-    "metadata": None
-},
-]
-
-model_name_to_id = {model["name"] : id_ for id_, model in enumerate(models)}
-
-for model in models:
-
-    model["cfg"] = get_cfg()
-    # model["cfg"].merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml"))
-    model["cfg"].MODEL.ROI_HEADS.NUM_CLASSES = len(model["classes"])
-    model["cfg"].MODEL.WEIGHTS = model["model_path"]
-    
-    model["metadata"] = MetadataCatalog.get(model["name"])
-    model["metadata"].thing_classes = model["classes"]
-
-    if not torch.cuda.is_available():
-      model["cfg"].MODEL.DEVICE = "cpu"
-
+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)
 
-def inference(image_url, image, min_score, model_name):
-    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 expect BGR!
-        im = image[:,:,::-1]
+model_name='COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml'
 
-    model_id = model_name_to_id[model_name]
+# model = model_zoo.get(model_name, trained=True)
 
-    models[model_id]["cfg"].MODEL.ROI_HEADS.SCORE_THRESH_TEST = min_score
-    print("model cfg", models[model_id]["cfg"])
-    predictor = DefaultPredictor(models[model_id]["cfg"])
-    print("predictor", predictor)
+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)
 
-    outputs = predictor(im)
+if not torch.cuda.is_available():
+    cfg.MODEL.DEVICE='cpu'
 
-    v = Visualizer(im, models[model_id]["metadata"], scale=0.8)
-    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
-
-    return out.get_image()
+predictor = DefaultPredictor(cfg)
 
 
-title = "# Instance Segmentation Demo"
-description = """
-This demo introduces an interactive playground for our trained Detectron2 model.
-Currently, a model is supported that was trained on manually annotated segments from petri dish images:
-* [Version 1 (1-class)](https://huggingface.co/dbmdz/detectron2-model): This model can detect *Object* segments on a given petri dish.
-"""
-footer = "Made by eyepop.ai with ❤️."
+def inference(image):
+    img = np.array(image.resize((1024,1024)))
+    outputs = predictor(img)
 
-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://huggingface.co/spaces/eyepop-ai/segmentation/blob/main/image_10_jpg.rf.46f8c223849190c943d24da8555d62fe.jpg")
+    v = Visualizer(img, MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
+    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
     
-    with gr.Tab("From Image"):
-        image_input = gr.Image(type="numpy", label="Input Image")
-
-    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")
+    return out.get_image()
 
-    output_image = gr.Image(type="pil", label="Output")
 
-    inference_button = gr.Button("Submit")
-    
-    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()