Rename init_predict_mock.py to model_mock.py and update main.py to use the Model class for predictions

inference_utils/{init_predict_mock.py → model_mock.py}

@@ -23,7 +23,7 @@ class Model:
         pass
 
     def predict(
-        image: Image, modality_type: str, targets: list[str]
+        self, image: Image, modality_type: str, targets: list[str]
     ) -> Tuple[Image, str]:
         # Randomly split targets into found and not found
         targets_found = random.sample(targets, k=len(targets) // 2)

main.py

@@ -23,9 +23,9 @@ DEV_MODE = True if os.getenv("DEV_MODE") else False
 import gradio as gr
 
 if DEV_MODE:
-    from inference_utils.init_predict_mock import init_model, predict
+    from inference_utils.model_mock import Model
 else:
-    from inference_utils.init_predict import init_model, predict
+    from inference_utils.model import Model
 
 
 gr.set_static_paths(["assets"])
@@ -93,8 +93,8 @@ DEFAULT_MODALITY = "CT-Abdomen"
 
 
 def run():
-    global model
-    model = init_model()
+    model = Model()
+    model.init()
 
     with gr.Blocks() as demo:
         gr.Markdown("# BiomedParse Demo")
@@ -126,7 +126,7 @@ def run():
 
         submit_btn = gr.Button("Submit")
         submit_btn.click(
-            fn=predict,
+            fn=model.predict,
             inputs=[input_image, input_modality_type, input_targets],
             outputs=[output_image, output_targets_not_found],
         )
@@ -134,8 +134,8 @@ def run():
         gr.Examples(
             examples=examples,
             inputs=[input_image, input_modality_type, input_targets],
-            outputs=output_image,
-            fn=predict,
+            outputs=[output_image, output_targets_not_found],
+            fn=model.predict,
             cache_examples=False,
         )