app.py

import gradio as gr
import numpy as np
import laspy
import trimesh
import tempfile
import os
import warnings

# Standard imports
import asyncio

def load_laz(file_path):
    """Load LAZ/LAS file and return all data."""
    las = laspy.read(file_path)
    
    # Get XYZ - center around origin for visualization
    points = np.vstack([las.x, las.y, las.z]).T
    centroid = points.mean(axis=0)
    points = points - centroid
    
    # Get RGB colors (convert 16-bit to 8-bit)
    rgb = None
    if hasattr(las, 'red') and len(las.red) > 0:
        rgb = np.vstack([las.red, las.green, las.blue]).T
        rgb = (rgb / 256).astype(np.uint8)  # 16-bit to 8-bit
    
    # Get PredInstance if available
    pred_instance = None
    if 'PredInstance' in las.point_format.extra_dimension_names:
        pred_instance = np.array(las['PredInstance'])
    
    # Get intensity
    intensity = None
    if hasattr(las, 'intensity'):
        intensity = np.array(las.intensity)
    
    return {
        'points': points,
        'rgb': rgb,
        'pred_instance': pred_instance,
        'intensity': intensity,
        'total_count': len(points)
    }

def get_instance_colors(pred_instance):
    """Generate distinct colors for each instance."""
    # Distinct colors for instances
    instance_colors = np.array([
        [128, 128, 128],  # -1: Gray (unclassified)
        [255, 0, 0],      # 0: Red
        [0, 255, 0],      # 1: Green
        [0, 0, 255],      # 2: Blue
        [255, 255, 0],    # 3: Yellow
        [255, 0, 255],    # 4: Magenta
        [0, 255, 255],    # 5: Cyan
        [255, 128, 0],    # 6: Orange
        [128, 0, 255],    # 7: Purple
        [0, 255, 128],    # 8: Spring Green
        [255, 128, 128],  # 9: Light Red
        [128, 255, 128],  # 10: Light Green
        [128, 128, 255],  # 11: Light Blue
        [255, 255, 128],  # 12: Light Yellow
        [255, 128, 255],  # 13: Pink
        [128, 255, 255],  # 14: Light Cyan
    ], dtype=np.uint8)
    
    # Map instance IDs to colors (-1 maps to index 0)
    colors = np.zeros((len(pred_instance), 3), dtype=np.uint8)
    for i, inst in enumerate(pred_instance):
        idx = int(inst) + 1  # -1 -> 0, 0 -> 1, etc.
        idx = max(0, min(idx, len(instance_colors) - 1))
        colors[i] = instance_colors[idx]
    
    return colors

def get_elevation_colors(points, colormap_name):
    """Apply elevation-based colormap."""
    import matplotlib.pyplot as plt
    
    z = points[:, 2]
    z_norm = (z - z.min()) / (z.max() - z.min() + 1e-8)
    
    cmap = plt.get_cmap(colormap_name)
    colors = (cmap(z_norm)[:, :3] * 255).astype(np.uint8)
    
    return colors

def get_intensity_colors(intensity):
    """Convert intensity to grayscale colors."""
    i_norm = (intensity - intensity.min()) / (intensity.max() - intensity.min() + 1e-8)
    gray = (i_norm * 255).astype(np.uint8)
    return np.stack([gray, gray, gray], axis=1)

def load_example_file(example_name):
    """Load example file and return the file path."""
    import os
    base_dir = os.path.dirname(__file__)
    example_path = os.path.join(base_dir, example_name)
    if os.path.exists(example_path):
        return example_path
    return None

