update application file

app.py

@@ -1,25 +1,49 @@
 import streamlit as st
+from transformers import MobileViTFeatureExtractor, MobileViTForSemanticSegmentation
 from PIL import Image
-import cv2
 import numpy as np
+import cv2
+import torch
 
-st.title("Segmentation and Blur Effects")
+# 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))
 
-uploaded_image = st.file_uploader("Upload an Image", type=["png", "jpg"])
-uploaded_depth = st.file_uploader("Upload a Depth Map", type=["png", "jpg"])
+# Function to load and process image for segmentation
+def segment_image(image):
+    feature_extractor = MobileViTFeatureExtractor.from_pretrained("apple/mobilevit-small")
+    model = MobileViTForSemanticSegmentation.from_pretrained("apple/mobilevit-small")
+    
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+    
+    # 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
 
-if uploaded_image and uploaded_depth:
-    image = Image.open(uploaded_image)
-    depth_map = Image.open(uploaded_depth)
+# Streamlit interface
+st.title("Image Segmentation and Blur Effects")
+st.write("Upload an image to apply segmentation, Gaussian blur, and depth-based blur.")
 
-    # Gaussian Blur
-    sigma = st.slider("Gaussian Blur Intensity", 5, 50, 15)
-    blurred_image = cv2.GaussianBlur(np.array(image), (0, 0), sigma)
+uploaded_file = st.file_uploader("Upload an Image (PNG, JPG, JPEG)", type=["png", "jpg", "jpeg"])
 
-    # Depth-based Blur (Dummy example)
-    depth_array = np.array(depth_map).astype(np.float32)
-    depth_array = cv2.normalize(depth_array, None, 0, sigma, cv2.NORM_MINMAX)
-    blurred_depth = cv2.GaussianBlur(np.array(image), (0, 0), int(np.max(depth_array)))
-
-    st.image(blurred_image, caption="Gaussian Blur", use_column_width=True)
-    st.image(blurred_depth, caption="Depth-based Blur", use_column_width=True)
+if uploaded_file:
+    image = Image.open(uploaded_file)
+    st.image(image, caption="Uploaded Image", use_column_width=True)
+    
+    # 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)
+    
+    # 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)