Update app.py

Added Model Selection

app.py

@@ -10,11 +10,42 @@ import numpy as np
 import requests
 import cv2
 
-# Load models once
-print("Loading segmentation model...")
-segmentation_model = pipeline("image-segmentation", model="nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
-print("Loading depth estimation model...")
-depth_estimator = pipeline("depth-estimation", model="Intel/zoedepth-nyu-kitti")
+# Dictionary of available segmentation models
+SEGMENTATION_MODELS = {
+    "NVIDIA SegFormer (Cityscapes)": "nvidia/segformer-b1-finetuned-cityscapes-1024-1024",
+    "NVIDIA SegFormer (ADE20K)": "nvidia/segformer-b0-finetuned-ade-512-512",
+    "Facebook MaskFormer (COCO)": "facebook/maskformer-swin-base-ade",
+    "OneFormer (COCO)": "shi-labs/oneformer_coco_swin_large",
+    "NVIDIA SegFormer (B5)": "nvidia/segformer-b5-finetuned-cityscapes-1024-1024"
+}
+
+# Dictionary of available depth estimation models
+DEPTH_MODELS = {
+    "Intel ZoeDepth (NYU-KITTI)": "Intel/zoedepth-nyu-kitti",
+    "DPT (Large)": "Intel/dpt-large",
+    "DPT (Hybrid)": "Intel/dpt-hybrid-midas",
+    "GLPDepth": "vinvino02/glpn-nyu"
+}
+
+# Initialize model placeholders
+segmentation_model = None
+depth_estimator = None
+
+def load_segmentation_model(model_name):
+    """Load the selected segmentation model"""
+    global segmentation_model
+    model_path = SEGMENTATION_MODELS[model_name]
+    print(f"Loading segmentation model: {model_path}...")
+    segmentation_model = pipeline("image-segmentation", model=model_path)
+    return f"Loaded segmentation model: {model_name}"
+
+def load_depth_model(model_name):
+    """Load the selected depth estimation model"""
+    global depth_estimator
+    model_path = DEPTH_MODELS[model_name]
+    print(f"Loading depth estimation model: {model_path}...")
+    depth_estimator = pipeline("depth-estimation", model=model_path)
+    return f"Loaded depth model: {model_name}"
 
 def lens_blur(image, radius):
     """
@@ -72,6 +103,10 @@ def process_image(input_image, method, blur_intensity, blur_type):
        - output_image: final composited image.
        - mask_image: the mask used (binary for segmentation, normalized depth for depth-based).
     """
+    # Check if models are loaded
+    if segmentation_model is None or depth_estimator is None:
+        return input_image, input_image.convert("L")
+    
     # Ensure image is in RGB mode
     input_image = input_image.convert("RGB")
     
@@ -135,22 +170,71 @@ def process_image(input_image, method, blur_intensity, blur_type):
 with gr.Blocks() as demo:
     gr.Markdown("## Image Processing App: Segmentation & Depth-based Blur")
     
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(label="Input Image", type="pil")
-            method = gr.Radio(label="Processing Method", 
-                              choices=["Segmented Background Blur", "Depth-based Variable Blur"],
-                              value="Segmented Background Blur")
-            blur_intensity = gr.Slider(label="Blur Intensity (Maximum Blur Radius)", minimum=1, maximum=30, step=1, value=15)
-            blur_type = gr.Dropdown(label="Blur Type", choices=["Gaussian Blur", "Lens Blur"], value="Gaussian Blur")
-            run_button = gr.Button("Process Image")
-        with gr.Column():
-            output_image = gr.Image(label="Output Image")
-            mask_output = gr.Image(label="Mask")
-    
-    run_button.click(fn=process_image, 
-                     inputs=[input_image, method, blur_intensity, blur_type], 
-                     outputs=[output_image, mask_output])
+    with gr.Tab("Model Selection"):
+        with gr.Row():
+            with gr.Column():
+                seg_model_dropdown = gr.Dropdown(
+                    label="Segmentation Model",
+                    choices=list(SEGMENTATION_MODELS.keys()),
+                    value=list(SEGMENTATION_MODELS.keys())[0]
+                )
+                seg_model_load_btn = gr.Button("Load Segmentation Model")
+                seg_model_status = gr.Textbox(label="Status", value="No model loaded")
+            
+            with gr.Column():
+                depth_model_dropdown = gr.Dropdown(
+                    label="Depth Estimation Model",
+                    choices=list(DEPTH_MODELS.keys()),
+                    value=list(DEPTH_MODELS.keys())[0]
+                )
+                depth_model_load_btn = gr.Button("Load Depth Model")
+                depth_model_status = gr.Textbox(label="Status", value="No model loaded")
+    
+    with gr.Tab("Image Processing"):
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(label="Input Image", type="pil")
+                method = gr.Radio(label="Processing Method", 
+                                choices=["Segmented Background Blur", "Depth-based Variable Blur"],
+                                value="Segmented Background Blur")
+                blur_intensity = gr.Slider(label="Blur Intensity (Maximum Blur Radius)", 
+                                        minimum=1, maximum=30, step=1, value=15)
+                blur_type = gr.Dropdown(label="Blur Type", 
+                                        choices=["Gaussian Blur", "Lens Blur"], 
+                                        value="Gaussian Blur")
+                run_button = gr.Button("Process Image")
+            with gr.Column():
+                output_image = gr.Image(label="Output Image")
+                mask_output = gr.Image(label="Mask")
+    
+    # Set up event handlers
+    seg_model_load_btn.click(
+        fn=load_segmentation_model,
+        inputs=[seg_model_dropdown],
+        outputs=[seg_model_status]
+    )
+    
+    depth_model_load_btn.click(
+        fn=load_depth_model,
+        inputs=[depth_model_dropdown],
+        outputs=[depth_model_status]
+    )
+    
+    run_button.click(
+        fn=process_image, 
+        inputs=[input_image, method, blur_intensity, blur_type], 
+        outputs=[output_image, mask_output]
+    )
+
+    # Load default models on startup
+    demo.load(
+        fn=lambda: (
+            load_segmentation_model(list(SEGMENTATION_MODELS.keys())[0]),
+            load_depth_model(list(DEPTH_MODELS.keys())[0])
+        ),
+        inputs=None,
+        outputs=[seg_model_status, depth_model_status]
+    )
 
 # Launch the app
 demo.launch()