Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import pipeline | |
| import numpy as np | |
| from PIL import Image | |
| # 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 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") | |
| 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 | |
| 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__": | |
| demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| share=False | |
| ) | |