# -*- coding: utf-8 -*- """ CS5330-HW4: Parallax Effect Gradio App Converted from Colab notebook. (V4: Final fix for halo/border artifact. Uses correct mask.) """ import torch import numpy as np import matplotlib.pyplot as plt from PIL import Image from transformers import DPTImageProcessor, DPTForDepthEstimation import cv2 import imageio.v2 as imageio import gradio as gr import time # To create unique filenames # ================================================================== # Global Transformer Setup # ================================================================== print("Loading Intel DPT depth estimation model...") processor = DPTImageProcessor.from_pretrained("Intel/dpt-large") model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large") model.eval() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) print(f"Model loaded on {device}. Gradio app is ready.") # ================================================================== # Helper Function 1: Get Depth Map # ================================================================== def get_depth_map(pil_image, processor, model, device): print("... (1/5) Extracting depth map") inputs = processor(images=pil_image, return_tensors="pt") inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(): outputs = model(**inputs) predicted_depth = outputs.predicted_depth prediction = torch.nn.functional.interpolate( predicted_depth.unsqueeze(1), size=pil_image.size[::-1], mode="bicubic", align_corners=False, ) depth_map = prediction.squeeze().cpu().numpy() depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min()) return depth_map # ================================================================== # Helper Function 2: Layer Separation # ================================================================== # This function returns mask_clean (hard) and mask_soft (soft/full-size) def separate_foreground_background(image, depth_map, *, assume_bgr_input=True, near_is_foreground=True, foreground_depth_is_high=True): print("... (2/5) Separating layers") if not isinstance(image, np.ndarray): image = np.array(image) if not isinstance(depth_map, np.ndarray): depth_map = np.array(depth_map) if assume_bgr_input and image.ndim == 3 and image.shape[2] == 3: image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) if image.ndim == 2: image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB) if depth_map.ndim == 3: depth_map = depth_map[:, :, 0] depth_norm = cv2.normalize(depth_map, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8) depth_smooth = cv2.GaussianBlur(depth_norm, (5, 5), 0) if near_is_foreground and foreground_depth_is_high: thresh_flag = cv2.THRESH_BINARY elif near_is_foreground and not foreground_depth_is_high: thresh_flag = cv2.THRESH_BINARY_INV elif (not near_is_foreground) and foreground_depth_is_high: thresh_flag = cv2.THRESH_BINARY_INV else: thresh_flag = cv2.THRESH_BINARY _, binary_mask = cv2.threshold(depth_smooth, 0, 255, thresh_flag + cv2.THRESH_OTSU) kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)) mask_clean = cv2.morphologyEx(binary_mask, cv2.MORPH_OPEN, kernel, iterations=1) mask_clean = cv2.morphologyEx(mask_clean, cv2.MORPH_CLOSE, kernel, iterations=2) num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask_clean, 8) if num_labels > 1: largest_label = 1 + np.argmax(stats[1:, cv2.CC_STAT_AREA]) mask_clean = (labels == largest_label).astype(np.uint8) * 255 # mask_soft is the full-size mask, which is key to fixing the artifact. mask_soft = cv2.GaussianBlur(mask_clean, (9, 9), 5).astype(np.float32) / 255.0 img_f = image.astype(np.float32) / 255.0 mask_3 = np.dstack([mask_soft]*3) foreground = np.clip(img_f * mask_3, 0, 1) background = np.clip(img_f * (1.0 - mask_3), 0, 1) foreground = (foreground * 255.0).astype(np.uint8) background = (background * 255.0).astype(np.uint8) return foreground, background, mask_clean, mask_soft # ================================================================== # Helper Function 3: Background Reconstruction # ================================================================== # This function returns final_bg (inpainted background) and alpha_no_halo (eroded mask) # Note: We no longer use alpha_no_halo for the animation, but the function is fine. def reconstruct_background(background, mask_hard, original_image_np): print("... (3/5) Reconstructing background") kernel = np.ones((7,7), np.uint8) mask_dilated = cv2.dilate(mask_hard, kernel, iterations=1) bg_inpainted = cv2.inpaint(background, mask_dilated, inpaintRadius=6, flags=cv2.INPAINT_TELEA) bg_smooth = cv2.bilateralFilter(bg_inpainted, d=9, sigmaColor=75, sigmaSpace=75) final_bg = np.where(mask_dilated[..., None] == 255, bg_smooth, background) k3 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3,3)) mask_erode = cv2.erode(mask_hard, k3, iterations=1) dist = cv2.distanceTransform(mask_erode, cv2.DIST_L2, 5) alpha_no_halo = dist / 6.0 alpha_no_halo = np.clip(alpha_no_halo, 0, 1).astype(np.float32) alpha_no_halo = alpha_no_halo[..., None] # HxWx1 return final_bg, alpha_no_halo # ================================================================== # Helper Function 4: Animation # ================================================================== # This is the animation function (from V2 logic), which is correct (uses normalization to prevent gaps). def create_multi_layer_animation( image_original, background_clean, alpha_mask, # KEY: We will pass the full-size mask_soft here depth_map, n_frames=60, parallax_strength=12, blur_strength=1.0, direction='right', zoom_center=1.10, zoom_peak=1.05 ): print(f"... (4/5) Generating {n_frames} animation frames") print(f" Params: Parallax={parallax_strength}px, Blur={blur_strength}x, Dir={direction}") h, w = image_original.shape[:2] # --- 1. Prepare motion and blur settings --- direction_map = {'right': (1, 0), 'left': (-1, 0), 'up': (0, -1), 'down': (0, 1)} dx, dy = direction_map.get(direction, (1, 0)) fg_shift = parallax_strength mid_shift = parallax_strength * 0.5 far_shift = parallax_strength * (2 / 12) base_mid_k = 9 base_far_k = 35 mid_k_raw = int(base_mid_k * blur_strength) far_k_raw = int(base_far_k * blur_strength) mid_k = (mid_k_raw + 1) if (mid_k_raw > 0 and mid_k_raw % 2 == 0) else max(1, mid_k_raw) far_k = (far_k_raw + 1) if (far_k_raw > 0 and far_k_raw % 2 == 0) else max(1, far_k_raw) mid_blur_ksize = (mid_k, mid_k) far_blur_ksize = (far_k, far_k) print(f" ...Using blur kernels: Mid={mid_blur_ksize}, Far={far_blur_ksize}") # --- 2. Prepare base masks (FG vs BG) --- # alpha_mask is now the full-size mask_soft if alpha_mask.max() > 1: alpha_mask = alpha_mask.astype(np.float32) / 255.0 if alpha_mask.ndim == 2: alpha_mask = alpha_mask[..., None] fg_mask_3ch = np.repeat(alpha_mask, 3, axis=2) # full-size foreground bg_mask_3ch = 1.0 - fg_mask_3ch # full-size background "hole" # --- 3. Create mid-ground / far-ground masks --- if depth_map.ndim == 3: depth_map = cv2.cvtColor(depth_map, cv2.COLOR_BGR2GRAY) # We find depth values inside the "background hole" (bg_mask_3ch) bg_depth_values = depth_map[alpha_mask[..., 0] < 0.5] if len(bg_depth_values) > 0: bg_split_threshold = np.percentile(bg_depth_values, 50) else: bg_split_threshold = 0.5 raw_mid_mask = (depth_map > bg_split_threshold).astype(np.float32) raw_mid_mask_smooth = cv2.GaussianBlur(raw_mid_mask, (21, 21), 0) if raw_mid_mask_smooth.ndim == 2: raw_mid_mask_smooth = raw_mid_mask_smooth[..., None] raw_mid_mask_smooth_3ch = np.repeat(raw_mid_mask_smooth, 3, axis=2) # --- 4. Generate the final 3 mutually exclusive masks --- # These three layers will perfectly cover the image with no gaps or overlaps. mid_mask_3ch = raw_mid_mask_smooth_3ch * bg_mask_3ch far_mask_3ch = (1.0 - raw_mid_mask_smooth_3ch) * bg_mask_3ch frames = [] # --- 5. Loop to generate each frame --- for i in range(n_frames): phase = (i / n_frames) * 2 * np.pi ease = np.sin(phase) zoom_range = zoom_center - zoom_peak scale = zoom_center - (zoom_range * abs(ease)) center = (w / 2, h / 2) M_scale = cv2.getRotationMatrix2D(center, 0, scale) M_fg_trans = np.float32([[1, 0, dx*ease*fg_shift], [0, 1, dy*ease*fg_shift]]) M_mid_trans = np.float32([[1, 0, dx*ease*mid_shift], [0, 1, dy*ease*mid_shift]]) M_far_trans = np.float32([[1, 0, dx*ease*far_shift], [0, 1, dy*ease*far_shift]]) # --- Layer Transforms --- fg_warped = cv2.warpAffine(image_original, M_fg_trans, (w,h), borderMode=cv2.BORDER_REFLECT_101) fg_final = cv2.warpAffine(fg_warped, M_scale, (w,h), borderMode=cv2.BORDER_REFLECT_101).astype(np.float32) mid_warped = cv2.warpAffine(background_clean, M_mid_trans, (w,h), borderMode=cv2.BORDER_REPLICATE) mid_warped_scaled = cv2.warpAffine(mid_warped, M_scale, (w,h), borderMode=cv2.BORDER_REPLICATE) mid_final = cv2.GaussianBlur(mid_warped_scaled, mid_blur_ksize, 0).astype(np.float32) far_warped = cv2.warpAffine(background_clean, M_far_trans, (w,h), borderMode=cv2.BORDER_REPLICATE) far_warped_scaled = cv2.warpAffine(far_warped, M_scale, (w,h), borderMode=cv2.BORDER_REPLICATE) far_final = cv2.GaussianBlur(far_warped_scaled, far_blur_ksize, 0).astype(np.float32) # --- Mask Transforms --- fg_mask_warped = cv2.warpAffine(fg_mask_3ch, M_fg_trans, (w,h)) fg_mask_warped = cv2.warpAffine(fg_mask_warped, M_scale, (w,h)) mid_mask_warped = cv2.warpAffine(mid_mask_3ch, M_mid_trans, (w,h)) mid_mask_warped = cv2.warpAffine(mid_mask_warped, M_scale, (w,h)) far_mask_warped = cv2.warpAffine(far_mask_3ch, M_far_trans, (w,h)) far_mask_warped = cv2.warpAffine(far_mask_warped, M_scale, (w,h)) # --- Final Composite (V2 normalization logic) --- # Re-normalize masks to prevent black borders or tiny gaps after warp. total_mask = fg_mask_warped + mid_mask_warped + far_mask_warped + 1e-6 fg_mask_warped /= total_mask mid_mask_warped /= total_mask far_mask_warped /= total_mask # Add the three layers, weighted by their masks. composite = (fg_final * fg_mask_warped) + \ (mid_final * mid_mask_warped) + \ (far_final * far_mask_warped) frame = np.clip(composite, 0, 255).astype(np.uint8) frames.append(frame) print(f"... (4/5) Frame generation complete.") return frames # ================================================================== # MAIN GRADIO FUNCTION (Ties everything together) # ================================================================== def generate_parallax_effect(input_image_np, parallax_strength, blur_strength, animation_direction): print("\n--- Processing new image ---") # --- 0. Image Preparation --- image_pil = Image.fromarray(input_image_np).convert('RGB') max_size = 640 if max(image_pil.size) > max_size: ratio = max_size / max(image_pil.size) new_size = tuple(int(dim * ratio) for dim in image_pil.size) image_pil = image_pil.resize(new_size, Image.LANCZOS) image_resized_np = np.array(image_pil) print(f"Image resized to: {image_pil.size}") # --- 1. Get Depth Map --- depth_map_0_1 = get_depth_map(image_pil, processor, model, device) # --- 2. Layer Separation --- # We get mask_soft (full-size mask) from this function. foreground, background, mask_hard, mask_soft = separate_foreground_background( image_pil, depth_map_0_1, assume_bgr_input=False, near_is_foreground=True, foreground_depth_is_high=True ) # --- 3. Background Reconstruction --- # We get final_bg (inpainted background) from this. # We also get alpha_no_halo, but we won't use it for the animation. final_bg, alpha_no_halo = reconstruct_background(background, mask_hard, image_resized_np) # --- 4. Animation --- # *** THIS IS THE KEY FIX *** # We use the V2-logic animation function (V4) with `mask_soft` (the full-size mask). multi_layer_frames = create_multi_layer_animation( image_original=image_resized_np, background_clean=final_bg, alpha_mask=mask_soft, # <-- KEY FIX: Pass the full-size soft mask depth_map=depth_map_0_1, n_frames=60, parallax_strength=parallax_strength, blur_strength=blur_strength, direction=animation_direction ) # --- 5. Save GIF and Return Path --- print("... (5/5) Saving final GIF") timestamp = int(time.time()) output_filename = f'parallax_final_{timestamp}.gif' # This saves the file to the SERVER'S disk. # It does NOT trigger a download in the user's browser. imageio.mimsave(output_filename, multi_layer_frames, duration=0.04, loop=0) print(f"--- Processing complete! Saved to {output_filename} ---") # MODIFIED: Only return the GIF filepath return output_filename # ================================================================== # Gradio Interface (Modified) # ================================================================== print("Creating Gradio interface...") # --- 1. Define Input Components --- input_image = gr.Image(label="1. Upload Your Image", type="numpy") param_parallax = gr.Slider( minimum=0, maximum=30, value=12, step=1, label="2. Parallax Strength (px)", info="Foreground motion in pixels. Higher = stronger 3D effect." ) param_blur = gr.Slider( minimum=0.0, maximum=2.0, value=1.0, step=0.1, label="3. Aperture / Blur Strength", info="Controls background blur (bokeh). 0 = no blur, 1 = default, 2 = max blur." ) param_direction = gr.Dropdown( choices=['right', 'left', 'up', 'down'], value='right', label="4. Animation Direction" ) # --- 2. Define Output Components --- # MODIFIED: Removed output_original output_gif = gr.Image(label="Generated Parallax GIF") # NOTE: The gr.Image component automatically provides a download button # in the top-right corner when displaying an image/GIF. This # fulfills the requirement for a "Gradio download button". # --- 4. Create Interface --- iface = gr.Interface( fn=generate_parallax_effect, inputs=[input_image, param_parallax, param_blur, param_direction], # MODIFIED: Only one output outputs=output_gif, title="📸 3D Parallax Photo Animator (CS5330-HW4)", description=""" Upload a photo (ideally with a clear foreground and background) to generate a 3D parallax and depth-of-field animation. 1. Upload an image. 2. Adjust the 3 parameters below. 3. Click "Submit". Processing may take 30-60 seconds. You can find the download button in the top-right corner of the generated GIF. """, delete_cache=(86400,86400) ) if __name__ == "__main__": iface.launch(share=False)