enhanced_app.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Enhanced Single-View Gradio App for Semantic Segmentation, Depth Estimation, and 3D Point Cloud
Processes one image and shows all outputs: segmentation, depth, and colored point cloud
Now with precomputed examples for demonstration
"""

import sys
import locale
import os
import datetime
from pathlib import Path

# Set UTF-8 encoding
if sys.version_info >= (3, 7):
    sys.stdout.reconfigure(encoding='utf-8')
    sys.stderr.reconfigure(encoding='utf-8')

# Set locale for proper Unicode support
try:
    locale.setlocale(locale.LC_ALL, 'en_US.UTF-8')
except locale.Error:
    try:
        locale.setlocale(locale.LC_ALL, 'C.UTF-8')
    except locale.Error:
        pass  # Use system default

import gradio as gr
import torch
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import io
import base64
from dataclasses import dataclass
from typing import Optional, List, Tuple, Dict, Any
import requests
import cv2
from abc import ABC, abstractmethod
from collections import namedtuple
import plotly.graph_objects as go
import plotly.io as pio
import open3d as o3d
import json
import subprocess

# Import DepthAnythingV2 (assuming it's in the same directory or installed)
try:
    from metric_depth.depth_anything_v2.dpt import DepthAnythingV2
    DEPTH_AVAILABLE = True
except ImportError:
    print("DepthAnythingV2 not available. Using precomputed examples only.")
    DEPTH_AVAILABLE = False
    
    
CUDA_AVAILABLE = torch.cuda.is_available() 

# Set environment variable to disable xFormers
os.environ['XFORMERS_DISABLED'] = '1'
os.environ['XFORMERS_MORE_DETAILS'] = '1'

# Output directory structure (mounted volume)
OUTPUT_DIR = Path("outputs")

def fix_lfs_on_startup():
    """Quick fix for LFS issues on HuggingFace startup."""
    print("Checking for LFS issues...")
    
    try:
        # Try to pull LFS files
        result = subprocess.run(['git', 'lfs', 'pull'], 
                              capture_output=True, text=True, timeout=30)
        if result.returncode == 0:
            print("LFS files pulled successfully")
        else:
            print(f"LFS pull failed: {result.stderr}")
            # Try checkout instead
            subprocess.run(['git', 'lfs', 'checkout'], 
                         capture_output=True, timeout=20)
    except Exception as e:
        print(f"LFS operations failed: {e}")

# =============================================================================
# Model Base Classes and Configurations
# =============================================================================

@dataclass
class ModelConfig:
    """Configuration for segmentation models."""
    model_name: str
    processor_name: str
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
    trust_remote_code: bool = True
    task_type: str = "semantic"

@dataclass
class DepthConfig:
    """Configuration for depth estimation models."""
    encoder: str = "vitl"  # 'vits', 'vitb', 'vitl'
    dataset: str = "vkitti"  # 'hypersim' for indoor, 'vkitti' for outdoor
    max_depth: int = 80  # 20 for indoor, 80 for outdoor
    weights_path: str = "depth_anything_v2_metric_vkitti_vitl.pth"
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
    


class BaseSegmentationModel(ABC):
    """Abstract base class for segmentation models."""
    
    def __init__(self, model_config):
        self.config = model_config
        self.model = None
        self.processor = None
        self.device = torch.device(model_config.device if torch.cuda.is_available() else "cpu")
        
    @abstractmethod
    def load_model(self):
        """Load the model and processor."""
        pass
    
    @abstractmethod
    def preprocess(self, image: Image.Image, **kwargs) -> Dict[str, torch.Tensor]:
        """Preprocess the input image."""
        pass
    
    @abstractmethod
    def predict(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        """Run inference on preprocessed inputs."""
        pass
    
    @abstractmethod
    def postprocess(self, outputs: Dict[str, torch.Tensor], target_size: Tuple[int, int]) -> np.ndarray:
        """Postprocess model outputs to segmentation map."""
        pass
    
    def segment_image(self, image: Image.Image, **kwargs) -> np.ndarray:
        """End-to-end segmentation pipeline."""
        if self.model is None:
            self.load_model()
            
        inputs = self.preprocess(image, **kwargs)
        outputs = self.predict(inputs)
        segmentation_map = self.postprocess(outputs, image.size[::-1])
        
        return segmentation_map

# =============================================================================
# OneFormer Model Implementation
# =============================================================================

class OneFormerModel(BaseSegmentationModel):
    """OneFormer model for universal segmentation."""
    
    def __init__(self, model_config):
        super().__init__(model_config)
        
    def load_model(self):
        """Load OneFormer model and processor."""
        print(f"Loading OneFormer model: {self.config.model_name}")
        
        try:
            from transformers import OneFormerProcessor, OneFormerForUniversalSegmentation
            
            self.processor = OneFormerProcessor.from_pretrained(
                self.config.processor_name,
                trust_remote_code=self.config.trust_remote_code
            )
            
            self.model = OneFormerForUniversalSegmentation.from_pretrained(
                self.config.model_name,
                trust_remote_code=self.config.trust_remote_code
            )
            
            self.model.to(self.device)
            self.model.eval()
            
            print(f"OneFormer model loaded successfully on {self.device}")
            
        except Exception as e:
            print(f"Error loading OneFormer model: {e}")
            raise
        
    def preprocess(self, image: Image.Image, task_inputs: List[str] = None) -> Dict[str, torch.Tensor]:
        """Preprocess image for OneFormer."""
        if task_inputs is None:
            task_inputs = [self.config.task_type]
            
        inputs = self.processor(
            images=image,
            task_inputs=task_inputs,
            return_tensors="pt"
        )
        
        # Move inputs to device
        inputs = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v 
                 for k, v in inputs.items()}
        
        return inputs
    
    def predict(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        """Run inference with OneFormer."""
        with torch.no_grad():
            outputs = self.model(**inputs)
            
        return outputs
    
    def postprocess(self, outputs: Dict[str, torch.Tensor], target_size: Tuple[int, int]) -> np.ndarray:
        """Postprocess OneFormer outputs."""
        predicted_semantic_map = self.processor.post_process_semantic_segmentation(
            outputs, 
            target_sizes=[target_size]
        )[0]
        
        return predicted_semantic_map.cpu().numpy()

# =============================================================================
# DepthAnythingV2 Model Implementation
# =============================================================================

class DepthAnythingV2Model:
    """DepthAnythingV2 model for depth estimation."""
    
    def __init__(self, depth_config: DepthConfig):
        self.config = depth_config
        self.model = None
        self.device = torch.device(depth_config.device if torch.cuda.is_available() else "cpu")
        
    def load_model(self):
        """Load DepthAnythingV2 model."""
        if not DEPTH_AVAILABLE:
            raise ImportError("DepthAnythingV2 is not available")
            
        print(f"Loading DepthAnythingV2 model: {self.config.encoder}")
        
        try:
            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]}
            }
            
            self.model = DepthAnythingV2(**{**model_configs[self.config.encoder], 'max_depth': self.config.max_depth})
            
            # Load weights
            if os.path.exists(self.config.weights_path):
                self.model.load_state_dict(torch.load(self.config.weights_path, map_location='cpu'))
                print(f"Loaded weights from {self.config.weights_path}")
            else:
                print(f"Warning: Weights file {self.config.weights_path} not found")
                
            self.model.to(self.device)
            self.model.eval()
            
            print(f"DepthAnythingV2 model loaded successfully on {self.device}")
            
        except Exception as e:
            print(f"Error loading DepthAnythingV2 model: {e}")
            raise
    
    def estimate_depth(self, image: Image.Image) -> np.ndarray:
        """Estimate depth from image."""
        if self.model is None:
            self.load_model()
            
        # Convert PIL to OpenCV format
        img_array = np.array(image)
        if len(img_array.shape) == 3:
            img_array = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
        
        # Infer depth
        depth_map = self.model.infer_image(img_array)
        
        return depth_map

# =============================================================================
# Cityscapes Label Definitions
# =============================================================================

Label = namedtuple('Label', [
    'name', 'id', 'trainId', 'category', 'categoryId', 
    'hasInstances', 'ignoreInEval', 'color'
])

labels = [
    Label('unlabeled',            0, 255, 'void',            0, False, True,  (0,  0,  0)),
    Label('ego vehicle',          1, 255, 'void',            0, False, True,  (0,  0,  0)),
    Label('rectification border', 2, 255, 'void',            0, False, True,  (0,  0,  0)),
    Label('out of roi',           3, 255, 'void',            0, False, True,  (0,  0,  0)),
    Label('static',               4, 255, 'void',            0, False, True,  (0,  0,  0)),
    Label('dynamic',              5, 255, 'void',            0, False, True,  (111, 74,  0)),
    Label('ground',               6, 255, 'void',            0, False, True,  (81,  0, 81)),
    Label('road',                 7,   0, 'flat',            1, False, False, (128, 64,128)),
    Label('sidewalk',             8,   1, 'flat',            1, False, False, (244, 35,232)),
    Label('parking',              9, 255, 'flat',            1, False, True,  (250,170,160)),
    Label('rail track',          10, 255, 'flat',            1, False, True,  (230,150,140)),
    Label('building',            11,   2, 'construction',    2, False, False, (70, 70, 70)),
    Label('wall',                12,   3, 'construction',    2, False, False, (102,102,156)),
    Label('fence',               13,   4, 'construction',    2, False, False, (190,153,153)),
    Label('guard rail',          14, 255, 'construction',    2, False, True,  (180,165,180)),
    Label('bridge',              15, 255, 'construction',    2, False, True,  (150,100,100)),
    Label('tunnel',              16, 255, 'construction',    2, False, True,  (150,120, 90)),
    Label('pole',                17,   5, 'object',          3, False, False, (153,153,153)),
    Label('polegroup',           18, 255, 'object',          3, False, True,  (153,153,153)),
    Label('traffic light',       19,   6, 'object',          3, False, False, (250,170, 30)),
    Label('traffic sign',        20,   7, 'object',          3, False, False, (220,220,  0)),
    Label('vegetation',          21,   8, 'nature',          4, False, False, (107,142, 35)),
    Label('terrain',             22,   9, 'nature',          4, False, False, (152,251,152)),
    Label('sky',                 23,  10, 'sky',             5, False, False, (70,130,180)),
    Label('person',              24,  11, 'human',           6, True,  False, (220, 20, 60)),
    Label('rider',               25,  12, 'human',           6, True,  False, (255,  0,  0)),
    Label('car',                 26,  13, 'vehicle',         7, True,  False, (0,  0,142)),
    Label('truck',               27,  14, 'vehicle',         7, True,  False, (0,  0, 70)),
    Label('bus',                 28,  15, 'vehicle',         7, True,  False, (0, 60,100)),
    Label('caravan',             29, 255, 'vehicle',         7, True,  True,  (0,  0, 90)),
    Label('trailer',             30, 255, 'vehicle',         7, True,  True,  (0,  0,110)),
    Label('train',               31,  16, 'vehicle',         7, True,  False, (0, 80,100)),
    Label('motorcycle',          32,  17, 'vehicle',         7, True,  False, (0,  0,230)),
    Label('bicycle',             33,  18, 'vehicle',         7, True,  False, (119, 11, 32)),
    Label('license plate',       -1,  -1, 'vehicle',         7, False, True,  (0,  0,142)),
]

# Sky trainId is 10
SKY_TRAIN_ID = 10

# =============================================================================
# Utility Functions
# =============================================================================

def get_color_map(labels):
    """Returns a color map dictionary for the given labels."""
    color_map = {label.trainId: label.color for label in labels if label.trainId != 255}
    return color_map

def apply_color_map(semantic_map, color_map):
    """Applies a color map to a semantic map."""
    height, width = semantic_map.shape
    color_mapped_image = np.zeros((height, width, 3), dtype=np.uint8)
    
    for trainId, color in color_map.items():
        mask = semantic_map == trainId
        color_mapped_image[mask] = color
    
    return color_mapped_image

def create_depth_visualization(depth_map: np.ndarray, colormap: str = 'magma') -> Image.Image:
    """Create a colored depth map visualization with exact dimensions."""
    # Normalize depth map to [0, 1]
    normalized_depth = depth_map / np.max(depth_map)
    
    # Apply colormap
    cmap = plt.get_cmap(colormap)
    colored_depth = cmap(normalized_depth)
    
    # Convert to 8-bit RGB (remove alpha channel)
    colored_depth_8bit = (colored_depth[:, :, :3] * 255).astype(np.uint8)
    
    return Image.fromarray(colored_depth_8bit)

def depth_to_point_cloud_with_segmentation(depth_map: np.ndarray, rgb_image: Image.Image, 
                                          semantic_map: np.ndarray,
                                          fx: float = 525.0, fy: float = 525.0, 
                                          cx: float = None, cy: float = None) -> o3d.geometry.PointCloud:
    """Convert depth map and RGB image to 3D point cloud with segmentation colors, excluding sky."""
    height, width = depth_map.shape
    
    if cx is None:
        cx = width / 2.0
    if cy is None:
        cy = height / 2.0
    
    # Create coordinate matrices
    u, v = np.meshgrid(np.arange(width), np.arange(height))
    
    # Convert to 3D coordinates
    z = depth_map
    x = (u - cx) * z / fx
    y = (v - cy) * z / fy
    
    # Stack coordinates
    points = np.stack([x, y, z], axis=-1).reshape(-1, 3)
    
    # Create mask to exclude sky points and invalid depths
    flat_semantic = semantic_map.flatten()
    flat_depth = z.flatten()
    
    # Filter out invalid points and sky points
    valid_mask = (flat_depth > 0) & (flat_depth < 1000) & (flat_semantic != SKY_TRAIN_ID)
    points = points[valid_mask]
    
    # Get segmentation colors for each point
    color_map = get_color_map(labels)
    seg_colors = np.zeros((len(flat_semantic), 3))
    
    for trainId, color in color_map.items():
        mask = flat_semantic == trainId
        seg_colors[mask] = color
    
    # Filter colors to match valid points
    colors = seg_colors[valid_mask] / 255.0  # Normalize to [0, 1]
    
    # Create Open3D point cloud
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    pcd.colors = o3d.utility.Vector3dVector(colors)
    
    return pcd

def create_plotly_pointcloud(pcd: o3d.geometry.PointCloud, downsample_factor: float = 0.1) -> go.Figure:
    """Create interactive Plotly 3D point cloud visualization."""
    # Downsample for performance
    if downsample_factor < 1.0:
        num_points = len(pcd.points)
        indices = np.random.choice(num_points, int(num_points * downsample_factor), replace=False)
        points = np.asarray(pcd.points)[indices]
        colors = np.asarray(pcd.colors)[indices]
    else:
        points = np.asarray(pcd.points)
        colors = np.asarray(pcd.colors)
    
    # Create 3D scatter plot
    fig = go.Figure(data=[go.Scatter3d(
        x=points[:, 0],
        y=points[:, 1], 
        z=points[:, 2],
        mode='markers',
        marker=dict(
            size=1,
            color=colors,
            opacity=0.8
        ),
        text=[f'Point {i}' for i in range(len(points))],
        hovertemplate='X: %{x:.2f}<br>Y: %{y:.2f}<br>Z: %{z:.2f}<extra></extra>'
    )])
    
    # Update layout for centered display
    fig.update_layout(
        scene=dict(
            xaxis_title='X (Horizontal)',
            yaxis_title='Y (Vertical)', 
            zaxis_title='Z (Depth)',
            aspectmode='data'
        ),
        title={
            'text': 'Interactive 3D Point Cloud (Colored by Segmentation, Sky Excluded)',
            'x': 0.5,
            'xanchor': 'center'
        },
        width=None,  # Let it auto-size to container
        height=600,
        margin=dict(l=0, r=0, t=40, b=0),  # Minimal margins
        autosize=True  # Enable auto-sizing to container
    )
    
    # Set camera for bird's eye view that clearly shows 3D structure
    fig.update_layout(scene_camera=dict(
        up=dict(x=0, y=0, z=1),        # Z-axis points up
        center=dict(x=0, y=0, z=0),    # Center at origin
        eye=dict(x=0.5, y=-2.5, z=1.5) # View from above-back position
    ))
    
    return fig

def create_overlay_plot(rgb_image: Image.Image, semantic_map: np.ndarray, alpha: float = 0.5):
    """Create segmentation overlay plot without title and borders."""
    rgb_array = np.array(rgb_image)
    color_map = get_color_map(labels)
    colored_semantic_map = apply_color_map(semantic_map, color_map)
    
    # Create figure with exact image dimensions
    height, width = rgb_array.shape[:2]
    dpi = 100
    fig, ax = plt.subplots(1, 1, figsize=(width/dpi, height/dpi), dpi=dpi)
    
    # Remove all margins and padding
    fig.subplots_adjust(left=0, right=1, top=1, bottom=0)
    
    ax.imshow(rgb_array)
    ax.imshow(colored_semantic_map, alpha=alpha)
    ax.axis('off')
    
    buf = io.BytesIO()
    plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0, dpi=dpi)
    buf.seek(0)
    plt.close(fig)
    
    return Image.open(buf)

class PrecomputedExamplesManager:
    """Manages precomputed examples from output folder structure."""
    
    def __init__(self, output_dir: Path):
        self.output_dir = output_dir
        self.rgb_dir = output_dir / "rgb"
        self.segmentation_dir = output_dir / "segmentation"
        self.depth_dir = output_dir / "depth" 
        self.pointclouds_dir = output_dir / "pointclouds"
        self.examples = self._load_examples()
    
    def _load_examples(self) -> Dict[str, Dict]:
        """Load all available precomputed examples from output structure."""
        examples = {}
        
        if not self.output_dir.exists():
            print(f"Output directory {self.output_dir} not found.")
            return {}
        
        # Find all timestamps by looking at RGB files (the inputs)
        if not self.rgb_dir.exists():
            print(f"RGB directory {self.rgb_dir} not found.")
            return {}
        
        # Get all RGB files and extract timestamps
        timestamps = set()
        for rgb_file in self.rgb_dir.glob("rgb_*.png"):
            # Extract timestamp from filename like "rgb_20241215_143022.png"
            filename = rgb_file.stem
            if filename.startswith("rgb_"):
                timestamp = filename.replace("rgb_", "")
                timestamps.add(timestamp)
        
        print(f"Found {len(timestamps)} RGB input images")
        
        # For each timestamp, try to load the complete example
        for timestamp in sorted(timestamps, reverse=True):  # Most recent first
            example_data = self._load_single_example(timestamp)
            if example_data:
                examples[timestamp] = example_data
        
        print(f"Loaded {len(examples)} precomputed examples from output directory")
        return examples
    
    def _load_single_example(self, timestamp: str) -> Optional[Dict]:
        """Load a single precomputed example by timestamp."""
        try:
            # Input file (required)
            rgb_path = self.rgb_dir / f"rgb_{timestamp}.png"
            
            # Output files (some may be optional)
            seg_path = self.segmentation_dir / f"segmentation_{timestamp}.png"
            depth_path = self.depth_dir / f"depth_{timestamp}.png"
            ply_path = self.pointclouds_dir / f"pointcloud_{timestamp}.ply"
            html_path = self.pointclouds_dir / f"pointcloud_{timestamp}.html"
            
            # Check if RGB input exists (required)
            if not rgb_path.exists():
                print(f"RGB input file missing for timestamp {timestamp}: {rgb_path}")
                return None
            
            # Check if at least segmentation output exists
            if not seg_path.exists():
                print(f"Segmentation output missing for timestamp {timestamp}: {seg_path}")
                return None
            
            # Create a display name from timestamp
            try:
                # Parse timestamp like "20241215_143022"
                if len(timestamp) >= 13 and "_" in timestamp:
                    date_part = timestamp[:8]
                    time_part = timestamp[9:15]
                    # Format as "Dec 15, 2024 14:30"
                    year = date_part[:4]
                    month = date_part[4:6]
                    day = date_part[6:8]
                    hour = time_part[:2]
                    minute = time_part[2:4]
                    
                    month_names = ["", "Jan", "Feb", "Mar", "Apr", "May", "Jun",
                                 "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"]
                    month_name = month_names[int(month)] if 1 <= int(month) <= 12 else month
                    
                    display_name = f"{month_name} {int(day)}, {year} {hour}:{minute}"
                else:
                    display_name = timestamp
            except:
                display_name = timestamp
            
            return {
                'name': display_name,
                'timestamp': timestamp,
                'rgb_path': rgb_path,  # Input image
                'segmentation_path': seg_path,  # Output
                'depth_path': depth_path if depth_path.exists() else None,  # Output (optional)
                'pointcloud_ply_path': ply_path if ply_path.exists() else None,  # Output (optional)
                'pointcloud_html_path': html_path if html_path.exists() else None,  # Output (optional)
                'preview_image': self._create_preview_image(rgb_path, timestamp)
            }
            
        except Exception as e:
            print(f"Error loading example {timestamp}: {e}")
            return None
    
    def _create_preview_image(self, rgb_path: Path, timestamp: str) -> Image.Image:
        """Create a preview thumbnail from RGB input image."""
        try:
            image = Image.open(rgb_path)
            image.thumbnail((600, 450), Image.Resampling.LANCZOS)
            return image
                
        except Exception as e:
            print(f"Error creating preview for {timestamp}: {e}")
            return Image.new('RGB', (200, 150), color=(128, 128, 128))
    
    def get_example_names(self) -> List[str]:
        """Get list of available example names."""
        return [data['name'] for data in self.examples.values()]
    
    def get_example_previews(self) -> List[Tuple[Image.Image, str]]:
        """Get preview images for all examples."""
        previews = []
        for timestamp, data in self.examples.items():
            previews.append((data['preview_image'], data['name']))
        return previews
    
    def get_timestamp_by_name(self, name: str) -> Optional[str]:
        """Get timestamp by display name."""
        for timestamp, data in self.examples.items():
            if data['name'] == name:
                return timestamp
        return None
    
    def load_example_results(self, example_name: str) -> Tuple[Optional[Image.Image], Optional[Image.Image], Optional[go.Figure], str]:
        """Load precomputed results for an example."""
        if not example_name:
            return None, None, None, "Please select an example."
        
        # Find the timestamp for this example name
        timestamp = self.get_timestamp_by_name(example_name)
        if not timestamp or timestamp not in self.examples:
            return None, None, None, f"Example '{example_name}' not found."
        
        example_data = self.examples[timestamp]
        
        try:
            # Load output images
            segmentation_image = Image.open(example_data['segmentation_path'])
            
            depth_image = None
            if example_data['depth_path'] and example_data['depth_path'].exists():
                depth_image = Image.open(example_data['depth_path'])
            
            # Load point cloud if available
            point_cloud_fig = None
            if example_data['pointcloud_ply_path'] and example_data['pointcloud_ply_path'].exists():
                try:
                    pcd = o3d.io.read_point_cloud(str(example_data['pointcloud_ply_path']))
                    if len(pcd.points) > 0:
                        point_cloud_fig = create_plotly_pointcloud(pcd, downsample_factor=1)
                    else:
                        print(f"Point cloud file {example_data['pointcloud_ply_path']} is empty")
                except Exception as e:
                    print(f"Error loading point cloud: {e}")
            
            return segmentation_image, depth_image, point_cloud_fig, ""
            
        except Exception as e:
            return None, None, None, f"Error loading example results: {str(e)}"

# =============================================================================
# Main Application Class
# =============================================================================

class EnhancedSingleViewApp:
    def __init__(self):
        # Model configurations
        self.oneformer_config = ModelConfig(
            model_name="shi-labs/oneformer_cityscapes_swin_large",
            processor_name="shi-labs/oneformer_cityscapes_swin_large",
            task_type="semantic"
        )
        
        self.depth_config = DepthConfig(
            encoder="vitl",
            dataset="vkitti",
            max_depth=80,
            weights_path="depth_anything_v2_metric_vkitti_vitl.pth"
        )
        
        # Models
        self.oneformer_model = None
        self.depth_model = None
        self.segmentation_loaded = False
        self.depth_loaded = False
        
        # Precomputed examples manager
        self.examples_manager = PrecomputedExamplesManager(OUTPUT_DIR)
        
        # Online sample images (fallback)
        self.sample_images = {
            "Street Scene 1": "https://images.unsplash.com/photo-1449824913935-59a10b8d2000?w=800",
            "Street Scene 2": "https://images.unsplash.com/photo-1502920917128-1aa500764cbd?w=800", 
            "Urban Road": "https://images.unsplash.com/photo-1516738901171-8eb4fc13bd20?w=800",
            "City View": "https://images.unsplash.com/photo-1477959858617-67f85cf4f1df?w=800",
            "Highway": "https://images.unsplash.com/photo-1544620347-c4fd4a3d5957?w=800",
        }
    
    def download_sample_image(self, image_url: str) -> Image.Image:
        """Download a sample image from URL."""
        try:
            response = requests.get(image_url, timeout=10)
            response.raise_for_status()
            return Image.open(io.BytesIO(response.content)).convert('RGB')
        except Exception as e:
            print(f"Error downloading image: {e}")
            return Image.new('RGB', (800, 600), color=(128, 128, 128))
    
    def create_overlay_plot(self, rgb_image: Image.Image, semantic_map: np.ndarray, alpha: float = 0.5):
        """Create segmentation overlay plot without title and borders."""
        rgb_array = np.array(rgb_image)
        color_map = get_color_map(labels)
        colored_semantic_map = apply_color_map(semantic_map, color_map)
        
        # Create figure with exact image dimensions
        height, width = rgb_array.shape[:2]
        dpi = 100
        fig, ax = plt.subplots(1, 1, figsize=(width/dpi, height/dpi), dpi=dpi)
        
        # Remove all margins and padding
        fig.subplots_adjust(left=0, right=1, top=1, bottom=0)
        
        ax.imshow(rgb_array)
        ax.imshow(colored_semantic_map, alpha=alpha)
        ax.axis('off')
        
        buf = io.BytesIO()
        plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0, dpi=dpi)
        buf.seek(0)
        plt.close(fig)
        
        return Image.open(buf)
    
    def process_complete_pipeline(self, image: Image.Image):
        """Process image through complete pipeline: segmentation + depth + point cloud."""
        if image is None:
            return None, None, None, "Please upload an image."
        
        # Default values
        overlay_alpha = 0.5
        depth_colormap = "magma"
        downsample_factor = 0.1
        
        try:
            # Auto-load models if not loaded
            if not self.segmentation_loaded:
                if self.oneformer_model is None:
                    self.oneformer_model = OneFormerModel(self.oneformer_config)
                    self.oneformer_model.load_model()
                self.segmentation_loaded = True
            
            if not self.depth_loaded and DEPTH_AVAILABLE:
                if self.depth_model is None:
                    self.depth_model = DepthAnythingV2Model(self.depth_config)
                    self.depth_model.load_model()
                self.depth_loaded = True
            
            # Resize if too large
            original_size = image.size
            if max(image.size) > 1024:
                image.thumbnail((1024, 1024), Image.Resampling.LANCZOS)
            
            # Step 1: Semantic Segmentation
            task_inputs = ["semantic"]
            semantic_map = self.oneformer_model.segment_image(image, task_inputs=task_inputs)
            segmentation_overlay = self.create_overlay_plot(image, semantic_map, overlay_alpha)
            
            # Step 2: Depth Estimation (if available)
            depth_vis = None
            point_cloud_fig = None
            pcd = None
            
            if DEPTH_AVAILABLE and self.depth_loaded:
                depth_map = self.depth_model.estimate_depth(image)
                depth_vis = create_depth_visualization(depth_map, depth_colormap)
                
                # Step 3: Point Cloud with Segmentation Colors
                pcd = depth_to_point_cloud_with_segmentation(depth_map, image, semantic_map)
                point_cloud_fig = create_plotly_pointcloud(pcd, downsample_factor)
            
            # Generate comprehensive info
            unique_classes = np.unique(semantic_map)
            class_info = []
            total_pixels = semantic_map.size
            
            for class_id in unique_classes:
                if class_id < len(labels) and class_id != 255:
                    label = labels[class_id]
                    pixel_count = np.sum(semantic_map == class_id)
                    percentage = (pixel_count / total_pixels) * 100
                    if percentage > 0.1:
                        class_info.append(f"- {label.name}: {percentage:.1f}%")
            
            # Point cloud statistics
            if point_cloud_fig is not None:
                num_points = len(pcd.points)
                downsampled_points = int(num_points * downsample_factor)
                point_cloud_info = f"""
