Spaces:
Running
Running
| import numpy as np | |
| import torch | |
| from torchvision.transforms.functional import normalize | |
| import torch.nn.functional as F | |
| from skimage import io | |
| from PIL import Image, ImageFilter | |
| from transformers import AutoModelForImageSegmentation, AutoImageProcessor, AutoModelForDepthEstimation | |
| import streamlit as st | |
| # Load models | |
| segmentation_model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4", trust_remote_code=True) | |
| depth_model = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf") | |
| depth_processor = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf") | |
| device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
| segmentation_model.to(device) | |
| depth_model.to(device) | |
| # Preprocessing function for segmentation | |
| def preprocess_image(im: np.ndarray, model_input_size: list) -> torch.Tensor: | |
| if len(im.shape) < 3: | |
| im = im[:, :, np.newaxis] | |
| im_tensor = torch.tensor(im, dtype=torch.float32).permute(2, 0, 1) | |
| im_tensor = F.interpolate(torch.unsqueeze(im_tensor, 0), size=model_input_size, mode='bilinear') | |
| image = torch.divide(im_tensor, 255.0) | |
| image = normalize(image, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0]) | |
| return image | |
| # Postprocessing function for segmentation mask | |
| def postprocess_image(result: torch.Tensor, im_size: list) -> np.ndarray: | |
| result = torch.squeeze(F.interpolate(result, size=im_size, mode='bilinear'), 0) | |
| ma = torch.max(result) | |
| mi = torch.min(result) | |
| result = (result - mi) / (ma - mi) | |
| im_array = (result * 255).permute(1, 2, 0).cpu().data.numpy().astype(np.uint8) | |
| im_array = np.squeeze(im_array) | |
| return im_array | |
| # Streamlit UI | |
| st.title("Blur Effects App") | |
| st.markdown("Choose between Gaussian blur (segmentation-based) or depth-based lens blur.") | |
| # File uploader | |
| uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"]) | |
| if uploaded_file: | |
| # Load the uploaded image | |
| orig_image = Image.open(uploaded_file).convert("RGB") | |
| orig_im_size = orig_image.size | |
| orig_im_np = np.array(orig_image) | |
| # Display original image | |
| st.image(orig_image, caption="Uploaded Image", use_column_width=True) | |
| # Effect selection | |
| effect_option = st.radio("Choose an effect", ("Gaussian Blur (Segmentation)", "Lens Blur (Depth Map)")) | |
| if effect_option == "Gaussian Blur (Segmentation)": | |
| # Gaussian Blur (Segmentation-Based) | |
| st.subheader("Gaussian Blur (Segmentation)") | |
| # Preprocess image for segmentation | |
| model_input_size = [256, 256] # Resize for model compatibility | |
| image = preprocess_image(orig_im_np, model_input_size).to(device) | |
| # Inference using the segmentation model | |
| with torch.no_grad(): | |
| result = segmentation_model(image) | |
| # Postprocess result to generate mask | |
| result_image = postprocess_image(result[0][0], orig_im_size) | |
| pil_mask_im = Image.fromarray(result_image) | |
| # Create binary mask for background and foreground separation | |
| binary_mask = pil_mask_im.point(lambda p: 255 if p > 170 else 0, '1') | |
| # Gaussian blur for background | |
| blur_intensity = st.slider("Blur Intensity (σ)", min_value=1, max_value=30, value=15) | |
| blurred_background = orig_image.filter(ImageFilter.GaussianBlur(blur_intensity)) | |
| # Combine blurred background and sharp foreground | |
| binary_mask_rgba = binary_mask.convert("L").resize(orig_image.size) | |
| foreground = Image.composite(orig_image, blurred_background, binary_mask_rgba) | |
| # Display results side by side | |
| st.image([orig_image, foreground], caption=["Original Image", "Blurred Background"], use_column_width=True) | |
| elif effect_option == "Lens Blur (Depth Map)": | |
| # Lens Blur (Depth-Based) | |
| st.subheader("Lens Blur (Depth Map)") | |
| # Resize image for depth model processing | |
| inputs = depth_processor(images=orig_image, return_tensors="pt") | |
| # Generate depth map | |
| with torch.no_grad(): | |
| outputs = depth_model(**inputs) | |
| depth = outputs.predicted_depth.squeeze().cpu().numpy() | |
| depth = np.array(Image.fromarray(depth).resize(orig_image.size, resample=Image.BICUBIC)) | |
| # Resize depth map back to the original image size | |
| depth_min, depth_max = np.min(depth), np.max(depth) | |
| manual_depth_min = depth_min/10 | |
| manual_depth_max = depth_max | |
| normalized_depth = (depth - depth_min) / (depth_max - depth_min) | |
| adjusted_depth = np.clip((normalized_depth * (manual_depth_max - manual_depth_min)) + manual_depth_min, 0, 1) | |
| # Create depth-based blur effect | |
| max_blur = st.slider("Max Blur Intensity", min_value=1, max_value=10, value=3) | |
| depth_array = ((1 - adjusted_depth) * max_blur).astype(np.uint8) | |
| # Create blurred images for depth-based blurring | |
| blurred_images = [orig_image.filter(ImageFilter.GaussianBlur(i)) for i in range(max_blur + 1)] | |
| final_image = Image.new("RGB", orig_image.size) | |
| # Apply depth-based blur pixel by pixel | |
| for y in range(orig_image.height): | |
| for x in range(orig_image.width): | |
| blur_level = min(max_blur, max(0, depth_array[y, x])) | |
| final_image.putpixel((x, y), blurred_images[blur_level].getpixel((x, y))) | |
| # Display results side by side | |
| st.image([orig_image, final_image], caption=["Original Image", "Lens Blur Image"], use_column_width=True) | |