Update app.py

app.py

@@ -9,17 +9,26 @@ from PIL import Image
 import io 
 from huggingface_hub import hf_hub_download
 
-# Download model at startup
-HF_REPO_ID = "myang4218/cellposemodel" 
-HF_MODEL_FILENAME = "cellposemodel.npy"
-model_path = hf_hub_download(repo_id=HF_REPO_ID, filename=HF_MODEL_FILENAME)
-print(f"Model downloaded to: {model_path}")
+HF_REPO_ID = "myang4218/cellposemodel"
+
+MODEL_OPTIONS = {
+    "Hemocytometer Model": "hemocytometermodel.npy",
+    "General Model": "generalmodel.npy"
+}
+
+loaded_models = {}
 
 @spaces.GPU
-def segment_and_count(image):
-    # Initialize model inside the GPU-decorated function
-    model = models.CellposeModel(gpu=True, pretrained_model=model_path)
+def segment_and_count(image, model_choice):
     
+    model_filename = MODEL_OPTIONS[model_choice]
+    model_path = hf_hub_download(repo_id=HF_REPO_ID, filename=model_filename)
+    if model_filename in loaded_models:
+        model = loaded_models[model_filename]
+    else:
+        model = models.CellposeModel(gpu=True, pretrained_model=model_path)
+        loaded_models[model_filename] = model
+        
     # Convert PIL Image to numpy array
     image_np = np.array(image)
     
@@ -64,13 +73,16 @@ def segment_and_count(image):
 # Gradio interface
 demo = gr.Interface(
     fn=segment_and_count,
-    inputs=gr.Image(type="pil"),
+    inputs=[
+        gr.Image(type="pil", label="Microscopy Image"),
+        gr.Dropdown(choices=list(MODEL_OPTIONS.keys()), label="Select Model", value="Hemocytometer Model")
+    ],
     outputs=[
-        gr.Number(label="Number of Cells"), 
+        gr.Number(label="Number of Cells"),
         gr.Image(type="pil", label="Segmented Overlay")
     ],
     title="Cell Counter with Cellpose",
-    description="Upload a microscopy image and count the number of cells using Cellpose segmentation.",
+    description="Upload a microscopy image and select a model to count the number of cells using Cellpose segmentation."
 )
 
 if __name__ == "__main__":