app.py

import gradio as gr
import os
import cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
import tempfile
import io

from depth_anything.dpt import DepthAnything_AC


def normalize_depth(disparity_tensor):
    """Standard normalization method to convert disparity to depth"""
    eps = 1e-6
    disparity_min = disparity_tensor.min()
    disparity_max = disparity_tensor.max()
    normalized_disparity = (disparity_tensor - disparity_min) / (disparity_max - disparity_min + eps)
    return normalized_disparity


def is_video_file(filepath):
    """Check if the given file is a video file based on its extension"""
    if filepath is None:
        return False
    video_extensions = ['.mp4', '.avi', '.mov', '.mkv', '.flv', '.wmv', '.webm', '.m4v']
    _, ext = os.path.splitext(filepath.lower())
    return ext in video_extensions


def load_model(model_path='checkpoints/depth_anything_AC_vits.pth', encoder='vits'):
    """Load trained depth estimation model"""
    model_configs = {
        'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024], 'version': 'v2'},
        'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768], 'version': 'v2'},
        'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384], 'version': 'v2'}
    }
    
    model = DepthAnything_AC(model_configs[encoder])
    
    if os.path.exists(model_path):
        checkpoint = torch.load(model_path, map_location='cpu')
        model.load_state_dict(checkpoint, strict=False)
    else:
        print(f"Warning: Model file {model_path} not found")
        
    model.eval()
    if torch.cuda.is_available():
        model.cuda()
    
    return model


def preprocess_image(image, target_size=518):
    """Preprocess input image (supports both PIL Image and numpy array)"""
    if isinstance(image, str):
        raw_image = cv2.imread(image)
        if raw_image is None:
            raise ValueError(f"Cannot read image: {image}")
        image = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
    elif isinstance(image, Image.Image):
        image = np.array(image)
        image = image.astype(np.float32) / 255.0
    elif isinstance(image, np.ndarray):
        if image.dtype == np.uint8:
            image = image.astype(np.float32) / 255.0
    else:
        raise ValueError(f"Unsupported image type: {type(image)}")
    
    if len(image.shape) == 3 and image.shape[2] == 3:
        pass
    elif len(image.shape) == 3 and image.shape[2] == 4:
        image = image[:, :, :3]
    
    h, w = image.shape[:2]
    scale = target_size / min(h, w)
    new_h, new_w = int(h * scale), int(w * scale)
    
    new_h = ((new_h + 13) // 14) * 14
    new_w = ((new_w + 13) // 14) * 14
    image = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_CUBIC)

    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    image = (image - mean) / std

    image = torch.from_numpy(image.transpose(2, 0, 1)).float()
    image = image.unsqueeze(0)
    
    return image, (h, w)


def postprocess_depth(depth_tensor, original_size):
    """Post-process depth map"""
    if depth_tensor.dim() == 3:
        depth_tensor = depth_tensor.unsqueeze(1)
    elif depth_tensor.dim() == 2:
        depth_tensor = depth_tensor.unsqueeze(0).unsqueeze(1)
    
    h, w = original_size
    depth = F.interpolate(depth_tensor, size=(h, w), mode='bilinear', align_corners=True)
    depth = depth.squeeze().cpu().numpy()
    
    return depth


def create_colored_depth_map(depth, colormap='spectral'):
    """Create colored depth map"""
    if colormap == 'inferno':
        depth_colored = cv2.applyColorMap((depth * 255).astype(np.uint8), cv2.COLORMAP_INFERNO)
        depth_colored = cv2.cvtColor(depth_colored, cv2.COLOR_BGR2RGB)
    elif colormap == 'spectral':
        from matplotlib import cm
        spectral_cmap = cm.get_cmap('Spectral_r')
        depth_colored = (spectral_cmap(depth) * 255).astype(np.uint8)
        depth_colored = depth_colored[:, :, :3]
    else:
        depth_colored = (depth * 255).astype(np.uint8)
        depth_colored = np.stack([depth_colored] * 3, axis=2)
    
    return depth_colored


def process_video(video_path, colormap_choice, progress=gr.Progress()):
    """Process video file for depth estimation"""
    try:
        print(f"Processing video: {video_path}")
        
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise ValueError(f"Cannot open video file: {video_path}")
        
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        input_fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        
        print(f"Video properties: {total_frames} frames, {input_fps} FPS, {width}x{height}")
        
        temp_output = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
        output_path = temp_output.name
        temp_output.close()
        
        fourcc = cv2.VideoWriter.fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, input_fps, (width, height))
        
        if not out.isOpened():
            cap.release()
            raise ValueError("Cannot create output video file")
        
        frame_count = 0
        
        try:
            while True:
                ret, frame = cap.read()
                if not ret:
                    break
                
                frame_count += 1
                
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                
                try:
                    image_tensor, original_size = preprocess_image(frame_rgb)
                    
                    if torch.cuda.is_available():
                        image_tensor = image_tensor.cuda()
                    
                    with torch.no_grad():
                        prediction = model(image_tensor)
                        disparity_tensor = prediction['out']
                        depth_tensor = normalize_depth(disparity_tensor)
                    
                    depth = postprocess_depth(depth_tensor, original_size)
                    
                    if depth is None:
                        if depth_tensor.dim() == 1:
                            h, w = original_size
                            expected_size = h * w
                            if depth_tensor.shape[0] == expected_size:
                                depth_tensor = depth_tensor.view(1, 1, h, w)
                            else:
                                import math
                                side_length = int(math.sqrt(depth_tensor.shape[0]))
                                if side_length * side_length == depth_tensor.shape[0]:
                                    depth_tensor = depth_tensor.view(1, 1, side_length, side_length)
                        depth = postprocess_depth(depth_tensor, original_size)
                    
                    if depth is None:
                        print(f"Warning: Frame {frame_count} processing failed, using black frame")
                        depth_frame = np.zeros((height, width, 3), dtype=np.uint8)
                    else:
                        if colormap_choice.lower() == 'inferno':
                            depth_frame = cv2.applyColorMap((depth * 255).astype(np.uint8), cv2.COLORMAP_INFERNO)
                        elif colormap_choice.lower() == 'spectral':
                            from matplotlib import cm
                            spectral_cmap = cm.get_cmap('Spectral_r')
                            depth_frame = (spectral_cmap(depth) * 255).astype(np.uint8)
                            depth_frame = depth_frame[:, :, :3]
                            depth_frame = cv2.cvtColor(depth_frame, cv2.COLOR_RGB2BGR)
                        else:
                            depth_frame = (depth * 255).astype(np.uint8)
                            depth_frame = cv2.cvtColor(depth_frame, cv2.COLOR_GRAY2BGR)
                    
                    out.write(depth_frame)
                    
                except Exception as e:
                    print(f"Error processing frame {frame_count}: {str(e)}")
                    black_frame = np.zeros((height, width, 3), dtype=np.uint8)
                    out.write(black_frame)
                
                progress((frame_count / total_frames), f"Processing progress: {frame_count}/{total_frames} frames")
                
        except Exception as e:
            print(f"Unexpected error during video processing: {str(e)}")
        finally:
            cap.release()
            out.release()
        
        print(f"Video processing completed! Output saved to: {output_path}")
        return output_path
        
    except Exception as e:
        print(f"Video processing failed: {str(e)}")
        return None


