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() @torch.no_grad() 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.""" @spaces.GPU 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()