meladeayol/Road_Segmentation_with_Depth_Estimation
0
1#!/usr/bin/env python32# -*- coding: utf-8 -*-3"""4Enhanced Single-View Gradio App for Semantic Segmentation, Depth Estimation, and 3D Point Cloud5Processes one image and shows all outputs: segmentation, depth, and colored point cloud6Now with precomputed examples for demonstration7"""8 9import sys10import locale11import os12import datetime13from pathlib import Path14 15# Set UTF-8 encoding16if sys.version_info >= (3, 7):17 sys.stdout.reconfigure(encoding='utf-8')18 sys.stderr.reconfigure(encoding='utf-8')19 20# Set locale for proper Unicode support21try:22 locale.setlocale(locale.LC_ALL, 'en_US.UTF-8')23except locale.Error:24 try:25 locale.setlocale(locale.LC_ALL, 'C.UTF-8')26 except locale.Error:27 pass # Use system default28 29import gradio as gr30import torch31import numpy as np32import matplotlib.pyplot as plt33from PIL import Image34import io35import base6436from dataclasses import dataclass37from typing import Optional, List, Tuple, Dict, Any38import requests39import cv240from abc import ABC, abstractmethod41from collections import namedtuple42import plotly.graph_objects as go43import plotly.io as pio44import open3d as o3d45import json46import subprocess47 48# Import DepthAnythingV2 (assuming it's in the same directory or installed)49try:50 from metric_depth.depth_anything_v2.dpt import DepthAnythingV251 DEPTH_AVAILABLE = True52except ImportError:53 print("DepthAnythingV2 not available. Using precomputed examples only.")54 DEPTH_AVAILABLE = False55 56 57CUDA_AVAILABLE = torch.cuda.is_available() 58 59# Set environment variable to disable xFormers60os.environ['XFORMERS_DISABLED'] = '1'61os.environ['XFORMERS_MORE_DETAILS'] = '1'62 63# Output directory structure (mounted volume)64OUTPUT_DIR = Path("outputs")65 66def fix_lfs_on_startup():67 """Quick fix for LFS issues on HuggingFace startup."""68 print("Checking for LFS issues...")69 70 try:71 # Try to pull LFS files72 result = subprocess.run(['git', 'lfs', 'pull'], 73 capture_output=True, text=True, timeout=30)74 if result.returncode == 0:75 print("LFS files pulled successfully")76 else:77 print(f"LFS pull failed: {result.stderr}")78 # Try checkout instead79 subprocess.run(['git', 'lfs', 'checkout'], 80 capture_output=True, timeout=20)81 except Exception as e:82 print(f"LFS operations failed: {e}")83 84# =============================================================================85# Model Base Classes and Configurations86# =============================================================================87 88@dataclass89class ModelConfig:90 """Configuration for segmentation models."""91 model_name: str92 processor_name: str93 device: str = "cuda" if torch.cuda.is_available() else "cpu"94 trust_remote_code: bool = True95 task_type: str = "semantic"96 97@dataclass98class DepthConfig:99 """Configuration for depth estimation models."""100 encoder: str = "vitl" # 'vits', 'vitb', 'vitl'101 dataset: str = "vkitti" # 'hypersim' for indoor, 'vkitti' for outdoor102 max_depth: int = 80 # 20 for indoor, 80 for outdoor103 weights_path: str = "depth_anything_v2_metric_vkitti_vitl.pth"104 device: str = "cuda" if torch.cuda.is_available() else "cpu"105 106 107 108class BaseSegmentationModel(ABC):109 """Abstract base class for segmentation models."""110 111 def __init__(self, model_config):112 self.config = model_config113 self.model = None114 self.processor = None115 self.device = torch.device(model_config.device if torch.cuda.is_available() else "cpu")116 117 @abstractmethod118 def load_model(self):119 """Load the model and processor."""120 pass121 122 @abstractmethod123 def preprocess(self, image: Image.Image, **kwargs) -> Dict[str, torch.Tensor]:124 """Preprocess the input image."""125 pass126 127 @abstractmethod128 def predict(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:129 """Run inference on preprocessed inputs."""130 pass131 132 @abstractmethod133 def postprocess(self, outputs: Dict[str, torch.Tensor], target_size: Tuple[int, int]) -> np.ndarray:134 """Postprocess model outputs to segmentation map."""135 pass136 137 def segment_image(self, image: Image.Image, **kwargs) -> np.ndarray:138 """End-to-end segmentation pipeline."""139 if self.model is None:140 self.load_model()141 142 inputs = self.preprocess(image, **kwargs)143 outputs = self.predict(inputs)144 segmentation_map = self.postprocess(outputs, image.size[::-1])145 146 return segmentation_map147 148# =============================================================================149# OneFormer Model Implementation150# =============================================================================151 152class OneFormerModel(BaseSegmentationModel):153 """OneFormer model for universal segmentation."""154 155 def __init__(self, model_config):156 super().__init__(model_config)157 158 def load_model(self):159 """Load OneFormer model and processor."""160 print(f"Loading OneFormer model: {self.config.model_name}")161 162 try:163 from transformers import OneFormerProcessor, OneFormerForUniversalSegmentation164 165 self.processor = OneFormerProcessor.from_pretrained(166 self.config.processor_name,167 trust_remote_code=self.config.trust_remote_code168 )169 170 self.model = OneFormerForUniversalSegmentation.from_pretrained(171 self.config.model_name,172 trust_remote_code=self.config.trust_remote_code173 )174 175 self.model.to(self.device)176 self.model.eval()177 178 print(f"OneFormer model loaded successfully on {self.device}")179 180 except Exception as e:181 print(f"Error loading OneFormer model: {e}")182 raise183 184 def preprocess(self, image: Image.Image, task_inputs: List[str] = None) -> Dict[str, torch.Tensor]:185 """Preprocess image for OneFormer."""186 if task_inputs is None:187 task_inputs = [self.config.task_type]188 189 inputs = self.processor(190 images=image,191 task_inputs=task_inputs,192 return_tensors="pt"193 )194 195 # Move inputs to device196 inputs = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v 197 for k, v in inputs.items()}198 199 return inputs200 201 def predict(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:202 """Run inference with OneFormer."""203 with torch.no_grad():204 outputs = self.model(**inputs)205 206 return outputs207 208 def postprocess(self, outputs: Dict[str, torch.Tensor], target_size: Tuple[int, int]) -> np.ndarray:209 """Postprocess OneFormer outputs."""210 predicted_semantic_map = self.processor.post_process_semantic_segmentation(211 outputs, 212 target_sizes=[target_size]213 )[0]214 215 return predicted_semantic_map.cpu().numpy()216 217# =============================================================================218# DepthAnythingV2 Model Implementation219# =============================================================================220 221class DepthAnythingV2Model:222 """DepthAnythingV2 model for depth estimation."""223 224 def __init__(self, depth_config: DepthConfig):225 self.config = depth_config226 self.model = None227 self.device = torch.device(depth_config.device if torch.cuda.is_available() else "cpu")228 229 def load_model(self):230 """Load DepthAnythingV2 model."""231 if not DEPTH_AVAILABLE:232 raise ImportError("DepthAnythingV2 is not available")233 234 print(f"Loading DepthAnythingV2 model: {self.config.encoder}")235 236 try:237 model_configs = {238 'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},239 'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},240 'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]}241 }242 243 self.model = DepthAnythingV2(**{**model_configs[self.config.encoder], 'max_depth': self.config.max_depth})244 245 # Load weights246 if os.path.exists(self.config.weights_path):247 self.model.load_state_dict(torch.load(self.config.weights_path, map_location='cpu'))248 print(f"Loaded weights from {self.config.weights_path}")249 else:250 print(f"Warning: Weights file {self.config.weights_path} not found")251 252 self.model.to(self.device)253 self.model.eval()254 255 print(f"DepthAnythingV2 model loaded successfully on {self.device}")256 257 except Exception as e:258 print(f"Error loading DepthAnythingV2 model: {e}")259 raise260 261 def estimate_depth(self, image: Image.Image) -> np.ndarray:262 """Estimate depth from image."""263 if self.model is None:264 self.load_model()265 266 # Convert PIL to OpenCV format267 img_array = np.array(image)268 if len(img_array.shape) == 3:269 img_array = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)270 271 # Infer depth272 depth_map = self.model.infer_image(img_array)273 274 return depth_map275 276# =============================================================================277# Cityscapes Label Definitions278# =============================================================================279 280Label = namedtuple('Label', [281 'name', 'id', 'trainId', 'category', 'categoryId', 282 'hasInstances', 'ignoreInEval', 'color'283])284 285labels = [286 Label('unlabeled', 0, 255, 'void', 0, False, True, (0, 0, 0)),287 Label('ego vehicle', 1, 255, 'void', 0, False, True, (0, 0, 0)),288 Label('rectification border', 2, 255, 'void', 0, False, True, (0, 0, 0)),289 Label('out of roi', 3, 255, 'void', 0, False, True, (0, 0, 0)),290 Label('static', 4, 255, 'void', 0, False, True, (0, 0, 0)),291 Label('dynamic', 5, 255, 'void', 0, False, True, (111, 74, 0)),292 Label('ground', 6, 255, 'void', 0, False, True, (81, 0, 81)),293 Label('road', 7, 0, 'flat', 1, False, False, (128, 64,128)),294 Label('sidewalk', 8, 1, 'flat', 1, False, False, (244, 35,232)),295 Label('parking', 9, 255, 'flat', 1, False, True, (250,170,160)),296 Label('rail track', 10, 255, 'flat', 1, False, True, (230,150,140)),297 Label('building', 11, 2, 'construction', 2, False, False, (70, 70, 70)),298 Label('wall', 12, 3, 'construction', 2, False, False, (102,102,156)),299 Label('fence', 13, 4, 'construction', 2, False, False, (190,153,153)),300 Label('guard rail', 14, 255, 'construction', 2, False, True, (180,165,180)),301 Label('bridge', 15, 255, 'construction', 2, False, True, (150,100,100)),302 Label('tunnel', 16, 255, 'construction', 2, False, True, (150,120, 90)),303 Label('pole', 17, 5, 'object', 3, False, False, (153,153,153)),304 Label('polegroup', 18, 255, 'object', 3, False, True, (153,153,153)),305 Label('traffic light', 19, 6, 'object', 3, False, False, (250,170, 30)),306 Label('traffic sign', 20, 7, 'object', 3, False, False, (220,220, 0)),307 Label('vegetation', 21, 8, 'nature', 4, False, False, (107,142, 35)),308 Label('terrain', 22, 9, 'nature', 4, False, False, (152,251,152)),309 Label('sky', 23, 10, 'sky', 5, False, False, (70,130,180)),310 Label('person', 24, 11, 'human', 6, True, False, (220, 20, 60)),311 Label('rider', 25, 12, 'human', 6, True, False, (255, 0, 0)),312 Label('car', 26, 13, 'vehicle', 7, True, False, (0, 0,142)),313 Label('truck', 27, 14, 'vehicle', 7, True, False, (0, 0, 70)),314 Label('bus', 28, 15, 'vehicle', 7, True, False, (0, 60,100)),315 Label('caravan', 29, 255, 'vehicle', 7, True, True, (0, 0, 90)),316 Label('trailer', 30, 255, 'vehicle', 7, True, True, (0, 0,110)),317 Label('train', 31, 16, 'vehicle', 7, True, False, (0, 80,100)),318 Label('motorcycle', 32, 17, 'vehicle', 7, True, False, (0, 0,230)),319 Label('bicycle', 33, 18, 'vehicle', 7, True, False, (119, 11, 32)),320 Label('license plate', -1, -1, 'vehicle', 7, False, True, (0, 0,142)),321]322 323# Sky trainId is 10324SKY_TRAIN_ID = 10325 326# =============================================================================327# Utility Functions328# =============================================================================329 330def get_color_map(labels):331 """Returns a color map dictionary for the given labels."""332 color_map = {label.trainId: label.color for label in labels if label.trainId != 255}333 return color_map334 335def apply_color_map(semantic_map, color_map):336 """Applies a color map to a semantic map."""337 height, width = semantic_map.shape338 color_mapped_image = np.zeros((height, width, 3), dtype=np.uint8)339 340 for trainId, color in color_map.items():341 mask = semantic_map == trainId342 color_mapped_image[mask] = color343 344 return color_mapped_image345 346def create_depth_visualization(depth_map: np.ndarray, colormap: str = 'magma') -> Image.Image:347 """Create a colored depth map visualization with exact dimensions."""348 # Normalize depth map to [0, 1]349 normalized_depth = depth_map / np.max(depth_map)350 351 # Apply colormap352 cmap = plt.get_cmap(colormap)353 colored_depth = cmap(normalized_depth)354 355 # Convert to 8-bit RGB (remove alpha channel)356 colored_depth_8bit = (colored_depth[:, :, :3] * 255).astype(np.uint8)357 358 return Image.fromarray(colored_depth_8bit)359 360def depth_to_point_cloud_with_segmentation(depth_map: np.ndarray, rgb_image: Image.Image, 361 semantic_map: np.ndarray,362 fx: float = 525.0, fy: float = 525.0, 363 cx: float = None, cy: float = None) -> o3d.geometry.PointCloud:364 """Convert depth map and RGB image to 3D point cloud with segmentation colors, excluding sky."""365 height, width = depth_map.shape366 367 if cx is None:368 cx = width / 2.0369 if cy is None:370 cy = height / 2.0371 372 # Create coordinate matrices373 u, v = np.meshgrid(np.arange(width), np.arange(height))374 375 # Convert to 3D coordinates376 z = depth_map377 x = (u - cx) * z / fx378 y = (v - cy) * z / fy379 380 # Stack coordinates381 points = np.stack([x, y, z], axis=-1).reshape(-1, 3)382 383 # Create mask to exclude sky points and invalid depths384 flat_semantic = semantic_map.flatten()385 flat_depth = z.flatten()386 387 # Filter out invalid points and sky points388 valid_mask = (flat_depth > 0) & (flat_depth < 1000) & (flat_semantic != SKY_TRAIN_ID)389 points = points[valid_mask]390 391 # Get segmentation colors for each point392 color_map = get_color_map(labels)393 seg_colors = np.zeros((len(flat_semantic), 3))394 395 for trainId, color in color_map.items():396 mask = flat_semantic == trainId397 seg_colors[mask] = color398 399 # Filter colors to match valid points400 colors = seg_colors[valid_mask] / 255.0 # Normalize to [0, 1]401 402 # Create Open3D point cloud403 pcd = o3d.geometry.PointCloud()404 pcd.points = o3d.utility.Vector3dVector(points)405 pcd.colors = o3d.utility.Vector3dVector(colors)406 407 return pcd408 409def create_plotly_pointcloud(pcd: o3d.geometry.PointCloud, downsample_factor: float = 0.1) -> go.Figure:410 """Create interactive Plotly 3D point cloud visualization."""411 # Downsample for performance412 if downsample_factor < 1.0:413 num_points = len(pcd.points)414 indices = np.random.choice(num_points, int(num_points * downsample_factor), replace=False)415 points = np.asarray(pcd.points)[indices]416 colors = np.asarray(pcd.colors)[indices]417 else:418 points = np.asarray(pcd.points)419 colors = np.asarray(pcd.colors)420 421 # Create 3D scatter plot422 fig = go.Figure(data=[go.Scatter3d(423 x=points[:, 0],424 y=points[:, 1], 425 z=points[:, 2],426 mode='markers',427 marker=dict(428 size=1,429 color=colors,430 opacity=0.8431 ),432 text=[f'Point {i}' for i in range(len(points))],433 hovertemplate='X: %{x:.2f}<br>Y: %{y:.2f}<br>Z: %{z:.2f}<extra></extra>'434 )])435 436 # Update layout for centered display437 fig.update_layout(438 scene=dict(439 xaxis_title='X (Horizontal)',440 yaxis_title='Y (Vertical)', 441 zaxis_title='Z (Depth)',442 aspectmode='data'443 ),444 title={445 'text': 'Interactive 3D Point Cloud (Colored by Segmentation, Sky Excluded)',446 'x': 0.5,447 'xanchor': 'center'448 },449 width=None, # Let it auto-size to container450 height=600,451 margin=dict(l=0, r=0, t=40, b=0), # Minimal margins452 autosize=True # Enable auto-sizing to container453 )454 455 # Set camera for bird's eye view that clearly shows 3D structure456 fig.update_layout(scene_camera=dict(457 up=dict(x=0, y=0, z=1), # Z-axis points up458 center=dict(x=0, y=0, z=0), # Center at origin459 eye=dict(x=0.5, y=-2.5, z=1.5) # View from above-back position460 ))461 462 return fig463 464def create_overlay_plot(rgb_image: Image.Image, semantic_map: np.ndarray, alpha: float = 0.5):465 """Create segmentation overlay plot without title and borders."""466 rgb_array = np.array(rgb_image)467 color_map = get_color_map(labels)468 colored_semantic_map = apply_color_map(semantic_map, color_map)469 470 # Create figure with exact image dimensions471 height, width = rgb_array.shape[:2]472 dpi = 100473 fig, ax = plt.subplots(1, 1, figsize=(width/dpi, height/dpi), dpi=dpi)474 475 # Remove all margins and padding476 fig.subplots_adjust(left=0, right=1, top=1, bottom=0)477 478 ax.imshow(rgb_array)479 ax.imshow(colored_semantic_map, alpha=alpha)480 ax.axis('off')481 482 buf = io.BytesIO()483 plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0, dpi=dpi)484 buf.seek(0)485 plt.close(fig)486 487 return Image.open(buf)488 489class PrecomputedExamplesManager:490 """Manages precomputed examples from output folder structure."""491 492 def __init__(self, output_dir: Path):493 self.output_dir = output_dir494 self.rgb_dir = output_dir / "rgb"495 self.segmentation_dir = output_dir / "segmentation"496 self.depth_dir = output_dir / "depth" 497 self.pointclouds_dir = output_dir / "pointclouds"498 self.examples = self._load_examples()499 500 def _load_examples(self) -> Dict[str, Dict]:501 """Load all available precomputed examples from output structure."""502 examples = {}503 504 if not self.output_dir.exists():505 print(f"Output directory {self.output_dir} not found.")506 return {}507 508 # Find all timestamps by looking at RGB files (the inputs)509 if not self.rgb_dir.exists():510 print(f"RGB directory {self.rgb_dir} not found.")511 return {}512 513 # Get all RGB files and extract timestamps514 timestamps = set()515 for rgb_file in self.rgb_dir.glob("rgb_*.png"):516 # Extract timestamp from filename like "rgb_20241215_143022.png"517 filename = rgb_file.stem518 if filename.startswith("rgb_"):519 timestamp = filename.replace("rgb_", "")520 timestamps.add(timestamp)521 522 print(f"Found {len(timestamps)} RGB input images")523 524 # For each timestamp, try to load the complete example525 for timestamp in sorted(timestamps, reverse=True): # Most recent first526 example_data = self._load_single_example(timestamp)527 if example_data:528 examples[timestamp] = example_data529 530 print(f"Loaded {len(examples)} precomputed examples from output directory")531 return examples532 533 def _load_single_example(self, timestamp: str) -> Optional[Dict]:534 """Load a single precomputed example by timestamp."""535 try:536 # Input file (required)537 rgb_path = self.rgb_dir / f"rgb_{timestamp}.png"538 539 # Output files (some may be optional)540 seg_path = self.segmentation_dir / f"segmentation_{timestamp}.png"541 depth_path = self.depth_dir / f"depth_{timestamp}.png"542 ply_path = self.pointclouds_dir / f"pointcloud_{timestamp}.ply"543 html_path = self.pointclouds_dir / f"pointcloud_{timestamp}.html"544 545 # Check if RGB input exists (required)546 if not rgb_path.exists():547 print(f"RGB input file missing for timestamp {timestamp}: {rgb_path}")548 return None549 550 # Check if at least segmentation output exists551 if not seg_path.exists():552 print(f"Segmentation output missing for timestamp {timestamp}: {seg_path}")553 return None554 555 # Create a display name from timestamp556 try:557 # Parse timestamp like "20241215_143022"558 if len(timestamp) >= 13 and "_" in timestamp:559 date_part = timestamp[:8]560 time_part = timestamp[9:15]561 # Format as "Dec 15, 2024 14:30"562 year = date_part[:4]563 month = date_part[4:6]564 day = date_part[6:8]565 hour = time_part[:2]566 minute = time_part[2:4]567 568 month_names = ["", "Jan", "Feb", "Mar", "Apr", "May", "Jun",569 "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"]570 month_name = month_names[int(month)] if 1 <= int(month) <= 12 else month571 572 display_name = f"{month_name} {int(day)}, {year} {hour}:{minute}"573 else:574 display_name = timestamp575 except:576 display_name = timestamp577 578 return {579 'name': display_name,580 'timestamp': timestamp,581 'rgb_path': rgb_path, # Input image582 'segmentation_path': seg_path, # Output583 'depth_path': depth_path if depth_path.exists() else None, # Output (optional)584 'pointcloud_ply_path': ply_path if ply_path.exists() else None, # Output (optional)585 'pointcloud_html_path': html_path if html_path.exists() else None, # Output (optional)586 'preview_image': self._create_preview_image(rgb_path, timestamp)587 }588 589 except Exception as e:590 print(f"Error loading example {timestamp}: {e}")591 return None592 593 def _create_preview_image(self, rgb_path: Path, timestamp: str) -> Image.Image:594 """Create a preview thumbnail from RGB input image."""595 try:596 image = Image.open(rgb_path)597 image.thumbnail((600, 450), Image.Resampling.LANCZOS)598 return image599 600 except Exception as e:601 print(f"Error creating preview for {timestamp}: {e}")602 return Image.new('RGB', (200, 150), color=(128, 128, 128))603 604 def get_example_names(self) -> List[str]:605 """Get list of available example names."""606 return [data['name'] for data in self.examples.values()]607 608 def get_example_previews(self) -> List[Tuple[Image.Image, str]]:609 """Get preview images for all examples."""610 previews = []611 for timestamp, data in self.examples.items():612 previews.append((data['preview_image'], data['name']))613 return previews614 615 def get_timestamp_by_name(self, name: str) -> Optional[str]:616 """Get timestamp by display name."""617 for timestamp, data in self.examples.items():618 if data['name'] == name:619 return timestamp620 return None621 622 def load_example_results(self, example_name: str) -> Tuple[Optional[Image.Image], Optional[Image.Image], Optional[go.Figure], str]:623 """Load precomputed results for an example."""624 if not example_name:625 return None, None, None, "Please select an example."626 627 # Find the timestamp for this example name628 timestamp = self.get_timestamp_by_name(example_name)629 if not timestamp or timestamp not in self.examples:630 return None, None, None, f"Example '{example_name}' not found."631 632 example_data = self.examples[timestamp]633 634 try:635 # Load output images636 segmentation_image = Image.open(example_data['segmentation_path'])637 638 depth_image = None639 if example_data['depth_path'] and example_data['depth_path'].exists():640 depth_image = Image.open(example_data['depth_path'])641 642 # Load point cloud if available643 point_cloud_fig = None644 if example_data['pointcloud_ply_path'] and example_data['pointcloud_ply_path'].exists():645 try:646 pcd = o3d.io.read_point_cloud(str(example_data['pointcloud_ply_path']))647 if len(pcd.points) > 0:648 point_cloud_fig = create_plotly_pointcloud(pcd, downsample_factor=1)649 else:650 print(f"Point cloud file {example_data['pointcloud_ply_path']} is empty")651 except Exception as e:652 print(f"Error loading point cloud: {e}")653 654 return segmentation_image, depth_image, point_cloud_fig, ""655 656 except Exception as e:657 return None, None, None, f"Error loading example results: {str(e)}"658 659# =============================================================================660# Main Application Class661# =============================================================================662 663class EnhancedSingleViewApp:664 def __init__(self):665 # Model configurations666 self.oneformer_config = ModelConfig(667 model_name="shi-labs/oneformer_cityscapes_swin_large",668 processor_name="shi-labs/oneformer_cityscapes_swin_large",669 task_type="semantic"670 )671 672 self.depth_config = DepthConfig(673 encoder="vitl",674 dataset="vkitti",675 max_depth=80,676 weights_path="depth_anything_v2_metric_vkitti_vitl.pth"677 )678 679 # Models680 self.oneformer_model = None681 self.depth_model = None682 self.segmentation_loaded = False683 self.depth_loaded = False684 685 # Precomputed examples manager686 self.examples_manager = PrecomputedExamplesManager(OUTPUT_DIR)687 688 # Online sample images (fallback)689 self.sample_images = {690 "Street Scene 1": "https://images.unsplash.com/photo-1449824913935-59a10b8d2000?w=800",691 "Street Scene 2": "https://images.unsplash.com/photo-1502920917128-1aa500764cbd?w=800", 692 "Urban Road": "https://images.unsplash.com/photo-1516738901171-8eb4fc13bd20?w=800",693 "City View": "https://images.unsplash.com/photo-1477959858617-67f85cf4f1df?w=800",694 "Highway": "https://images.unsplash.com/photo-1544620347-c4fd4a3d5957?w=800",695 }696 697 def download_sample_image(self, image_url: str) -> Image.Image:698 """Download a sample image from URL."""699 try:700 response = requests.get(image_url, timeout=10)701 response.raise_for_status()702 return Image.open(io.BytesIO(response.content)).convert('RGB')703 except Exception as e:704 print(f"Error downloading image: {e}")705 return Image.new('RGB', (800, 600), color=(128, 128, 128))706 707 def create_overlay_plot(self, rgb_image: Image.Image, semantic_map: np.ndarray, alpha: float = 0.5):708 """Create segmentation overlay plot without title and borders."""709 rgb_array = np.array(rgb_image)710 color_map = get_color_map(labels)711 colored_semantic_map = apply_color_map(semantic_map, color_map)712 713 # Create figure with exact image dimensions714 height, width = rgb_array.shape[:2]715 dpi = 100716 fig, ax = plt.subplots(1, 1, figsize=(width/dpi, height/dpi), dpi=dpi)717 718 # Remove all margins and padding719 fig.subplots_adjust(left=0, right=1, top=1, bottom=0)720 721 ax.imshow(rgb_array)722 ax.imshow(colored_semantic_map, alpha=alpha)723 ax.axis('off')724 725 buf = io.BytesIO()726 plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0, dpi=dpi)727 buf.seek(0)728 plt.close(fig)729 730 return Image.open(buf)731 732 def process_complete_pipeline(self, image: Image.Image):733 """Process image through complete pipeline: segmentation + depth + point cloud."""734 if image is None:735 return None, None, None, "Please upload an image."736 737 # Default values738 overlay_alpha = 0.5739 depth_colormap = "magma"740 downsample_factor = 0.1741 742 try:743 # Auto-load models if not loaded744 if not self.segmentation_loaded:745 if self.oneformer_model is None:746 self.oneformer_model = OneFormerModel(self.oneformer_config)747 self.oneformer_model.load_model()748 self.segmentation_loaded = True749 750 if not self.depth_loaded and DEPTH_AVAILABLE:751 if self.depth_model is None:752 self.depth_model = DepthAnythingV2Model(self.depth_config)753 self.depth_model.load_model()754 self.depth_loaded = True755 756 # Resize if too large757 original_size = image.size758 if max(image.size) > 1024:759 image.thumbnail((1024, 1024), Image.Resampling.LANCZOS)760 761 # Step 1: Semantic Segmentation762 task_inputs = ["semantic"]763 semantic_map = self.oneformer_model.segment_image(image, task_inputs=task_inputs)764 segmentation_overlay = self.create_overlay_plot(image, semantic_map, overlay_alpha)765 766 # Step 2: Depth Estimation (if available)767 depth_vis = None768 point_cloud_fig = None769 pcd = None770 771 if DEPTH_AVAILABLE and self.depth_loaded:772 depth_map = self.depth_model.estimate_depth(image)773 depth_vis = create_depth_visualization(depth_map, depth_colormap)774 775 # Step 3: Point Cloud with Segmentation Colors776 pcd = depth_to_point_cloud_with_segmentation(depth_map, image, semantic_map)777 point_cloud_fig = create_plotly_pointcloud(pcd, downsample_factor)778 779 # Generate comprehensive info780 unique_classes = np.unique(semantic_map)781 class_info = []782 total_pixels = semantic_map.size783 784 for class_id in unique_classes:785 if class_id < len(labels) and class_id != 255:786 label = labels[class_id]787 pixel_count = np.sum(semantic_map == class_id)788 percentage = (pixel_count / total_pixels) * 100789 if percentage > 0.1:790 class_info.append(f"- {label.name}: {percentage:.1f}%")791 792 # Point cloud statistics793 if point_cloud_fig is not None:794 num_points = len(pcd.points)795 downsampled_points = int(num_points * downsample_factor)796 point_cloud_info = f"""7973D Point Cloud:798- Total points: {num_points:,}799- Displayed points: {downsampled_points:,} ({downsample_factor*100:.0f}%)800- Sky points excluded801- Colors match segmentation classes"""802 else:803 point_cloud_info = "Point cloud not available (DepthAnythingV2 required)"804 805 # Depth statistics806 if depth_vis is not None and DEPTH_AVAILABLE:807 depth_stats = {808 'min': np.min(depth_map),809 'max': np.max(depth_map),810 'mean': np.mean(depth_map),811 'std': np.std(depth_map)812 }813 depth_info = f"""814Depth Estimation:815- Min depth: {depth_stats['min']:.2f}m816- Max depth: {depth_stats['max']:.2f}m 817- Mean depth: {depth_stats['mean']:.2f}m818- Std deviation: {depth_stats['std']:.2f}m819- Colormap: {depth_colormap}"""820 else:821 depth_info = "Depth estimation not available"822 823 info_text = f"""Complete vision pipeline processed successfully!824 825Models Used:826- OneFormer (Semantic Segmentation)827{f"- DepthAnythingV2 ({self.depth_config.encoder.upper()})" if DEPTH_AVAILABLE else "- DepthAnythingV2 (Not Available)"}828 829Image Processing:830- Original size: {original_size[0]}x{original_size[1]}831- Processed size: {image.size[0]}x{image.size[1]}832- Overlay transparency: {overlay_alpha:.1f}833 834Detected Classes:835{chr(10).join(class_info)}836{depth_info}837{point_cloud_info}838 839The point cloud shows 3D structure with each point colored according to its segmentation class. Sky points are excluded for better visualization."""840 841 return segmentation_overlay, depth_vis, point_cloud_fig, info_text842 843 except Exception as e:844 return None, None, None, f"Error processing pipeline: {str(e)}"845 846# Initialize the app847app = EnhancedSingleViewApp()848 849def process_uploaded_image(image):850 try:851 return app.process_complete_pipeline(image)852 except:853 return None, None, None854 855def process_sample_image(sample_choice):856 """Process sample image through complete pipeline.""" 857 if sample_choice and sample_choice in app.sample_images:858 image_url = app.sample_images[sample_choice]859 image = app.download_sample_image(image_url)860 return app.process_complete_pipeline(image)861 return None, None, None, "Please select a sample image."862 863def load_precomputed_example(evt: gr.SelectData):864 """Load precomputed example results from gallery selection."""865 if evt.index is not None:866 example_names = app.examples_manager.get_example_names()867 if evt.index < len(example_names):868 example_name = example_names[evt.index]869 seg_image, depth_image, pc_fig, info_text = app.examples_manager.load_example_results(example_name)870 return seg_image, depth_image, pc_fig871 return None, None, None872 873def get_example_previews():874 """Get preview images for the gallery."""875 previews = app.examples_manager.get_example_previews()876 if not previews:877 return []878 return previews879 880# =============================================================================881# Create Gradio Interface882# =============================================================================883 884def create_gradio_interface():885 """Create and return the enhanced single-view Gradio interface."""886 887 with gr.Blocks(888 title="Enhanced Computer Vision Pipeline",889 theme=gr.themes.Default()890 ) as demo:891 892 gr.Markdown("""893 # Street Scene 3D Reconstruction894 895 Upload an image or select an example to see:896 - **Semantic Segmentation** - Identify roads, buildings, vehicles, people, and other scene elements897 - **Depth Estimation** - Generate metric depth maps showing distance to objects 898 - **3D Point Cloud** - Interactive 3D reconstruction with semantic colors)899 """)900 901 with gr.Row():902 # Left Column: Controls and Input903 with gr.Column(scale=1):904 if CUDA_AVAILABLE:905 gr.Markdown("### Upload Image")906 907 uploaded_image = gr.Image(908 type="pil",909 label="Upload Image"910 )911 upload_btn = gr.Button("Process Image", variant="primary", size="lg")912 else:913 uploaded_image = gr.Image(visible=False) # Hidden placeholder914 upload_btn = gr.Button(visible=False) # Hidden placeholder915 916 gr.Markdown("### CPU Mode")917 gr.Markdown("⚠️ **Upload disabled**: DepthAnythingV2 requires CUDA. Using precomputed examples only.")918 919 gr.Markdown("### Examples")920 gr.Markdown("Click on an image to load the example:")921 922 # Example gallery (always visible)923 example_gallery = gr.Gallery(924 value=get_example_previews(),925 label="Example Images",926 show_label=False,927 elem_id="example_gallery",928 columns=2,929 rows=3,930 height="auto",931 object_fit="cover"932 )933 934 # Right Column: Results935 with gr.Column(scale=2):936 gr.Markdown("### Results")937 938 # Segmentation and Depth side by side939 with gr.Row():940 with gr.Column():941 gr.Markdown("#### Semantic Segmentation")942 segmentation_output = gr.Image(label="Segmentation Overlay")943 944 with gr.Column():945 gr.Markdown("#### Depth Estimation")946 depth_output = gr.Image(label="Depth Map")947 948 # Point Cloud below949 gr.Markdown("#### 3D Point Cloud")950 pointcloud_output = gr.Plot(label="Interactive 3D Point Cloud (Colored by Segmentation)")951 952 if CUDA_AVAILABLE:953 upload_btn.click(954 fn=process_uploaded_image,955 inputs=[uploaded_image],956 outputs=[segmentation_output, depth_output, pointcloud_output]957 )958 959 # Gallery selection loads example directly960 example_gallery.select(961 fn=load_precomputed_example,962 outputs=[segmentation_output, depth_output, pointcloud_output]963 )964 965 return demo966 967# =============================================================================968# Main Execution969# =============================================================================970 971if __name__ == "__main__":972 fix_lfs_on_startup()973 # Create and launch the interface974 demo = create_gradio_interface()975 976 print("Starting Enhanced Single-View Computer Vision App...")977 print("Complete Pipeline: Segmentation + Depth + 3D Point Cloud")978 print("Device:", "CUDA" if torch.cuda.is_available() else "CPU")979 print("Depth Available:", "YES" if DEPTH_AVAILABLE else "NO")980 print("Point Cloud Colors: Segmentation-based (Sky Excluded)")981 print(f"Output Directory: {OUTPUT_DIR.absolute()}")982 print(f"Available Examples: {len(app.examples_manager.examples)}")983 984 # Launch the app985 demo.launch(986 share=True, # Creates a public link987 debug=True, # Enable debugging988 server_name="0.0.0.0", # Allow external connections989 server_port=7860, # Default port990 show_error=True, # Show errors in the interface991 quiet=False # Show startup logs992 )