print("Loading model...")
model = load_model()
print("Model loaded successfully!")


def predict_depth(input_file, colormap_choice):
    """Main depth prediction function for both images and videos"""
    try:
        if input_file is None:
            return None, gr.update(visible=False)
            
        if is_video_file(input_file):
            output_path = process_video(input_file, colormap_choice)
            if output_path:
                return output_path, gr.update(visible=True, value=output_path)
            else:
                return None, gr.update(visible=False)
        else:
            if isinstance(input_file, str):
                input_image = Image.open(input_file)
            else:
                input_image = input_file
                
            image_tensor, original_size = preprocess_image(input_image)
            
            if torch.cuda.is_available():
                image_tensor = image_tensor.cuda()
            
            with torch.no_grad():
                prediction = model(image_tensor)
                disparity_tensor = prediction['out']
                depth_tensor = normalize_depth(disparity_tensor)
            
            depth = postprocess_depth(depth_tensor, original_size)
            depth_colored = create_colored_depth_map(depth, colormap_choice.lower())
            
            result = Image.fromarray(depth_colored)
            
            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
            result.save(temp_file.name)
            
            return result, gr.update(visible=True, value=temp_file.name)
        
    except Exception as e:
        print(f"Error during inference: {str(e)}")
        return None, gr.update(visible=False)


def capture_and_predict(camera_image, colormap_choice):
    """Capture image from camera and predict depth"""
    return predict_depth(camera_image, colormap_choice)


