Implement chromostereopsis effect using Depth Anything V2 Large model

app.py

@@ -1,30 +1,237 @@
 import gradio as gr
 from transformers import pipeline
+import numpy as np
 from PIL import Image
 
-# Load the depth estimation pipeline using Depth Anything V2 Large model
+# Load the depth estimation pipeline using the Depth Anything V2 Large model
 depth_pipeline = pipeline(task="depth-estimation", model="depth-anything/Depth-Anything-V2-Large-hf")
 
+# Global variables to store current state
+current_original_image = None
+current_depth_image = None
 
-def predict_depth(image):
-    # image is a PIL Image
-    result = depth_pipeline(image)
-    depth_image = result["depth"]
-    return depth_image
+def process_depth_to_pil(depth_pil):
+    """
+    Normalize the depth PIL image to a grayscale 0-255 PIL image.
+    """
+    depth_array = np.array(depth_pil, dtype=np.float32)
+    # If the depth image has multiple channels, take the first channel
+    if depth_array.ndim == 3:
+        depth_array = depth_array[..., 0]
+    depth_min = depth_array.min()
+    depth_max = depth_array.max()
+    depth_norm = (depth_array - depth_min) / (depth_max - depth_min + 1e-6)
+    depth_255 = depth_norm * 255.0
+    depth_img = Image.fromarray(depth_255.astype(np.uint8), mode="L")
+    return depth_img
 
+def apply_chromo_stereopsis(
+    original_img: Image.Image,
+    depth_img: Image.Image,
+    threshold: float,
+    feather: float,
+    black_level: float,
+    white_level: float
+) -> Image.Image:
+    """
+    Apply chromostereopsis red-blue effect to the original image using a depth map.
+    """
+    # Convert original image to grayscale
+    gray = np.array(original_img.convert("L"), dtype=np.float32)
+    # Levels adjustment
+    denom = (white_level - black_level) if (white_level > black_level) else 1e-6
+    adjusted_gray = (gray - black_level) / denom * 255.0
+    adjusted_gray = np.clip(adjusted_gray, 0, 255)
+    adjusted_gray_01 = adjusted_gray / 255.0
+    # Depth-based blend factor
+    depth_arr = np.array(depth_img, dtype=np.float32)
+    half_feather = feather / 2.0
+    blend = np.clip((depth_arr - (threshold - half_feather)) / (feather + 1e-6), 0, 1)
+    # Red / Blue channels
+    red = blend * adjusted_gray_01 * 255.0
+    blue = (1.0 - blend) * adjusted_gray_01 * 255.0
+    # Create the output image
+    h, w = gray.shape
+    output = np.zeros((h, w, 3), dtype=np.uint8)
+    output[..., 0] = np.clip(red, 0, 255).astype(np.uint8)   # Red channel
+    output[..., 1] = 0                                       # Green channel (always 0)
+    output[..., 2] = np.clip(blue, 0, 255).astype(np.uint8)  # Blue channel
+    return Image.fromarray(output, mode="RGB")
 
-title = "Depth Anything V2 Large - Depth Map Generator"
-description = "Upload an image to generate a depth map using the Depth Anything V2 Large model (Relative depth)."
+def generate_depth_map(input_image):
+    """
+    Generate a depth map from the input image and create an initial chromostereopsis output.
+    """
+    global current_original_image, current_depth_image
+    if input_image is None:
+        current_original_image = None
+        current_depth_image = None
+        return None, None
+    try:
+        # Store original image
+        current_original_image = input_image
+        # Run depth estimation
+        result = depth_pipeline(input_image)
+        depth_pil = result["depth"]
+        # Normalize depth map
+        current_depth_image = process_depth_to_pil(depth_pil)
+        # Use default parameters for initial chromostereopsis effect
+        threshold = 50.0 * 255.0 / 100.0
+        feather = 10.0 * 255.0 / 100.0
+        black_level = 0.0
+        white_level = 255.0
+        chromostereo_result = apply_chromo_stereopsis(
+            current_original_image,
+            current_depth_image,
+            threshold,
+            feather,
+            black_level,
+            white_level
+        )
+        return current_depth_image.convert("RGB"), chromostereo_result
+    except Exception as e:
+        print(f"Error during depth generation: {e}")
+        current_original_image = None
+        current_depth_image = None
+        return None, None
 