def visualize(file, color_mode, colormap, max_points):
    """Main visualization function."""
    if file is None:
        return None, "⚠️ Please upload a LAZ/LAS file"
    
    try:
        # Load data
        data = load_laz(file)
        points = data['points']
        total = data['total_count']
        
        # Subsample if needed
        if len(points) > max_points:
            indices = np.random.choice(len(points), int(max_points), replace=False)
            points = points[indices]
            rgb = data['rgb'][indices] if data['rgb'] is not None else None
            pred_instance = data['pred_instance'][indices] if data['pred_instance'] is not None else None
            intensity = data['intensity'][indices] if data['intensity'] is not None else None
        else:
            indices = None
            rgb = data['rgb']
            pred_instance = data['pred_instance']
            intensity = data['intensity']
        
        # Get colors based on mode
        if color_mode == "RGB (Original)":
            if rgb is not None:
                colors = rgb
            else:
                colors = get_elevation_colors(points, colormap)
        elif color_mode == "Instance Segmentation":
            if pred_instance is not None:
                colors = get_instance_colors(pred_instance)
            else:
                return None, "❌ This file does not have PredInstance data"
        elif color_mode == "Elevation":
            colors = get_elevation_colors(points, colormap)
        elif color_mode == "Intensity":
            if intensity is not None:
                colors = get_intensity_colors(intensity)
            else:
                colors = get_elevation_colors(points, colormap)
        else:
            colors = get_elevation_colors(points, colormap)
        
        # Add alpha channel
        alpha = np.full((len(colors), 1), 255, dtype=np.uint8)
        colors_rgba = np.hstack([colors, alpha])
        
        # Create point cloud and export
        cloud = trimesh.PointCloud(points, colors=colors_rgba)
        
        tmp = tempfile.NamedTemporaryFile(suffix='.glb', delete=False)
        tmp.close()
        cloud.export(tmp.name, file_type='glb')
        
        # Stats
        instance_info = ""
        if pred_instance is not None and color_mode == "Instance Segmentation":
            unique, counts = np.unique(pred_instance, return_counts=True)
            instance_info = "\n\n**Instance Breakdown:**\n"
            for u, c in sorted(zip(unique, counts), key=lambda x: -x[1])[:10]:
                instance_info += f"- Instance {int(u)}: {c:,} pts\n"
        
        stats = f"""
### 📊 Point Cloud Statistics
| Property | Value |
|----------|-------|
| Total Points | {total:,} |
| Displayed | {len(points):,} |
| X Range | {points[:,0].min():.2f} to {points[:,0].max():.2f} |
| Y Range | {points[:,1].min():.2f} to {points[:,1].max():.2f} |
| Z Range | {points[:,2].min():.2f} to {points[:,2].max():.2f} |
| Color Mode | {color_mode} |
| Has RGB | {'✅' if data['rgb'] is not None else '❌'} |
| Has Segmentation | {'✅' if data['pred_instance'] is not None else '❌'} |
{instance_info}
        """
        
        return tmp.name, stats
        
    except Exception as e:
        import traceback
        return None, f"❌ Error: {str(e)}\n```\n{traceback.format_exc()}\n```"

# UI
with gr.Blocks(title="GeoSpatial-LiDAR-3D Point Cloud Visualizer") as demo:
    gr.Markdown("# 🌍 GeoSpatial-LiDAR-3D Point Cloud Visualizer")
    gr.Markdown("Upload LAZ/LAS files with support for RGB colors and instance segmentation")
    
    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(
                label="📂 Upload LAS/LAZ File",
                file_types=[".las", ".laz"]
            )
            
            gr.Markdown("### 📋 Or use an example:")
            example_dropdown = gr.Dropdown(
                choices=["tree_raw.laz", "tree_segmentation.laz"],
                label="📂 Examples",
                value="tree_raw.laz"
            )
            btn_load_example = gr.Button("📂 Load Example", size="sm")
            
            color_mode = gr.Radio(
                choices=["RGB (Original)", "Instance Segmentation", "Elevation", "Intensity"],
                value="RGB (Original)",
                label="🎨 Color Mode"
            )
            
            colormap = gr.Dropdown(
                choices=["viridis", "terrain", "plasma", "inferno", "Spectral", "coolwarm"],
                value="viridis",
                label="Elevation Colormap (for Elevation mode)",
                visible=True
            )
            
            max_points = gr.Slider(
                50000, 2000000,
                value=500000,
                step=50000,
                label="📊 Max Points to Display"
            )
            
            btn = gr.Button("🚀 Visualize", variant="primary", size="lg")
            
            gr.Markdown("""
            ---
            ### 🎨 Color Modes
            - **RGB**: Original colors from file
            - **Instance Segmentation**: Color by PredInstance (if available)
            - **Elevation**: Color by height (Z)
            - **Intensity**: Grayscale by intensity
            """)
        
        with gr.Column(scale=2):
            model_output = gr.Model3D(
                label="3D Visualization",
                height=550,
                clear_color=(0.1, 0.1, 0.15, 1.0)
            )
            stats_output = gr.Markdown("*Upload a file and click Visualize*")
    
    btn.click(
        visualize,
        inputs=[file_input, color_mode, colormap, max_points],
        outputs=[model_output, stats_output]
    )
    
    # Example file handlers
    def load_and_visualize_example(example_name):
        import os
        if example_name is None:
            return None, "⚠️ Please select an example file"
        base_dir = os.path.dirname(os.path.abspath(__file__))
        example_path = os.path.join(base_dir, "example", example_name)
        if os.path.exists(example_path):
            return visualize(example_path, "RGB (Original)", "viridis", 500000)
        return None, f"❌ Example file '{example_name}' not found at {example_path}"
    
    btn_load_example.click(
        load_and_visualize_example,
        inputs=[example_dropdown],
        outputs=[model_output, stats_output]
    )

if __name__ == "__main__":
    # Launch with MCP server support for integration with other tools
    # To enable MCP server, set mcp_server=True or use GRADIO_MCP_SERVER env var
    demo.launch(ssr_mode=False, mcp_server=False)