app: rewrite in Gradio Blocks with multi-model support

app.py

@@ -22,21 +22,40 @@ from detectron2.utils.visualizer import Visualizer
 from detectron2.data import MetadataCatalog
 
 
-model_path = "https://huggingface.co/dbmdz/detectron2-model/resolve/main/model_final.pth"
-
-cfg = get_cfg()
-cfg.merge_from_file("./configs/detectron2/faster_rcnn_R_50_FPN_3x.yaml")
-cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2
-cfg.MODEL.WEIGHTS = model_path
-
-my_metadata = MetadataCatalog.get("dbmdz_coco_all")
-my_metadata.thing_classes = ["Illumination", "Illustration"]
-
-if not torch.cuda.is_available():
-    cfg.MODEL.DEVICE = "cpu"
-
-
-def inference(image_url, image, min_score):
+models = [
+{
+    "name": "Version 1 (2-class)",
+    "model_path": "https://huggingface.co/dbmdz/detectron2-model/resolve/main/model_final.pth",
+    "classes": ["Illumination", "Illustration"],
+    "cfg": None,
+    "metadata": None
+},
+{
+    "name": "Version 2 (4-class)",
+    "model_path": "https://huggingface.co/dbmdz/detectron2-v2-model/resolve/main/model_final.pth",
+    "classes": ["ILLUSTRATION", "OTHER", "STAMP", "INITIAL"],
+    "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("./configs/detectron2/faster_rcnn_R_50_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"
+
+
+def inference(image_url, image, min_score, model_name):
     if image_url:
         r = requests.get(image_url)
         if r:
@@ -46,29 +65,63 @@ def inference(image_url, image, min_score):
         # Model expect BGR!
         im = image[:,:,::-1]
 
-    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = min_score
-    predictor = DefaultPredictor(cfg)
+    model_id = model_name_to_id[model_name]
+
+    models[model_id]["cfg"].MODEL.ROI_HEADS.SCORE_THRESH_TEST = min_score
+    predictor = DefaultPredictor(models[model_id]["cfg"])
 
     outputs = predictor(im)
 
-    v = Visualizer(im, my_metadata, scale=1.2)
+    v = Visualizer(im, models[model_id]["metadata"], scale=1.2)
     out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
 
     return out.get_image()
 
 
-title = "DBMDZ Detectron2 Model Demo"
-description = "This demo introduces an interactive playground for our trained Detectron2 model. <br>The model was trained on manually annotated segments from digitized books to detect Illustration or Illumination segments on a given page."
-article = '<p>Detectron model is available from our repository <a href="">here</a> on the Hugging Face Model Hub.</p>'
-
-gr.Interface(
-    inference,
-    [gr.inputs.Textbox(label="Image URL", placeholder="https://api.digitale-sammlungen.de/iiif/image/v2/bsb10483966_00008/full/500,/0/default.jpg"),
-     gr.inputs.Image(type="numpy", label="Input Image"),
-     gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Minimum score"),
-    ], 
-    gr.outputs.Image(type="pil", label="Output"),
-    title=title,
-    description=description,
-    article=article,
-    examples=[]).launch()
+title = "# DBMDZ Detectron2 Model Demo"
+description = """
+This demo introduces an interactive playground for our trained Detectron2 model.
+
+Currently, two models are supported that were trained on manually annotated segments from digitized books:
+
+* [Version 1 (2-class)](https://huggingface.co/dbmdz/detectron2-model): This model can detect *Illustration* or *Illumination* segments on a given page.
+* [Version 2 (4-class)](https://huggingface.co/dbmdz/detectron2-v2-model): This model is more powerful and can detect *Illustration*, *Stamp*, *Initial* or *Other* segments on a given page.
+"""
+footer = "Made in Munich with ❤️ and 🥨."
+
+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://api.digitale-sammlungen.de/iiif/image/v2/bsb10483966_00008/full/500,/0/default.jpg")
+    
+    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")
+
+    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)
+
+demo.launch()
+
+#gr.Interface(
+#    inference,
+#    [gr.inputs.Textbox(label="Image URL", placeholder="https://api.digitale-sammlungen.de/iiif/image/v2/bsb10483966_00008/full/500,/0/default.jpg"),
+#     gr.inputs.Image(type="numpy", label="Input Image"),
+#     gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Minimum score"),
+#     gr.Radio(choices=[model["name"] for model in models], value=models[0]["name"], label="Select Detectron2 model"),
+#    ], 
+#    gr.outputs.Image(type="pil", label="Output"),
+#    title=title,
+#    description=description,
+#    article=article,
+#    examples=[]).launch()