Update app.py

app.py

@@ -1,34 +1,65 @@
 import gradio as gr
+from transformers import OneFormerProcessor, OneFormerForUniversalSegmentation
 from PIL import Image, ImageFilter
 import numpy as np
 import torch
-import cv2
-from transformers import AutoImageProcessor, AutoModelForDepthEstimation, OneFormerProcessor, OneFormerForUniversalSegmentation
+from scipy.ndimage import gaussian_filter
+
+# Load the OneFormer processor and model globally (to avoid reloading for each request)
+processor = None
+model = None
+try:
+    processor = OneFormerProcessor.from_pretrained("shi-labs/oneformer_coco_swin_large")
+    model = OneFormerForUniversalSegmentation.from_pretrained("shi-labs/oneformer_coco_swin_large")
+except Exception as e:
+    print(f"Error loading OneFormer model: {e}")
+
+def apply_gaussian_blur(image, mask, radius):
+    """Applies Gaussian blur to the background of the image."""
+    blurred_background = image.filter(ImageFilter.GaussianBlur(radius=radius))
+    img_array = np.array(image)
+    blurred_array = np.array(blurred_background)
+    foreground_mask = mask > 0
+    foreground_mask_3d = np.stack([foreground_mask] * 3, axis=-1)
+    final_image_array = np.where(foreground_mask_3d, img_array, blurred_array)
+    return Image.fromarray(final_image_array.astype(np.uint8))
+
+def apply_lens_blur(image, mask, strength):
+    """Placeholder for Lens Blur function. Will be implemented later."""
+    # Convert PIL Image to NumPy array
+    img_array = np.array(image)
+    mask_array = np.array(mask) / 255.0  # Normalize mask to 0-1
 
-# Load depth estimation model
-image_processor = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
-depth_model = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
+    # Apply a simple blur based on the mask (this is a very basic placeholder)
+    blurred_image = gaussian_filter(img_array, sigma=strength * mask_array[:, :, np.newaxis])
 
-# Load OneFormer processor and model
-processor = OneFormerProcessor.from_pretrained("shi-labs/oneformer_coco_swin_large")
-segmentation_model = OneFormerForUniversalSegmentation.from_pretrained("shi-labs/oneformer_coco_swin_large")
+    return Image.fromarray(blurred_image.astype(np.uint8))
+
+def segment_and_blur(input_image, blur_type, gaussian_radius=15, lens_strength=5):
+    """Segments the input image and applies the selected blur."""
+    if processor is None or model is None:
+        return "Error: OneFormer model not loaded."
+
+    image = input_image.convert("RGB")
+    # Rotate the image (assuming this is still needed)
+    image = image.rotate(-90, expand=True)
 
-def apply_gaussian_blur(image, foreground_label='person'):
-    """Applies Gaussian blur to the background based on a segmentation mask for the foreground."""
-    
     # Prepare input for semantic segmentation
     inputs = processor(images=image, task_inputs=["semantic"], return_tensors="pt")
 
     # Semantic segmentation
     with torch.no_grad():
-        outputs = segmentation_model(**inputs)
+        outputs = model(**inputs)
 
     # Processing semantic segmentation output
     predicted_semantic_map = processor.post_process_semantic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
     segmentation_mask = predicted_semantic_map.cpu().numpy()
 
-    # Get the mapping of class IDs to labels from the processor
-    id2label = segmentation_model.config.id2label
+    # Get the mapping of class IDs to labels
+    id2label = model.config.id2label
+
+    # Set foreground label to person
+    foreground_label = 'person'
     foreground_class_id = None
     for id, label in id2label.items():
         if label == foreground_label:
@@ -36,108 +67,39 @@ def apply_gaussian_blur(image, foreground_label='person'):
             break
 
     if foreground_class_id is None:
-        print(f"Error: Could not find the label '{foreground_label}' in the model's class mapping.")
-        return image  # Return original image if foreground label is not found
+        return f"Error: Could not find the label '{foreground_label}' in the model's class mapping."
 
-    # Create a black background mask and set the pixels corresponding to the foreground object to white
+    # Black background mask
     output_mask_array = np.zeros(segmentation_mask.shape, dtype=np.uint8)
-    output_mask_array[segmentation_mask == foreground_class_id] = 255
-
-    # Convert the output mask to a PIL Image (Grayscale)
-    mask_pil = Image.fromarray(output_mask_array, mode='L')
-    
-    # Resize the mask to match the image size
-    mask_pil = mask_pil.resize(image.size)
-    output_mask_array = np.array(mask_pil)
-
-    # Create a blurred version of the input image
-    blurred_background = image.filter(ImageFilter.GaussianBlur(radius=15))
 
