Update app.py

app.py

@@ -23,49 +23,49 @@ depth_model.to(device)
 
 def process_image(image_pil):
     image = ImageOps.exif_transpose(image_pil).resize((512, 512)).convert("RGB")
-
-    # ---------- Part 1: Segmentation ----------
-    seg_inputs = seg_extractor(images=image, return_tensors="pt").to(device)
+    
+    # ---- Segmentation ----
+    seg_inputs = seg_extractor(images=image, return_tensors="pt", do_resize=True, do_normalize=True)
     with torch.no_grad():
-        seg_output = seg_model(**seg_inputs).logits
+        seg_output = seg_model(**seg_inputs.to(device)).logits
     seg_mask = torch.argmax(seg_output, dim=1)[0].cpu().numpy()
     binary_mask = np.where(seg_mask > 0, 255, 0).astype(np.uint8)
     foreground_mask = Image.fromarray(binary_mask).convert("L")
 
-    # ---------- Part 2: Gaussian blur to background ----------
-    blurred_background = image.filter(ImageFilter.GaussianBlur(15))
-    blurred_background = blurred_background.convert("RGBA")
+    # ---- Blur Background ----
     image_rgba = image.convert("RGBA")
-    output_blur = Image.composite(image_rgba, blurred_background, foreground_mask)
+    blurred = image.filter(ImageFilter.GaussianBlur(15)).convert("RGBA")
+    composite_blur = Image.composite(image_rgba, blurred, foreground_mask)
 
-    # ---------- Part 3: Depth Estimation ----------
+    # ---- Depth ----
     image_np = np.array(image)
-    depth_inputs = depth_extractor(images=image_np, return_tensors="pt").to(device)
+    depth_inputs = depth_extractor(images=image_np, return_tensors="pt")
     with torch.no_grad():
-        depth_output = depth_model(**depth_inputs)
+        depth_output = depth_model(**depth_inputs.to(device))
     predicted_depth = depth_output.predicted_depth.squeeze().cpu().numpy()
     normalized_depth = (predicted_depth - predicted_depth.min()) / (predicted_depth.max() - predicted_depth.min())
 
-    # ---------- Part 4: Depth-Based Variable Gaussian Blur ----------
-    image_np_float = image_np.astype(np.float32)
+    # ---- Depth-Based Blur ----
+    image_np = np.array(image).astype(np.float32)
     resized_depth = cv2.resize(normalized_depth, (image_np.shape[1], image_np.shape[0]))
     inverted_depth = 1.0 - resized_depth
-    total_blur_levels = 4
-    blurred_versions = []
-    for i in range(total_blur_levels):
+    blur_levels = 4
+    blurred_variants = []
+    for i in range(blur_levels):
         sigma = i * 3
-        blurred = cv2.GaussianBlur(image_np_float, (15, 15), sigmaX=sigma, sigmaY=sigma) if sigma > 0 else image_np_float.copy()
-        blurred_versions.append(blurred)
+        blurred = cv2.GaussianBlur(image_np, (15, 15), sigmaX=sigma, sigmaY=sigma) if sigma > 0 else image_np.copy()
+        blurred_variants.append(blurred)
 
-    blur_indices = (inverted_depth * (total_blur_levels - 1)).astype(np.uint8)
-    final_blurred_np = np.zeros_like(image_np_float)
-    for i in range(total_blur_levels):
+    blur_indices = (inverted_depth * (blur_levels - 1)).astype(np.uint8)
+    final_blur = np.zeros_like(image_np)
+    for i in range(blur_levels):
         mask = (blur_indices == i)
         for c in range(3):
-            final_blurred_np[:, :, c][mask] = blurred_versions[i][:, :, c][mask]
-    depth_blur_img = Image.fromarray(np.clip(final_blurred_np, 0, 255).astype(np.uint8))
+            final_blur[:, :, c][mask] = blurred_variants[i][:, :, c][mask]
+    lens_blur_pil = Image.fromarray(np.clip(final_blur, 0, 255).astype(np.uint8))
+
+    return image, composite_blur.convert("RGB"), lens_blur_pil
 
-    return image, output_blur.convert("RGB"), depth_blur_img
 
 # Gradio Interface
 gr.Interface(