3D Point Cloud:
- Total points: {num_points:,}
- Displayed points: {downsampled_points:,} ({downsample_factor*100:.0f}%)
- Sky points excluded
- Colors match segmentation classes"""
            else:
                point_cloud_info = "Point cloud not available (DepthAnythingV2 required)"
            
            # Depth statistics
            if depth_vis is not None and DEPTH_AVAILABLE:
                depth_stats = {
                    'min': np.min(depth_map),
                    'max': np.max(depth_map),
                    'mean': np.mean(depth_map),
                    'std': np.std(depth_map)
                }
                depth_info = f"""
Depth Estimation:
- Min depth: {depth_stats['min']:.2f}m
- Max depth: {depth_stats['max']:.2f}m  
- Mean depth: {depth_stats['mean']:.2f}m
- Std deviation: {depth_stats['std']:.2f}m
- Colormap: {depth_colormap}"""
            else:
                depth_info = "Depth estimation not available"
            
            info_text = f"""Complete vision pipeline processed successfully!

Models Used:
- OneFormer (Semantic Segmentation)
{f"- DepthAnythingV2 ({self.depth_config.encoder.upper()})" if DEPTH_AVAILABLE else "- DepthAnythingV2 (Not Available)"}

Image Processing:
- Original size: {original_size[0]}x{original_size[1]}
- Processed size: {image.size[0]}x{image.size[1]}
- Overlay transparency: {overlay_alpha:.1f}

Detected Classes:
{chr(10).join(class_info)}
{depth_info}
{point_cloud_info}

The point cloud shows 3D structure with each point colored according to its segmentation class. Sky points are excluded for better visualization."""
            
            return segmentation_overlay, depth_vis, point_cloud_fig, info_text
            
        except Exception as e:
            return None, None, None, f"Error processing pipeline: {str(e)}"

# Initialize the app
app = EnhancedSingleViewApp()

def process_uploaded_image(image):
        try:
            return app.process_complete_pipeline(image)
        except:
            return None, None, None

def process_sample_image(sample_choice):
    """Process sample image through complete pipeline.""" 
    if sample_choice and sample_choice in app.sample_images:
        image_url = app.sample_images[sample_choice]
        image = app.download_sample_image(image_url)
        return app.process_complete_pipeline(image)
    return None, None, None, "Please select a sample image."

def load_precomputed_example(evt: gr.SelectData):
    """Load precomputed example results from gallery selection."""
    if evt.index is not None:
        example_names = app.examples_manager.get_example_names()
        if evt.index < len(example_names):
            example_name = example_names[evt.index]
            seg_image, depth_image, pc_fig, info_text = app.examples_manager.load_example_results(example_name)
            return seg_image, depth_image, pc_fig
    return None, None, None

def get_example_previews():
    """Get preview images for the gallery."""
    previews = app.examples_manager.get_example_previews()
    if not previews:
        return []
    return previews

# =============================================================================
# Create Gradio Interface
# =============================================================================

def create_gradio_interface():
    """Create and return the enhanced single-view Gradio interface."""
    
    with gr.Blocks(
        title="Enhanced Computer Vision Pipeline",
        theme=gr.themes.Default()
    ) as demo:
        
        gr.Markdown("""
        # Street Scene 3D Reconstruction
        
        Upload an image or select an example to see:
        - **Semantic Segmentation** - Identify roads, buildings, vehicles, people, and other scene elements
        - **Depth Estimation** - Generate metric depth maps showing distance to objects  
        - **3D Point Cloud** - Interactive 3D reconstruction with semantic colors)
        """)
        
        with gr.Row():
            # Left Column: Controls and Input
            with gr.Column(scale=1):
                if CUDA_AVAILABLE:
                    gr.Markdown("### Upload Image")
                    
                    uploaded_image = gr.Image(
                        type="pil",
                        label="Upload Image"
                    )
                    upload_btn = gr.Button("Process Image", variant="primary", size="lg")
                else:
                    uploaded_image = gr.Image(visible=False)  # Hidden placeholder
                    upload_btn = gr.Button(visible=False)     # Hidden placeholder
                    
                    gr.Markdown("### CPU Mode")
                    gr.Markdown("⚠️ **Upload disabled**: DepthAnythingV2 requires CUDA. Using precomputed examples only.")
                
                gr.Markdown("### Examples")
                gr.Markdown("Click on an image to load the example:")
                
                # Example gallery (always visible)
                example_gallery = gr.Gallery(
                    value=get_example_previews(),
                    label="Example Images",
                    show_label=False,
                    elem_id="example_gallery",
                    columns=2,
                    rows=3,
                    height="auto",
                    object_fit="cover"
                )
            
            # Right Column: Results
            with gr.Column(scale=2):
                gr.Markdown("### Results")
                
                # Segmentation and Depth side by side
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("#### Semantic Segmentation")
                        segmentation_output = gr.Image(label="Segmentation Overlay")
                    
                    with gr.Column():
                        gr.Markdown("#### Depth Estimation")
                        depth_output = gr.Image(label="Depth Map")
                
                # Point Cloud below
                gr.Markdown("#### 3D Point Cloud")
                pointcloud_output = gr.Plot(label="Interactive 3D Point Cloud (Colored by Segmentation)")
        
        if CUDA_AVAILABLE:
          upload_btn.click(
              fn=process_uploaded_image,
              inputs=[uploaded_image],
              outputs=[segmentation_output, depth_output, pointcloud_output]
          )
        
        # Gallery selection loads example directly
        example_gallery.select(
            fn=load_precomputed_example,
            outputs=[segmentation_output, depth_output, pointcloud_output]
        )
    
    return demo

# =============================================================================
# Main Execution
# =============================================================================

if __name__ == "__main__":
    fix_lfs_on_startup()
    # Create and launch the interface
    demo = create_gradio_interface()
    
    print("Starting Enhanced Single-View Computer Vision App...")
    print("Complete Pipeline: Segmentation + Depth + 3D Point Cloud")
    print("Device:", "CUDA" if torch.cuda.is_available() else "CPU")
    print("Depth Available:", "YES" if DEPTH_AVAILABLE else "NO")
    print("Point Cloud Colors: Segmentation-based (Sky Excluded)")
    print(f"Output Directory: {OUTPUT_DIR.absolute()}")
    print(f"Available Examples: {len(app.examples_manager.examples)}")
    
    # Launch the app
    demo.launch(
        share=True,           # Creates a public link
        debug=True,           # Enable debugging
        server_name="0.0.0.0", # Allow external connections
        server_port=7860,     # Default port
        show_error=True,      # Show errors in the interface
        quiet=False           # Show startup logs
    )