import gradio as gr import numpy as np import json import os import cv2 import sys import platform import torch from PIL import Image import model_utils import wandb from dotenv import load_dotenv # Load credentials from .env if present load_dotenv() # --- Configuration & State --- APP_CONFIG = { "system_os": platform.system(), "python_version": sys.version.split()[0], "torch_version": torch.__version__, "device": "cuda" if torch.cuda.is_available() else "cpu" } # Initialize W&B with system config INFERENCE_TABLE = None try: wandb.init( project="terrain-safety-v1", job_type="inference", config=APP_CONFIG ) # Define table schema once INFERENCE_TABLE = wandb.Table(columns=[ "model", "score", "safe_pct", "hazard_pct", "time_ms", "top_class", "confidence", "image_ref" ]) except Exception as e: print(f"Warning: W&B init failed: {e}. Logging disabled.") DEFAULT_MAPPING = { "safe": model_utils.SAFE_LABELS_DEFAULT, "hazard": model_utils.HAZARD_LABELS_DEFAULT } MODELS = { "SegFormer B0 (Fast)": "nvidia/segformer-b0-finetuned-ade-512-512", "SegFormer B2 (Balanced)": "nvidia/segformer-b2-finetuned-ade-512-512" } # --- Core Logic --- # --- Core Logic --- def process_image(image, opacity, mapping_json_str, enable_depth, enable_path, show_3d, max_range, rover_width, camera_hfov): if image is None: return None, None, None, None, "Please upload an image.", None # 1. Load Model (B2 - Balanced) model_name = MODELS["SegFormer B2 (Balanced)"] feature_extractor, model, device = model_utils.load_model(model_name) if model is None: return None, None, None, None, "Error loading model.", None # 2. Parse Mapping try: mapping_config = json.loads(mapping_json_str) except: mapping_config = DEFAULT_MAPPING # 3. Inference import time start_time = time.time() mask, logits = model_utils.predict_mask(image, (feature_extractor, model, device)) end_time = time.time() inference_time_ms = round((end_time - start_time) * 1000, 2) # 4. Safety Mapping id2label = model.config.id2label safety_mask = model_utils.map_classes_to_safety(mask, id2label, mapping_config) # 5. Depth & Pathfinding (Optional) path_coords = None depth_map_norm = None fig_3d = None depth_overlay = None # New Output if enable_depth or enable_path: # Load depth model (Small) depth_pipe = model_utils.load_depth_model("small") # Inference if depth_pipe: depth_map_norm = model_utils.estimate_depth(image, depth_pipe) # Pathfinding if enable_path and depth_map_norm is not None: path_coords = model_utils.compute_path(safety_mask, depth_map_norm, rover_width_m=rover_width, max_range=max_range, camera_hfov_deg=camera_hfov) elif enable_path: # Path enabled but depth failed or disabled -> use mask only path_coords = model_utils.compute_path(safety_mask, None, rover_width_m=rover_width, max_range=max_range, camera_hfov_deg=camera_hfov) # Create Depth Overlay if depth enabled (Now including path!) if enable_depth and depth_map_norm is not None: depth_overlay = model_utils.create_depth_plotly(image, depth_map_norm, max_range=max_range, path_coords=path_coords) # 6. Overlays hud_image = model_utils.create_hud(image, safety_mask, opacity=opacity, path_coords=path_coords) # Handle "No Safe Path Found" text if enabled but failed if enable_path and path_coords is None: import PIL.ImageDraw draw = PIL.ImageDraw.Draw(hud_image) try: # simple default font draw.text((20, 20), "No Safe Path Found", fill=(255, 0, 0)) except: pass # 7. 3D Visualizer if show_3d and depth_map_norm is not None: # Use new Mesh visualizer fig_3d = model_utils.create_3d_terrain(image, safety_mask, depth_map_norm, path_coords=path_coords) # 8. Stats stats = model_utils.compute_stats(mask, safety_mask, id2label, upsampled_logits=logits) # 9. Visualization colors mask_colored = Image.fromarray(mask.astype(np.uint8)).convert("P") # 10. W&B Log try: meta = { "model_id": model_name, "inference_time_ms": inference_time_ms, "image_size": image.size, "device": device, "depth_enabled": enable_depth, "path_enabled": enable_path } model_utils.log_inference_to_wandb(image, hud_image, stats, meta, table=INFERENCE_TABLE) except Exception as e: print(f"W&B Log failed: {e}") # Prepare outputs score_text = f"Safety Score: {stats['safety_score']}%" json_output = stats return hud_image, depth_overlay, mask_colored, json_output, score_text, fig_3d # --- UI Definition --- def create_demo(): with gr.Blocks(title="Terrain Safety Analysis") as demo: gr.Markdown("# 🛡️ Terrain Safety Analysis with SegFormer (v2)") gr.Markdown("Upload a terrain image to analyze Safe vs Hazard regions boundaries.") with gr.Row(): with gr.Column(scale=1): input_image = gr.Image(type="pil", label="Input Image") with gr.Accordion("Settings", open=True): opacity_slider = gr.Slider(0, 1, value=0.4, label="HUD Opacity") with gr.Accordion("Class Mapping (JSON)", open=False): mapping_editor = gr.Code( value=json.dumps(DEFAULT_MAPPING, indent=2), language="json", label="Safe/Hazard Definition" ) # Advanced Options with gr.Accordion("Advanced Features", open=True): enable_depth = gr.Checkbox(label="Enable Depth", value=False) with gr.Row(): enable_path = gr.Checkbox(label="Enable Pathfinding", value=False) show_3d = gr.Checkbox(label="Show 3D View", value=False) with gr.Accordion("Rover Settings", open=False): max_range_slider = gr.Slider(5, 200, value=50, step=5, label="Max Visible Range (meters)") rover_width_slider = gr.Slider(0.1, 2.0, value=0.45, step=0.05, label="Rover Width (meters)") camera_hfov_slider = gr.Slider(30, 120, value=35, step=5, label="Camera HFOV (degrees)") run_btn = gr.Button("Analyze Terrain", variant="primary") with gr.Column(scale=2): with gr.Row(): output_hud = gr.Image(label="HUD Prediction (with Path)", type="pil") output_mask = gr.Image(label="Raw Mask", type="pil") output_depth = gr.Plot(label="Depth Map (Hover for Distance)") output_3d = gr.Plot(label="3D Terrain View") score_display = gr.Label(label="Safety Score") output_json = gr.JSON(label="Detailed Stats") run_btn.click( process_image, inputs=[input_image, opacity_slider, mapping_editor, enable_depth, enable_path, show_3d, max_range_slider, rover_width_slider, camera_hfov_slider], outputs=[output_hud, output_depth, output_mask, output_json, score_display, output_3d] ) gr.Examples( examples=[["assets/sample.jpg"]], inputs=input_image ) return demo if __name__ == "__main__": demo = create_demo() demo.queue().launch(share=False)