Refactor main.py to reorganize functions and enhance demo interface with additional examples

main.py

@@ -28,20 +28,52 @@ 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)
-    for mask, color in zip(masks, colors):
-        overlay[mask > 0] = (overlay[mask > 0] * 0.4 + np.array(color) * 0.6).astype(
-            np.uint8
-        )
-    return Image.fromarray(overlay)
+def run():
+    global model
+    model = init_model()
 
+    demo = gr.Interface(
+        fn=predict,
+        inputs=[
+            gr.Image(type="pil", label="Input Image"),
+            gr.Textbox(
+                label="Prompts",
+                placeholder="Enter prompts separated by commas (e.g., neoplastic cells, inflammatory cells)",
+            ),
+        ],
+        outputs=gr.Image(type="pil", label="Prediction"),
+        title="BiomedParse Demo",
+        description=description,
+        allow_flagging="never",
+        examples=[
+            ["examples/144DME_as_F.jpeg", "edema"],
+            ["examples/C3_EndoCV2021_00462.jpg", "polyp"],
+            ["examples/CT-abdomen.png", "liver, pancreas, spleen"],
+            ["examples/covid_1585.png", "left lung"],
+            ["examples/covid_1585.png", "right lung"],
+            ["examples/covid_1585.png", "COVID-19 infection"],
+            ["examples/ISIC_0015551.jpg", "lesion"],
+            ["examples/LIDC-IDRI-0140_143_280_CT_lung.png", "lung nodule"],
+            ["examples/LIDC-IDRI-0140_143_280_CT_lung.png", "COVID-19 infection"],
+            [
+                "examples/Part_1_516_pathology_breast.png",
+                "connective tissue cells",
+            ],
+            [
+                "examples/Part_1_516_pathology_breast.png",
+                "neoplastic cells",
+            ],
+            [
+                "examples/Part_1_516_pathology_breast.png",
+                "neoplastic cells, inflammatory cells",
+            ],
+            ["examples/T0011.jpg", "optic disc"],
+            ["examples/T0011.jpg", "optic cup"],
+            ["examples/TCGA_HT_7856_19950831_8_MRI-FLAIR_brain.png", "glioma"],
+        ],
+    )
 
-def generate_colors(n):
-    cmap = plt.get_cmap("tab10")
-    colors = [tuple(int(255 * val) for val in cmap(i)[:3]) for i in range(n)]
-    return colors
+    demo.launch(server_name="0.0.0.0", server_port=7860)
 
 
 def init_model():
@@ -112,52 +144,20 @@ This Space is based on the [BiomedParse model](https://microsoft.github.io/Biome
 """
 
 
-def run():
-    global model
-    model = init_model()
+def generate_colors(n):
+    cmap = plt.get_cmap("tab10")
+    colors = [tuple(int(255 * val) for val in cmap(i)[:3]) for i in range(n)]
+    return colors
 
-    demo = gr.Interface(
-        fn=predict,
-        inputs=[
-            gr.Image(type="pil", label="Input Image"),
-            gr.Textbox(
-                label="Prompts",
-                placeholder="Enter prompts separated by commas (e.g., neoplastic cells, inflammatory cells)",
-            ),
-        ],
-        outputs=gr.Image(type="pil", label="Prediction"),
-        title="BiomedParse Demo",
-        description=description,
-        allow_flagging="never",
-        examples=[
-            ["examples/144DME_as_F.jpeg", "edema"],
-            ["examples/C3_EndoCV2021_00462.jpg", "polyp"],
-            ["examples/CT-abdomen.png", "liver, pancreas, spleen"],
-            ["examples/covid_1585.png", "left lung"],
-            ["examples/covid_1585.png", "right lung"],
-            ["examples/covid_1585.png", "COVID-19 infection"],
-            ["examples/ISIC_0015551.jpg", "lesion"],
-            ["examples/LIDC-IDRI-0140_143_280_CT_lung.png", "lung nodule"],
-            ["examples/LIDC-IDRI-0140_143_280_CT_lung.png", "COVID-19 infection"],
-            [
-                "examples/Part_1_516_pathology_breast.png",
-                "connective tissue cells",
-            ],
-            [
-                "examples/Part_1_516_pathology_breast.png",
-                "neoplastic cells",
-            ],
-            [
-                "examples/Part_1_516_pathology_breast.png",
-                "neoplastic cells, inflammatory cells",
-            ],
-            ["examples/T0011.jpg", "optic disc"],
-            ["examples/T0011.jpg", "optic cup"],
-            ["examples/TCGA_HT_7856_19950831_8_MRI-FLAIR_brain.png", "glioma"],
-        ],
-    )
 
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+def overlay_masks(image, masks, colors):
+    overlay = image.copy()
+    overlay = np.array(overlay, dtype=np.uint8)
+    for mask, color in zip(masks, colors):
+        overlay[mask > 0] = (overlay[mask > 0] * 0.4 + np.array(color) * 0.6).astype(
+            np.uint8
+        )
+    return Image.fromarray(overlay)
 
 
 if __name__ == "__main__":