Update app.py

app.py

@@ -2,45 +2,11 @@ import gradio as gr
 from transformers import pipeline
 from PIL import Image, ImageFilter, ImageOps
 import numpy as np
-import requests
 import cv2
 
-# 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}"
+# Initialize models with fixed choices
+segmentation_model = pipeline("image-segmentation", model="nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
+depth_estimator = pipeline("depth-estimation", model="Intel/zoedepth-nyu-kitti")
 
 def lens_blur(image, radius):
     """
@@ -98,10 +64,6 @@ 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")
     
@@ -114,24 +76,24 @@ def process_image(input_image, method, blur_intensity, blur_type):
         blur_fn = lambda img, rad: img.filter(ImageFilter.GaussianBlur(radius=rad))
     
     if method == "Segmented Background Blur":
-        # Use segmentation to obtain a foreground mask.
+        # Use segmentation to obtain a foreground mask
         results = segmentation_model(input_image)
-        # Assume the last result is the main foreground object.
+        # Assume the last result is the main foreground object
         foreground_mask = results[-1]["mask"]
-        # Ensure the mask is grayscale.
+        # Ensure the mask is grayscale
         foreground_mask = foreground_mask.convert("L")
-        # Threshold to create a binary mask.
+        # Threshold to create a binary mask
         binary_mask = foreground_mask.point(lambda p: 255 if p > 128 else 0)
         
-        # Blur the background using the selected blur function.
+        # Blur the background using the selected blur function
         blurred_background = blur_fn(input_image, blur_intensity)
         
-        # Composite the final image: keep foreground and use blurred background elsewhere.
+        # Composite the final image: keep foreground and use blurred background elsewhere
         output_image = Image.composite(input_image, blurred_background, binary_mask)
         mask_image = binary_mask
         
     elif method == "Depth-based Variable Blur":
-        # Generate depth map.
+        # Generate depth map
         depth_results = depth_estimator(input_image)
         depth_map = depth_results["depth"]
         
@@ -141,16 +103,16 @@ def process_image(input_image, method, blur_intensity, blur_type):
         normalized_depth = (norm * 255).astype(np.uint8)
         mask_image = Image.fromarray(normalized_depth)
         
-        # Create fully blurred version using the selected blur function.
+        # Create fully blurred version using the selected blur function
         blurred_image = blur_fn(input_image, blur_intensity)
         
-        # Convert images to arrays for blending.
+        # Convert images to arrays for blending
         orig_np = np.array(input_image).astype(np.float32)
         blur_np = np.array(blurred_image).astype(np.float32)
-        # Reshape mask for broadcasting.
+        # Reshape mask for broadcasting
         alpha = normalized_depth[..., np.newaxis] / 255.0
         
-        # Blend pixels: 0 = original; 1 = fully blurred.
+        # Blend pixels: 0 = original; 1 = fully blurred
         blended_np = (1 - alpha) * orig_np + alpha * blur_np
         blended_np = np.clip(blended_np, 0, 255).astype(np.uint8)
         output_image = Image.fromarray(blended_np)
@@ -165,71 +127,28 @@ 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.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]
-    )
-    
+    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 handler
     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()