Blame for app.py

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()