Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import cv2 | |
| import matplotlib | |
| import numpy as np | |
| import os | |
| from PIL import Image | |
| import spaces | |
| import torch | |
| import tempfile | |
| from gradio_imageslider import ImageSlider | |
| from huggingface_hub import hf_hub_download | |
| from transformers import DepthProImageProcessorFast, DepthProForDepthEstimation | |
| from depth_anything_v2.dpt import DepthAnythingV2 | |
| # Add this at the top to help with ZeroGPU authentication | |
| os.environ['SPACES_ZERO_GPU'] = '1' | |
| css = """ | |
| #img-display-container { | |
| max-height: 100vh; | |
| } | |
| #img-display-input { | |
| max-height: 80vh; | |
| } | |
| #img-display-output { | |
| max-height: 80vh; | |
| } | |
| #download { | |
| height: 62px; | |
| } | |
| """ | |
| DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| model_configs = { | |
| 'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]}, | |
| 'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]}, | |
| 'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]}, | |
| 'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]} | |
| } | |
| encoder2name = { | |
| 'vits': 'Small', | |
| 'vitb': 'Base', | |
| 'vitl': 'Large', | |
| 'vitg': 'Giant', | |
| 'depth_pro': 'Detailed', | |
| } | |
| class DepthProWrapper: | |
| def __init__(self, device): | |
| self.processor = DepthProImageProcessorFast.from_pretrained("apple/DepthPro-hf") | |
| self.model = DepthProForDepthEstimation.from_pretrained("apple/DepthPro-hf").to(device).eval() | |
| def infer_image(self, raw_image): | |
| rgb_image = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB) | |
| pil_image = Image.fromarray(rgb_image) | |
| device = next(self.model.parameters()).device | |
| inputs = self.processor(images=pil_image, return_tensors="pt").to(device) | |
| outputs = self.model(**inputs) | |
| post_processed = self.processor.post_process_depth_estimation( | |
| outputs, | |
| target_sizes=[(pil_image.height, pil_image.width)], | |
| ) | |
| depth = post_processed[0]["predicted_depth"].detach().cpu().numpy() | |
| inverse_depth = 1.0 / np.clip(depth, 0.1, 250.0) | |
| inv_min, inv_max = inverse_depth.min(), inverse_depth.max() | |
| if inv_max - inv_min > 1e-8: | |
| inverse_depth = (inverse_depth - inv_min) / (inv_max - inv_min) * 65535.0 | |
| return inverse_depth | |
| models = {} | |
| for encoder_key in model_configs.keys(): | |
| if encoder_key != 'vitg': | |
| try: | |
| model_name = encoder2name[encoder_key] | |
| model = DepthAnythingV2(**model_configs[encoder_key]) | |
| filepath = hf_hub_download(repo_id=f"depth-anything/Depth-Anything-V2-{model_name}", filename=f"depth_anything_v2_{encoder_key}.pth", repo_type="model") | |
| state_dict = torch.load(filepath, map_location="cpu") | |
| model.load_state_dict(state_dict) | |
| model = model.to(DEVICE).eval() | |
| models[encoder_key] = model | |
| print(f"Loaded {model_name} model successfully") | |
| except Exception as e: | |
| print(f"Failed to load {encoder2name[encoder_key]} model: {e}") | |
| try: | |
| models['depth_pro'] = DepthProWrapper(DEVICE) | |
| print("Loaded Detailed (DepthPro) model successfully") | |
| except Exception as e: | |
| print(f"Failed to load Detailed (DepthPro) model: {e}") | |
| default_model = 'vitl' if 'vitl' in models else list(models.keys())[0] | |
| title = "# Depthinator" | |
| description = """Edited Space for Uluç's personal works.""" | |
| def predict_depth(image, model_choice=None): | |
| if model_choice is None: | |
| model_choice = default_model | |
| return models[model_choice].infer_image(image) | |
| def process_multiple_images(files, model_choice=None, include_bw=True): | |
| if not files: | |
| return [], [] | |
| if model_choice is None: | |
| model_choice = default_model | |
| results = [] | |
| processed_files = [] | |
| for file in files: | |
| try: | |
| image = Image.open(file.name) | |
| image_np = np.array(image) | |
| if len(image_np.shape) == 3: | |
| original_image = image_np.copy() | |
| # Call the GPU function for each image | |
| depth = predict_depth(image_np[:, :, ::-1], model_choice) | |
| # Create colored depth map | |
| depth_normalized = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0 | |
| depth_normalized = depth_normalized.astype(np.uint8) | |
| colored_depth = (cmap(depth_normalized)[:, :, :3] * 255).astype(np.uint8) | |
| # Save processed files | |
| base_name = os.path.splitext(os.path.basename(file.name))[0] | |
| model_suffix = f"_{model_choice}" | |
| # Save original image temporarily for gallery | |
| original_img = Image.fromarray(original_image) | |
| original_path = tempfile.NamedTemporaryFile(suffix=f'_{base_name}_original.png', delete=False) | |
| original_img.save(original_path.name) | |
| # Save colored depth | |
| colored_depth_img = Image.fromarray(colored_depth) | |
| colored_depth_path = tempfile.NamedTemporaryFile(suffix=f'_{base_name}{model_suffix}_colored_depth.png', delete=False) | |
| colored_depth_img.save(colored_depth_path.name) | |
| # Save grayscale depth | |
| gray_depth_img = Image.fromarray(depth_normalized) | |
| gray_depth_path = tempfile.NamedTemporaryFile(suffix=f'_{base_name}{model_suffix}_gray_depth.png', delete=False) | |
| gray_depth_img.save(gray_depth_path.name) | |
| # Save black and white depth if requested | |
| if include_bw: | |
| bw_depth_img = Image.fromarray(depth_normalized).convert('L') | |
| bw_depth_path = tempfile.NamedTemporaryFile(suffix=f'_{base_name}{model_suffix}_bw_depth.png', delete=False) | |
| bw_depth_img.save(bw_depth_path.name) | |
| processed_files.append(bw_depth_path.name) | |
| # Save raw depth | |
| raw_depth_img = Image.fromarray(depth.astype('uint16')) | |
| raw_depth_path = tempfile.NamedTemporaryFile(suffix=f'_{base_name}{model_suffix}_raw_depth.png', delete=False) | |
| raw_depth_img.save(raw_depth_path.name) | |
| # Add to results gallery - use file paths instead of numpy arrays | |
| results.append((original_path.name, f"Original: {base_name}")) | |
| results.append((colored_depth_path.name, f"Depth: {base_name}")) | |
| processed_files.extend([colored_depth_path.name, gray_depth_path.name, raw_depth_path.name]) | |
| print(f"Successfully processed {base_name}") | |
| except Exception as err: | |
| print(f"Error processing {file.name}: {str(err)}") | |
| continue | |
| return results, processed_files | |
| with gr.Blocks(css=css) as demo: | |
| gr.Markdown(title) | |
| gr.Markdown(description) | |
| model_choices = [(f"{encoder2name[k]} ({k})", k) for k in models.keys()] | |
| with gr.Tabs(): | |
| with gr.TabItem("Single Image"): | |
| gr.Markdown("### Depth Prediction demo") | |
| with gr.Row(): | |
| model_dropdown_single = gr.Dropdown(choices=model_choices, value=default_model, label="Model Detail Level") | |
| include_bw_single = gr.Checkbox(value=True, label="Include Black & White Output") | |
| with gr.Row(): | |
| input_image = gr.Image(label="Input Image", type='numpy', elem_id='img-display-input') | |
| depth_image_slider = ImageSlider(label="Depth Map with Slider View", elem_id='img-display-output', position=0.5) | |
| submit = gr.Button(value="Compute Depth") | |
| gray_depth_file = gr.File(label="Grayscale depth map", elem_id="download",) | |
| bw_depth_file = gr.File(label="Black & White depth map", elem_id="download",) | |
| raw_file = gr.File(label="16-bit raw output (can be considered as disparity)", elem_id="download",) | |
| with gr.TabItem("Multiple Images"): | |
| gr.Markdown("### Batch Depth Prediction") | |
| with gr.Row(): | |
| model_dropdown_multiple = gr.Dropdown(choices=model_choices, value=default_model, label="Model Detail Level") | |
| include_bw_multiple = gr.Checkbox(value=True, label="Include Black & White Outputs") | |
| multiple_files = gr.File(label="Upload Multiple Images", file_count="multiple", file_types=["image"]) | |
| submit_multiple = gr.Button(value="Process All Images") | |
| gallery_output = gr.Gallery(label="Results (Original | Depth)", columns=2, height="auto") | |
| download_files = gr.File(label="Download All Processed Files", file_count="multiple", elem_id="download") | |
| cmap = matplotlib.colormaps.get_cmap('Spectral_r') | |
| def on_submit(image, model_choice, include_bw): | |
| original_image = image.copy() | |
| depth = predict_depth(image[:, :, ::-1], model_choice) | |
| raw_depth = Image.fromarray(depth.astype('uint16')) | |
| tmp_raw_depth = tempfile.NamedTemporaryFile(suffix='.png', delete=False) | |
| raw_depth.save(tmp_raw_depth.name) | |
| depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0 | |
| depth = depth.astype(np.uint8) | |
| colored_depth = (cmap(depth)[:, :, :3] * 255).astype(np.uint8) | |
| gray_depth = Image.fromarray(depth) | |
| tmp_gray_depth = tempfile.NamedTemporaryFile(suffix='.png', delete=False) | |
| gray_depth.save(tmp_gray_depth.name) | |
| tmp_bw_depth = None | |
| if include_bw: | |
| bw_depth = Image.fromarray(depth).convert('L') | |
| tmp_bw_depth = tempfile.NamedTemporaryFile(suffix='.png', delete=False) | |
| bw_depth.save(tmp_bw_depth.name) | |
| return [(original_image, colored_depth), tmp_gray_depth.name, tmp_bw_depth.name if tmp_bw_depth else None, tmp_raw_depth.name] | |
| submit.click(on_submit, inputs=[input_image, model_dropdown_single, include_bw_single], outputs=[depth_image_slider, gray_depth_file, bw_depth_file, raw_file]) | |
| submit_multiple.click(process_multiple_images, inputs=[multiple_files, model_dropdown_multiple, include_bw_multiple], outputs=[gallery_output, download_files]) | |
| if __name__ == '__main__': | |
| demo.queue().launch() |