Refactor predict function to accept modality type and targets; update main.py to integrate modality selection and target loading

inference_utils/init_predict_mock.py

@@ -15,5 +15,5 @@ def init_model():
     return None
 
 
-def predict(image, prompts):
+def predict(image, modality_type: str, targets: list[str]):
     return image

main.py

@@ -10,7 +10,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import os
+from pathlib import Path
+from typing import Dict, List
+
+from inference_utils.target_dist import modality_targets_from_target_dist
 
 # If True, then mock init_model() and predict() functions will be used.
 DEV_MODE = True if os.getenv("DEV_MODE") else False
@@ -50,27 +55,43 @@ This Space is based on the [BiomedParse model](https://microsoft.github.io/Biome
 
 
 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/144DME_as_F.jpeg", "OCT", []],
+    ["examples/C3_EndoCV2021_00462.jpg", "Endoscopy", []],
+    ["examples/CT-abdomen.png", "CT-Abdomen", []],
+    ["examples/covid_1585.png", "X-Ray-Chest", []],
+    ["examples/ISIC_0015551.jpg", "Dermoscopy", []],
+    [
+        "examples/LIDC-IDRI-0140_143_280_CT_lung.png",
+        "CT-Chest",
+        [],
+    ],
     [
         "examples/Part_1_516_pathology_breast.png",
-        "neoplastic cells, inflammatory cells",
+        "Pathology",
+        [],
+    ],
+    ["examples/T0011.jpg", "Fundus", []],
+    [
+        "examples/TCGA_HT_7856_19950831_8_MRI-FLAIR_brain.png",
+        "MRI-FLAIR-Brain",
+        [],
     ],
-    ["examples/T0011.jpg", "optic disc"],
-    ["examples/T0011.jpg", "optic cup"],
-    ["examples/TCGA_HT_7856_19950831_8_MRI-FLAIR_brain.png", "glioma"],
 ]
 
 
+def load_modality_targets() -> Dict[str, List[str]]:
+    target_dist_json_path = Path("inference_utils/target_dist.json")
+    with open(target_dist_json_path, "r") as f:
+        target_dist = json.load(f)
+
+    modality_targets = modality_targets_from_target_dist(target_dist)
+    return modality_targets
+
+
+MODALITY_TARGETS = load_modality_targets()
+DEFAULT_MODALITY = "CT-Abdomen"
+
+
 def run():
     global model
     model = init_model()
@@ -82,23 +103,34 @@ def run():
         with gr.Row():
             with gr.Column():
                 input_image = gr.Image(type="pil", label="Input Image")
-                input_text = gr.Textbox(
-                    label="Prompts",
-                    placeholder="Enter prompts separated by commas (e.g., neoplastic cells, inflammatory cells)",
+                input_modality_type = gr.Dropdown(
+                    choices=list(MODALITY_TARGETS.keys()),
+                    label="Modality Type",
+                    value=DEFAULT_MODALITY,
+                )
+                input_targets = gr.CheckboxGroup(
+                    choices=MODALITY_TARGETS[DEFAULT_MODALITY],
+                    label="Targets",
                 )
             with gr.Column():
                 output_image = gr.Image(type="pil", label="Prediction")
 
-        predict_btn = gr.Button("Submit")
-        predict_btn.click(
+        input_modality_type.change(
+            fn=update_input_targets,
+            inputs=input_modality_type,
+            outputs=input_targets,
+        )
+
+        submit_btn = gr.Button("Submit")
+        submit_btn.click(
             fn=predict,
-            inputs=[input_image, input_text],
+            inputs=[input_image, input_modality_type, input_targets],
             outputs=output_image,
         )
 
         gr.Examples(
             examples=examples,
-            inputs=[input_image, input_text],
+            inputs=[input_image, input_modality_type, input_targets],
             outputs=output_image,
             fn=predict,
             cache_examples=False,
@@ -107,6 +139,14 @@ def run():
     return demo
 
 
+def update_input_targets(input_modality_type):
+    return gr.CheckboxGroup(
+        choices=MODALITY_TARGETS[input_modality_type],
+        value=[],
+        label="Targets",
+    )
+
+
 demo = run()
 
 if __name__ == "__main__":