Update app.py

app.py

@@ -1,62 +1,103 @@
+import os
+import torch
 import gradio as gr
 from PIL import Image, ImageFilter
-import torch
-from transformers import DepthProImageProcessorFast, DepthProForDepthEstimation
+import torchvision.transforms as transforms
+from transformers import AutoModelForImageSegmentation, DepthProImageProcessorFast, DepthProForDepthEstimation
 import numpy as np
+import io
 
-# Load the device (use CPU or GPU)
+# Load Models
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Initialize the model and processor
-image_processor = DepthProImageProcessorFast.from_pretrained("apple/DepthPro-hf")
-model = DepthProForDepthEstimation.from_pretrained("apple/DepthPro-hf").to(device)
+HF_model_name = 'BiRefNet'
+birefnet = AutoModelForImageSegmentation.from_pretrained(f'zhengpeng7/{HF_model_name}', trust_remote_code=True).to(device).eval()
+print('BiRefNet (Segmentation) is ready to use.')
 
-# Function to apply background blur based on depth
-def apply_background_blur(image: Image):
-    # Convert the uploaded image to RGB if necessary
-    image = image.convert("RGB")
-    
-    # Process the image with DepthPro model
-    inputs = image_processor(images=image, return_tensors="pt").to(device)
-    
-    with torch.no_grad():
-        outputs = model(**inputs)
-    
-    post_processed_output = image_processor.post_process_depth_estimation(
-        outputs, target_sizes=[(image.height, image.width)],
-    )
-
-    # Get the predicted depth and normalize it
-    depth = post_processed_output[0]["predicted_depth"]
-    depth_np = depth.detach().cpu().numpy().squeeze()
-    depth_normalized = (depth_np - depth_np.min()) / (depth_np.max() - depth_np.min())
-
-    # Create a blurred image
-    blurred_image = image.copy()
+depth_processor = DepthProImageProcessorFast.from_pretrained("apple/DepthPro-hf")
+depth_model = DepthProForDepthEstimation.from_pretrained("apple/DepthPro-hf").to(device).eval()
+print('DepthPro (Blur) is ready to use.')
+
+# Combined Image Transform
+transform_image = transforms.Compose([
+    transforms.Resize((1024, 1024)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
 
-    # Apply variable Gaussian blur based on depth
-    blur_strength = 20  # You can adjust this for overall blur strength
-    blur_map = (depth_normalized * blur_strength).astype(int)
+# Refine Foreground (Placeholder)
+def refine_foreground(image, mask):
+    return image # Implement your refinement logic here
+
+# Segmentation Function
+def segment_image(image):
+    input_image = transform_image(image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        pred = birefnet(input_image)[-1].sigmoid().cpu()[0].squeeze()
+    mask = transforms.ToPILImage()(pred).resize(image.size)
+    image_masked = refine_foreground(image.copy(), mask)
+    image_masked.putalpha(mask)
+    return image_masked
 
-    for radius in range(1, blur_strength + 1):
+# Blur Function
+def apply_background_blur(image):
+    inputs = depth_processor(images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        depth = depth_processor.post_process_depth_estimation(depth_model(**inputs), target_sizes=[(image.height, image.width)])[0]["predicted_depth"][0].cpu().squeeze()
+    depth_normalized = (depth - depth.min()) / (depth.max() - depth.min())
+    blur_map = (depth_normalized * 20).astype(int)
+    blurred_image = image.copy()
+    for radius in range(1, 21):
         mask = (blur_map == radius)
         if np.any(mask):
-            temp_image = image.copy()
-            temp_image = temp_image.filter(ImageFilter.GaussianBlur(radius))
-            blurred_image = Image.composite(temp_image, blurred_image, Image.fromarray((mask * 255).astype(np.uint8)))
-    
+            blurred_image = Image.composite(image.copy().filter(ImageFilter.GaussianBlur(radius)), blurred_image, Image.fromarray((mask * 255).astype(np.uint8)))
     return blurred_image
 
-# Create Gradio interface
-def create_interface():
-    # Gradio interface with image upload input and output for processed image
-    gr.Interface(
-        fn=apply_background_blur,
-        inputs=gr.Image(type="pil", label="Upload Image"),
-        outputs=gr.Image(type="pil", label="Blurred Image"),
-        live=True
-    ).launch()
-
-# Start the app
-if __name__ == "__main__":
-    create_interface()
+# Process Image Function
+def process_image(image, action):
+    image = image.convert("RGB")
+    if action == "Segmentation":
+        return segment_image(image)
+    elif action == "Blur":
+        return apply_background_blur(image)
+    elif action == "Both":
+        return segment_image(image), apply_background_blur(image)
+    else:
+        return None
+
+# Download Function
+def download_image(image):
+    if image is None:
+        return None
+    if isinstance(image, tuple):
+        images = []
+        for img in image:
+            img_byte_arr = io.BytesIO()
+            img.save(img_byte_arr, format='PNG')
+            images.append(img_byte_arr.getvalue())
+        return images
+    else:
+        img_byte_arr = io.BytesIO()
+        image.save(img_byte_arr, format='PNG')
+        return img_byte_arr.getvalue()
+
+# Gradio Interface
+def gradio_interface(image, action):
+    result = process_image(image, action)
+    if action == "Both":
+        return download_image(result), result[0], result[1]
+    else:
+        return download_image(result), result
+
+interface = gr.Interface(
+    fn=gradio_interface,
+    inputs=[gr.Image(type="pil", label="Upload Image"), gr.Dropdown(["Segmentation", "Blur", "Both"], label="Select Action")],
+    outputs=[
+        gr.File(label="Download Output"),
+        gr.Image(label="Output Image 1"),
+        gr.Image(label="Output Image 2", visible=False)
+    ],
+    live=False,
+)
+
+interface.launch()