with gr.Blocks(title="Depth Anything AC - Depth Estimation Demo", theme=gr.themes.Soft(), css="""
    .image-container { 
        display: flex !important; 
        align-items: flex-start !important; 
        justify-content: center !important; 
    }
    .gradio-image { 
        vertical-align: top !important; 
    }
""") as demo:
    gr.Markdown("""
    # 🌊 Depth Anything AC - Depth Estimation Demo
    
    Upload an image or use your camera to generate corresponding depth maps! Different colors in the depth map represent different distances, allowing you to see the three-dimensional structure of the image.
    
    ## How to Use
    1. **Upload Mode**: Click the upload area to select an image or video file
    2. **Camera Mode**: Use your camera to capture a live image
    3. Choose your preferred colormap style
    4. Click the "Generate Depth Map" button
    5. View the results and download
    """)
    
    with gr.Row():
        input_source = gr.Radio(
            choices=["Upload Image", "Upload Video", "Use Camera"],
            value="Upload Image",
            label="Input Source"
        )
        colormap_choice = gr.Dropdown(
            choices=["Spectral", "Inferno", "Gray"],
            value="Spectral",
            label="Colormap Style"
        )
        submit_btn = gr.Button(
            "🎯 Generate Depth Map",
            variant="primary",
            size="lg"
        )
    
    with gr.Row():
        gr.HTML("<h3 style='text-align: center; margin: 10px;'>📷 Input Image</h3>")
        gr.HTML("<h3 style='text-align: center; margin: 10px;'>🌊 Depth Map Result</h3>")
    
    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            # Image input component for preview and examples
            upload_image = gr.Image(
                type="pil",
                height=450,
                visible=True,
                show_label=False,
                container=False,
                label="Upload Image"
            )
            
            # File component for video uploads
            upload_file = gr.File(
                file_types=["video"],
                height=200,
                visible=False,
                show_label=False,
                container=False,
                label="Upload Video"
            )
            
            # Camera component
            camera_image = gr.Image(
                type="pil", 
                sources=["webcam"],
                height=450,
                visible=False,
                show_label=False,
                container=False
            )
            
        with gr.Column(scale=1):
            output_file = gr.File(
                height=450,
                show_label=False,
                container=False,
                visible=False
            )
            
            output_image = gr.Image(
                type="pil",
                height=450,
                show_label=False,
                container=False,
                visible=True
            )
            
            download_btn = gr.DownloadButton(
                label="📥 Download Result",
                variant="secondary",
                size="sm",
                visible=False
            )
    
    def switch_input_source(source):
        if source == "Upload Image":
            return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
        elif source == "Upload Video":
            return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
        else:  # Use Camera
            return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
    
    input_source.change(
        fn=switch_input_source,
        inputs=[input_source],
        outputs=[upload_image, upload_file, camera_image]
    )
    
    def handle_prediction(input_source, upload_img, upload_file_path, camera_img, colormap):
        if input_source == "Upload Image":
            if upload_img is None:
                return None, None, gr.update(visible=False), gr.update(visible=False)
            
            result, download_update = predict_depth(upload_img, colormap)
            return result, None, gr.update(visible=True), download_update
            
        elif input_source == "Upload Video":
            if upload_file_path is None:
                return None, None, gr.update(visible=False), gr.update(visible=False)
            
            result, download_update = predict_depth(upload_file_path, colormap)
            
            if isinstance(result, str) and is_video_file(result):
                return None, result, gr.update(visible=False), download_update
            else:
                return result, None, gr.update(visible=True), download_update
        else:  # Use Camera
            result, download_update = predict_depth(camera_img, colormap)
            return result, None, gr.update(visible=True), download_update
    
    # Separate image and video examples
    image_examples = []
    video_examples = []
    if os.path.exists("toyset"):
        for img_file in ["1.png", "2.png", "good.png"]:
            if os.path.exists(f"toyset/{img_file}"):
                image_examples.append([f"toyset/{img_file}", "Spectral"])
        
        for vid_file in ["fog_2_processed_1s-6s_1.0x.mp4", "snow_processed_1s-6s_1.0x.mp4"]:
            if os.path.exists(f"toyset/{vid_file}"):
                video_examples.append([f"toyset/{vid_file}", "Spectral"])
    
    # Function to handle video example selection and auto-switch mode
    def handle_video_example(video_path, colormap):
        # Auto-switch to video mode and return the necessary updates
        return (
            "Upload Video",  # input_source
            gr.update(visible=False),  # upload_image
            gr.update(visible=True, value=video_path),  # upload_file
            gr.update(visible=False)  # camera_image
        )
    
    # Function to handle image example selection and auto-switch mode
    def handle_image_example(image, colormap):
        # Auto-switch to image mode and process the image
        result = predict_depth(image, colormap)
        output_image = result[0] if result[0] is not None else None
        return (
            "Upload Image",  # input_source
            gr.update(visible=True, value=image),  # upload_image
            gr.update(visible=False),  # upload_file
            gr.update(visible=False),  # camera_image
            output_image  # output_image
        )
    
    if image_examples:
        gr.Examples(
            examples=image_examples,
            inputs=[upload_image, colormap_choice],
            outputs=[input_source, upload_image, upload_file, camera_image, output_image],
            fn=handle_image_example,
            cache_examples=False,
            label="Try these example images"
        )
    
    if video_examples:
        gr.Examples(
            examples=video_examples,
            inputs=[upload_file, colormap_choice],
            outputs=[input_source, upload_image, upload_file, camera_image],
            fn=handle_video_example,
            cache_examples=False,
            label="Try these example videos"
        )
    
    submit_btn.click(
        fn=handle_prediction,
        inputs=[input_source, upload_image, upload_file, camera_image, colormap_choice],
        outputs=[output_image, output_file, output_image, download_btn],
        show_progress=True
    )
    
    gr.Markdown("""
    ## 📝 Colormap Description
    - **Spectral**: Rainbow spectrum, with clear contrast between near and far
    - **Inferno**: Fire spectrum, warm tones
    - **Gray**: Classic grayscale depth representation
    
    ## 📷 Camera Usage Tips
    - Ensure camera access is allowed when prompted
    - Click the camera button to capture the current frame
    - The captured image will be used as input for depth estimation
    
    ## 🎬 Video Processing Tips
    - Supports multiple video formats (MP4, AVI, MOV, etc.)
    - Video processing may take some time, please be patient
    - Processing progress will be displayed in real-time
    - The output video will maintain the same frame rate as the input
    """)


if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )