Add images for ontology hierarchy and modality types; update demo description

main.py

@@ -25,6 +25,9 @@ from utilities.constants import BIOMED_CLASSES
 from inference_utils.inference import interactive_infer_image
 
 
+gr.set_static_paths(["assets"])
+
+
 def overlay_masks(image, masks, colors):
     overlay = image.copy()
     overlay = np.array(overlay, dtype=np.uint8)
@@ -76,7 +79,6 @@ def predict(image, prompts):
 
     # Get predictions
     pred_mask = interactive_infer_image(model, image, prompts)
-    print(f"Predicted masks: {pred_mask}")
 
     # Generate visualization
     colors = generate_colors(len(prompts))
@@ -87,6 +89,18 @@ def predict(image, prompts):
     return pred_overlay
 
 
+description = """Upload a biomedical image and enter prompts (separated by commas) to detect specific features.
+
+The model understands these prompts:
+![gpt4_ontology_hierarchy.png](file/assets/gpt4_ontology_hierarchy.png)
+
+The model understands these types of biomedical images:
+![modality_bar_plot.png](file/assets/modality_bar_plot.png)
+
+Above figures are from the [BiomedParse paper](https://arxiv.org/abs/2405.12971).
+"""
+
+
 def run():
     global model
     model = init_model()
@@ -102,8 +116,8 @@ def run():
         ],
         outputs=gr.Image(type="pil", label="Prediction"),
         title="BiomedParse Demo",
-        description="Upload a biomedical image and enter prompts (separated by commas) to detect specific features.",
-        allow_flagging=False,
+        description=description,
+        allow_flagging="never",
         examples=[
             ["examples/144DME_as_F.jpeg", "edema"],
             ["examples/C3_EndoCV2021_00462.jpg", "polyp"],