-    # Convert images to NumPy arrays
-    img_array = np.array(image)
-    blurred_array = np.array(blurred_background)
-
-    # Create a boolean mask (foreground = True, background = False)
-    foreground_mask = output_mask_array > 0
-    foreground_mask_3d = np.stack([foreground_mask] * 3, axis=-1)
-
-    # Blend the original image with the blurred background
-    final_image_array = np.where(foreground_mask_3d, img_array, blurred_array)
-    final_image = Image.fromarray(final_image_array.astype(np.uint8))
-
-    return final_image
-
-def apply_lens_blur(image):
-    """Applies depth-based lens blur using a pre-trained model."""
-    
-    # Resize image to 512x512 for processing
-    resized_image = image.resize((512, 512))
-    image_np = np.array(resized_image)
+    # Set the pixels corresponding to the foreground object to white (255)
+    output_mask_array[segmentation_mask == foreground_class_id] = 255
 
-    # Prepare image for the model
-    inputs = image_processor(images=resized_image, return_tensors="pt")
+    # Convert the NumPy array to a PIL Image and resize to match input
+    mask_pil = Image.fromarray(output_mask_array, mode='L').resize(image.size)
+    mask_array = np.array(mask_pil)
 
-    with torch.no_grad():
-        outputs = depth_model(**inputs)
-        predicted_depth = outputs.predicted_depth
-
-    # Interpolate to the original size
-    prediction = torch.nn.functional.interpolate(
-        predicted_depth.unsqueeze(1),
-        size=resized_image.size[::-1],
-        mode="bicubic",
-        align_corners=False,
-    ).squeeze()
-
-    # Convert prediction to a NumPy array
-    depth_map = prediction.cpu().numpy()
-
-    # Normalize the depth map to the range 0-1
-    depth_norm = (depth_map - np.min(depth_map)) / (np.max(depth_map) - np.min(depth_map))
-
-    num_blur_levels = 5
-    blurred_layers = []
-    
-    for i in range(num_blur_levels):
-        sigma = i * 0.5
-        if sigma == 0:
-            blurred = image_np
-        else:
-            blurred = cv2.GaussianBlur(image_np, (15, 15), sigmaX=sigma, sigmaY=sigma, borderType=cv2.BORDER_REPLICATE)
-        blurred_layers.append(blurred)
-
-    depth_indices = ((1 - depth_norm) * (num_blur_levels - 1)).astype(np.uint8)
-
-    final_blurred_image = np.zeros_like(image_np)
-    
-    for y in range(image_np.shape[0]):
-        for x in range(image_np.shape[1]):
-            depth_index = depth_indices[y, x]
-            final_blurred_image[y, x] = blurred_layers[depth_index][y, x]
-
-    # Convert the final blurred image back to a PIL Image
-    final_blurred_pil_image = Image.fromarray(final_blurred_image)
-
-    return final_blurred_pil_image
-
-def process_image(image, blur_type, foreground_label='person'):
-    """Processes the image based on the selected blur type."""
-    if blur_type == "Gaussian Blur":
-        return apply_gaussian_blur(image, foreground_label=foreground_label)
+    if blur_type == "Gaussian":
+        blurred_image = apply_gaussian_blur(image, mask_array, gaussian_radius)
+    elif blur_type == "Lens":
+        blurred_image = apply_lens_blur(image, mask_array, lens_strength)
     else:
-        return apply_lens_blur(image)
+        return "Error: Invalid blur type selected."
+
+    return blurred_image
 
-interface = gr.Interface(
-    fn=process_image,
+iface = gr.Interface(
+    fn=segment_and_blur,
     inputs=[
-        gr.Image(type="pil", label="Upload an Image"),
-        gr.Radio(["Gaussian Blur", "Lens Blur"], label="Choose Blur Effect"),
-        gr.Textbox(label="Foreground Label (for Gaussian Blur)", default="person")
+        gr.Image(label="Input Image"),
+        gr.Radio(["Gaussian", "Lens"], label="Blur Type", value="Gaussian"),
+        gr.Slider(0, 30, step=1, default=15, label="Gaussian Blur Radius"),
+        gr.Slider(0, 10, step=1, default=5, label="Lens Blur Strength"),
     ],
-    outputs=[gr.Image(type="pil"), gr.Image(type="pil")],
-    title="Gaussian & Lens Blur Effects",
-    description="Upload an image and select either Gaussian blur (with foreground segmentation) or depth-based lens blur."
+    outputs=gr.Image(label="Output Image"),
+    title="Image Background Blur App",
+    description="Upload an image, select a blur type (Gaussian or Lens), and adjust the blur parameters to blur the background while keeping the person in focus."
 )
 
 if __name__ == "__main__":
-    interface.launch()
+    iface.launch()