Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from transformers import MobileViTFeatureExtractor, MobileViTForSemanticSegmentation, pipeline | |
| from PIL import Image, ImageFilter, ImageOps | |
| import numpy as np | |
| # Function to apply Gaussian Blur | |
| def apply_gaussian_blur(image, predicted_mask, sigma=15): | |
| 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 get_segmentation_mask(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 | |
| 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") | |
| st.write("Upload an image to apply segmentation, Gaussian blur, and depth-based blur.") | |
| uploaded_file = st.file_uploader("Upload an Image (PNG, JPG, JPEG)", type=["png", "jpg", "jpeg"]) | |
| if uploaded_file: | |
| image = Image.open(uploaded_file) | |
| 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, predicted_mask) | |
| st.image(blurred_image, caption="Gaussian Blurred Image", use_container_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) |