+def update_chromostereopsis(threshold_percent, feather_percent, black_level, white_level):
+    """
+    Update the chromostereopsis effect with new parameters for live preview.
+    """
+    global current_original_image, current_depth_image
+    if current_original_image is None or current_depth_image is None:
+        return None
+    try:
+        threshold = (threshold_percent / 100.0) * 255.0
+        feather = (feather_percent / 100.0) * 255.0
+        chromostereo_result = apply_chromo_stereopsis(
+            current_original_image,
+            current_depth_image,
+            threshold,
+            feather,
+            black_level,
+            white_level
+        )
+        return chromostereo_result
+    except Exception as e:
+        print(f"Error during chromostereopsis update: {e}")
+        return None
 
-iface = gr.Interface(
-    fn=predict_depth,
-    inputs=gr.Image(type="pil", label="Input Image"),
-    outputs=gr.Image(type="pil", label="Predicted Depth Map"),
-    title=title,
-    description=description,
-)
+def clear_results():
+    """
+    Clear stored images and reset outputs.
+    """
+    global current_original_image, current_depth_image
+    current_original_image = None
+    current_depth_image = None
+    return None, None
 
+with gr.Blocks(title="ChromoStereoizer", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# ChromoStereoizer")
+    with gr.Row():
+        # Left column: Input image and generate button
+        with gr.Column(scale=1):
+            input_image = gr.Image(
+                label="Upload Image",
+                type="pil",
+                height=400
+            )
+            generate_btn = gr.Button(
+                "Generate Depth Map",
+                variant="primary",
+                size="lg"
+            )
+        # Right column: Outputs and controls
+        with gr.Column(scale=1):
+            gr.Markdown("**Depth Map**")
+            depth_output = gr.Image(
+                type="pil",
+                height=400,
+                interactive=False,
+                show_download_button=True,
+                show_label=False
+            )
+            gr.Markdown("**ChromoStereoizer Result**")
+            chromo_output = gr.Image(
+                type="pil",
+                height=400,
+                interactive=False,
+                show_download_button=True,
+                show_label=False
+            )
+            with gr.Row():
+                threshold_slider = gr.Slider(
+                    minimum=1,
+                    maximum=100,
+                    value=50,
+                    step=1,
+                    label="Threshold (%)"
+                )
+                feather_slider = gr.Slider(
+                    minimum=0,
+                    maximum=100,
+                    value=10,
+                    step=1,
+                    label="Feather (%)"
+                )
+            with gr.Row():
+                black_level_slider = gr.Slider(
+                    minimum=0,
+                    maximum=255,
+                    value=0,
+                    step=1,
+                    label="Black Level"
+                )
+                white_level_slider = gr.Slider(
+                    minimum=0,
+                    maximum=255,
+                    value=255,
+                    step=1,
+                    label="White Level"
+                )
+            with gr.Row():
+                clear_btn = gr.Button("Clear", variant="secondary")
+    # Event bindings
+    generate_btn.click(
+        fn=generate_depth_map,
+        inputs=[input_image],
+        outputs=[depth_output, chromo_output],
+        show_progress=True
+    )
+    threshold_slider.change(
+        fn=update_chromostereopsis,
+        inputs=[threshold_slider, feather_slider, black_level_slider, white_level_slider],
+        outputs=chromo_output,
+        show_progress=False
+    )
+    feather_slider.change(
+        fn=update_chromostereopsis,
+        inputs=[threshold_slider, feather_slider, black_level_slider, white_level_slider],
+        outputs=chromo_output,
+        show_progress=False
+    )
+    black_level_slider.change(
+        fn=update_chromostereopsis,
+        inputs=[threshold_slider, feather_slider, black_level_slider, white_level_slider],
+        outputs=chromo_output,
+        show_progress=False
+    )
+    white_level_slider.change(
+        fn=update_chromostereopsis,
+        inputs=[threshold_slider, feather_slider, black_level_slider, white_level_slider],
+        outputs=chromo_output,
+        show_progress=False
+    )
+    clear_btn.click(
+        fn=clear_results,
+        inputs=[],
+        outputs=[depth_output, chromo_output]
+    )
 
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )