added lens blur

app.py

@@ -1,18 +1,28 @@
 import streamlit as st
-from transformers import MobileViTFeatureExtractor, MobileViTForSemanticSegmentation
-from PIL import Image
+from transformers import MobileViTFeatureExtractor, MobileViTForSemanticSegmentation, pipeline
+from PIL import Image, ImageFilter, ImageOps
 import numpy as np
-import cv2
-import torch
 
 # Function to apply Gaussian Blur
-def apply_gaussian_blur(image, sigma=15):
-    image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-    blurred = cv2.GaussianBlur(image_cv, (0, 0), sigma)
-    return Image.fromarray(cv2.cvtColor(blurred, cv2.COLOR_BGR2RGB))
+def apply_gaussian_blur(image, predicted_mask, sigma=15):
+    print(f"Predicted Mask: {predicted_mask}")
+    mask = Image.fromarray(
+        (predicted_mask.cpu().numpy() * 255).astype(np.uint8)
+    ).resize(image.size).convert('L')
+
+    foreground = Image.composite(
+        image, Image.new('RGB', image.size), mask
+    )
+    background = Image.composite(
+        image.filter(ImageFilter.GaussianBlur(sigma)), image,
+        ImageOps.invert(mask)
+    )
+
+    final_image = Image.composite(foreground, background, mask)
+    return final_image
 
 # Function to load and process image for segmentation
-def segment_image(image):
+def get_segmentation_mask(image):
     feature_extractor = MobileViTFeatureExtractor.from_pretrained("apple/mobilevit-small")
     model = MobileViTForSemanticSegmentation.from_pretrained("apple/mobilevit-small")
     
@@ -21,11 +31,35 @@ def segment_image(image):
     
     # Get segmentation mask
     logits = outputs.logits
-    upsampled_logits = torch.nn.functional.interpolate(
-        logits, size=image.size[::-1], mode="bilinear", align_corners=False
-    )
-    segmentation = upsampled_logits.argmax(dim=1).squeeze().detach().cpu().numpy()
-    return segmentation
+    predicted_mask = logits.argmax(1).squeeze(0)
+    return predicted_mask
+
+def get_depth_mask(image):
+    pipe = pipeline(task="depth-estimation", model="Intel/dpt-beit-base-384")
+
+    result = pipe(image)
+    depth_map = result["depth"]
+
+    return np.array(depth_map)
+
+def add_depth_based_blur(depth_array, image):
+    depth_normalized = (depth_array - depth_array.min()) / (depth_array.max() - depth_array.min()) * 15
+
+    image_array = np.array(image)
+
+    blurred_images = [
+        np.array(image.filter(ImageFilter.GaussianBlur(radius)))
+        for radius in range(16)
+    ]
+
+    depth_blurred_array = np.zeros_like(image_array)
+
+    for i in range(depth_normalized.shape[0]):
+        for j in range(depth_normalized.shape[1]):
+            blur_lvl = 15 - int(depth_normalized[i, j])
+            depth_blurred_array[i, j] = blurred_images[blur_lvl][i, j]
+
+    return Image.fromarray(depth_blurred_array.astype(np.uint8))
 
 # Streamlit interface
 st.title("Image Segmentation and Blur Effects")
@@ -35,15 +69,19 @@ uploaded_file = st.file_uploader("Upload an Image (PNG, JPG, JPEG)", type=["png"
 
 if uploaded_file:
     image = Image.open(uploaded_file)
-    st.image(image, caption="Uploaded Image", use_column_width=True)
+    image = image.resize((512, 512))
+    st.image(image, caption="Uploaded Image", use_container_width=True)
+
+    predicted_mask = get_segmentation_mask(image)
     
     # Apply Gaussian Blur
     sigma = st.slider("Gaussian Blur Intensity", 5, 50, 15)
-    blurred_image = apply_gaussian_blur(image, sigma)
-    st.image(blurred_image, caption="Gaussian Blurred Image", use_column_width=True)
+    blurred_image = apply_gaussian_blur(image, predicted_mask, sigma)
+    st.image(blurred_image, caption="Gaussian Blurred Image", use_container_width=True)
     
-    # Perform segmentation
-    if st.button("Perform Segmentation"):
-        st.write("Segmenting the image...")
-        segmentation = segment_image(image)
-        st.image(segmentation, caption="Segmentation Mask", use_column_width=True)
+    # Perform lens blur
+    st.write("Calculating depth and applying lens blur...")
+    with st.spinner("Applying lens blur... This might take a few moments."):
+        depth_array = get_depth_mask(image)
+        lens_blurred_img = add_depth_based_blur(depth_array, image)
+    st.image(lens_blurred_img, caption="Lens Blur Effect Image", use_container_width=True)

requirements.txt

@@ -1,5 +1,4 @@
 streamlit
 transformers
 torch
-pillow
-opencv-python
+pillow