Update app.py

app.py

@@ -1,64 +1,92 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+import torch
+import numpy as np
+from transformers import AutoImageProcessor, AutoModelForDepthEstimation
+from PIL import Image, ImageFilter
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+def load_depth_model():
+    """
+    Loads the depth estimation model and processor.
+    Returns (processor, model, device).
+    """
+    global processor, model, device
+    if "model" not in globals():
+        processor = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2")
+        model = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2")
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model.to(device).eval()
+    return processor, model, device
 
+def compute_depth_map(image: Image.Image, scale_factor: float) -> np.ndarray:
+    """
+    Computes the depth map for a PIL image.
+    Inverts the map (i.e. force invert_depth=True) and scales it.
+    Returns a NumPy array in [0, 1]*scale_factor.
+    """
+    processor, model, device = load_depth_model()
+    inputs = processor(images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        predicted_depth = outputs.predicted_depth
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+    prediction = torch.nn.functional.interpolate(
+        predicted_depth.unsqueeze(1),
+        size=image.size[::-1],  # PIL image size: (width, height)
+        mode="bicubic",
+        align_corners=False,
+    )
+    depth_min = prediction.min()
+    depth_max = prediction.max()
+    depth_vis = (prediction - depth_min) / (depth_max - depth_min + 1e-8)
+    depth_map = depth_vis.squeeze().cpu().numpy()
+    # Always invert depth so that near=0 and far=1
+    depth_map = 1.0 - depth_map
+    depth_map *= scale_factor
+    return depth_map
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+def layered_blur(image: Image.Image, depth_map: np.ndarray, num_layers: int, max_blur: float) -> Image.Image:
+    """
+    Creates multiple blurred versions of 'image' (radii from 0 to max_blur)
+    and composites them based on the depth map split into num_layers bins.
+    """
+    blur_radii = np.linspace(0, max_blur, num_layers)
+    blur_versions = [image.filter(ImageFilter.GaussianBlur(r)) for r in blur_radii]
+    upper_bound = depth_map.max()
+    thresholds = np.linspace(0, upper_bound, num_layers + 1)
+    final_image = blur_versions[-1]
+    for i in range(num_layers - 1, -1, -1):
+        mask_array = np.logical_and(
+            depth_map >= thresholds[i],
+            depth_map < thresholds[i + 1]
+        ).astype(np.uint8) * 255
+        mask_image = Image.fromarray(mask_array, mode="L")
+        final_image = Image.composite(blur_versions[i], final_image, mask_image)
+    return final_image
 
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
+def process_depth_blur(uploaded_image, max_blur_value, scale_factor, num_layers):
+    """
+    Processes the image with a depth-based blur.
+    The image is resized to 512x512, its depth is computed (with invert_depth always True),
+    and a layered blur is applied.
+    """
+    if not isinstance(uploaded_image, Image.Image):
+        uploaded_image = Image.open(uploaded_image)
+    image = uploaded_image.convert("RGB").resize((512, 512))
+    depth_map = compute_depth_map(image, scale_factor)
+    final_image = layered_blur(image, depth_map, int(num_layers), max_blur_value)
+    return final_image
 
+with gr.Blocks() as demo:
+    gr.Markdown("# Depth-Based Lens Blur")
+    depth_img = gr.Image(type="pil", label="Upload Image")
+    depth_max_blur = gr.Slider(1.0, 5.0, value=3.0, step=0.1, label="Maximum Blur Radius")
+    depth_scale = gr.Slider(0.1, 1.0, value=0.5, step=0.1, label="Depth Scale Factor")
+    depth_layers = gr.Slider(2, 20, value=8, step=1, label="Number of Layers")
+    depth_out = gr.Image(label="Depth-Based Blurred Image")
+    depth_button = gr.Button("Process Depth Blur")
+    depth_button.click(process_depth_blur, 
+                       inputs=[depth_img, depth_max_blur, depth_scale, depth_layers],
+                       outputs=depth_out)
 
 if __name__ == "__main__":
     demo.launch()