Add development environment setup with mock model functions and Gradio interface

.gitignore

@@ -0,0 +1,2 @@
+.venv
+__pycache__

.python-version

@@ -0,0 +1 @@
+3.9

Makefile

@@ -0,0 +1,2 @@
+run:
+	DEV_MODE=1 uv run gradio main.py

README.md

@@ -37,3 +37,9 @@ All rights to the model and its underlying research are held by the original aut
 
 - **Code in this Space**: Licensed under the [Apache License 2.0](https://spdx.org/licenses/Apache-2.0.html), as per the original BiomedParse GitHub repository.
 - **Model**: The BiomedParse model is licensed under the [Creative Commons Attribution Non Commercial Share Alike 4.0](https://spdx.org/licenses/CC-BY-NC-SA-4.0.html). Ensure that your use complies with the terms of this license.
+
+## Development of the Gradio interface
+
+To develop the Gradio interface locally against a mock ML model, execute
+
+    make run

inference_utils/init_predict.py

@@ -0,0 +1,76 @@
+from PIL import Image
+from huggingface_hub import hf_hub_download
+import matplotlib.pyplot as plt
+import numpy as np
+from inference_utils.inference import interactive_infer_image
+from main import model
+
+
+import gradio as gr
+
+from modeling import build_model
+from modeling.BaseModel import BaseModel
+from utilities.arguments import load_opt_from_config_files
+from utilities.constants import BIOMED_CLASSES
+from utilities.distributed import init_distributed
+
+
+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
+
+
+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 predict(image: gr.Image, prompts: str):
+    if not prompts:
+        return None
+
+    # Convert string input to list
+    prompts = [p.strip() for p in prompts.split(",")]
+
+    # Convert to RGB if needed
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+
+    # Get predictions
+    pred_mask = interactive_infer_image(model, image, prompts)
+
+    # Generate visualization
+    colors = generate_colors(len(prompts))
+    pred_overlay = overlay_masks(
+        image, [1 * (pred_mask[i] > 0.5) for i in range(len(prompts))], colors
+    )
+
+    return pred_overlay
+
+
+def init_model():
+    # Download model
+    model_file = hf_hub_download(
+        repo_id="microsoft/BiomedParse",
+        filename="biomedparse_v1.pt",
+        token=os.getenv("HF_TOKEN"),
+    )
+
+    # Initialize model
+    conf_files = "configs/biomedparse_inference.yaml"
+    opt = load_opt_from_config_files([conf_files])
+    opt = init_distributed(opt)
+
+    model = BaseModel(opt, build_model(opt)).from_pretrained(model_file).eval().cuda()
+    with torch.no_grad():
+        model.model.sem_seg_head.predictor.lang_encoder.get_text_embeddings(
+            BIOMED_CLASSES + ["background"], is_eval=True
+        )
+
+    return model

inference_utils/init_predict_dev.py

@@ -0,0 +1,6 @@
+def init_model():
+    return None
+
+
+def predict(image, prompts):
+    return image

main.py

@@ -11,27 +11,23 @@
 # limitations under the License.
 
 import os
-import gradio as gr
-import torch
-from PIL import Image
-import numpy as np
-import matplotlib.pyplot as plt
-from huggingface_hub import hf_hub_download
-from modeling.BaseModel import BaseModel
-from modeling import build_model
-from utilities.distributed import init_distributed
-from utilities.arguments import load_opt_from_config_files
-from utilities.constants import BIOMED_CLASSES
-from inference_utils.inference import interactive_infer_image
 
 # If True, then mock init_model() and predict() functions will be used.
-DEV_MODE = False
+DEV_MODE = True if os.getenv("DEV_MODE") else False
+
+import gradio as gr
+
+if DEV_MODE:
+    from inference_utils.init_predict_dev import init_model, predict
+else:
+    from inference_utils.init_predict import init_model, predict
+
 
 gr.set_static_paths(["assets"])
 
 
 def run():
-    global model
+    global model, demo
     model = init_model()
 
     demo = gr.Interface(
@@ -75,68 +71,6 @@ def run():
         ],
     )
 
-    demo.launch(server_name="0.0.0.0", server_port=7860)
-
-
-def init_model_prod():
-    # Download model
-    model_file = hf_hub_download(
-        repo_id="microsoft/BiomedParse",
-        filename="biomedparse_v1.pt",
-        token=os.getenv("HF_TOKEN"),
-    )
-
-    # Initialize model
-    conf_files = "configs/biomedparse_inference.yaml"
-    opt = load_opt_from_config_files([conf_files])
-    opt = init_distributed(opt)
-
-    model = BaseModel(opt, build_model(opt)).from_pretrained(model_file).eval().cuda()
-    with torch.no_grad():
-        model.model.sem_seg_head.predictor.lang_encoder.get_text_embeddings(
-            BIOMED_CLASSES + ["background"], is_eval=True
-        )
-
-    return model
-
-
-def init_model_dev():
-    return None
-
-
-def predict_prod(image: gr.Image, prompts: str):
-    if not prompts:
-        return None
-
-    # Convert string input to list
-    prompts = [p.strip() for p in prompts.split(",")]
-
-    # Convert to RGB if needed
-    if image.mode != "RGB":
-        image = image.convert("RGB")
-
-    # Get predictions
-    pred_mask = interactive_infer_image(model, image, prompts)
-
-    # Generate visualization
-    colors = generate_colors(len(prompts))
-    pred_overlay = overlay_masks(
-        image, [1 * (pred_mask[i] > 0.5) for i in range(len(prompts))], colors
-    )
-
-    return pred_overlay
-
-
-def predict_dev(image: gr.Image, prompts: str):
-    return image
-
-
-if DEV_MODE:
-    init_model = init_model_dev
-    predict = predict_dev
-else:
-    init_model = init_model_prod
-    predict = predict_prod
 
 description = """Upload a biomedical image and enter prompts (separated by commas) to detect specific features.
 
@@ -161,21 +95,6 @@ This Space is based on the [BiomedParse model](https://microsoft.github.io/Biome
 """
 
 
-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
-
-
-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__":
     run()
+    demo.launch(server_name="0.0.0.0", server_port=7860)

pyproject.toml

@@ -0,0 +1,9 @@
+[project]
+name = "biomedparse"
+version = "0.1.0"
+description = "Python environment to develop the Gradio interface. Mocks the actual ML model."
+readme = "README.md"
+requires-python = ">=3.9"
+dependencies = [
+    "gradio==4.44.1",
+]