from __future__ import annotations from typing import Dict, Tuple import numpy as np from PIL import Image, ImageDraw from .depth_pipeline import visualize_depth GRAD_ALPHA = 0.35 FLAT_ALPHA = 0.25 def _outline_mask(mask: np.ndarray | None) -> np.ndarray: """Compute a 1px outline around a boolean mask without external deps.""" if mask is None: return np.zeros((1, 1), dtype=bool) mask_bool = mask.astype(bool) h, w = mask_bool.shape padded = np.pad(mask_bool, 1, mode="edge") dilated = np.zeros_like(mask_bool, dtype=bool) for dy in (-1, 0, 1): for dx in (-1, 0, 1): if dx == 0 and dy == 0: continue dilated |= padded[1 + dy : 1 + dy + h, 1 + dx : 1 + dx + w] border = dilated & (~mask_bool) return border def make_safety_heatmap( rgb: Image.Image, safe_mask: np.ndarray, hazard_mask: np.ndarray, risk_map: np.ndarray, risk_threshold: float = 0.35, ): safe = np.clip(safe_mask.astype(np.float32), 0.0, 1.0) hazard = hazard_mask.astype(bool) risk = np.clip(risk_map.astype(np.float32), 0.0, 1.0) h, w = safe.shape safe_overlay = np.zeros((h, w, 4), dtype=np.uint8) safe_mask_bool = safe > 0.0 # Semi-transparent fill plus thicker outline for visibility (#00BF00) at ~80% alpha safe_overlay[safe_mask_bool, :] = (0, 191, 0, 204) safe_border = _outline_mask(safe_mask_bool) safe_border |= _outline_mask(safe_border) # thicken outline to ~2px safe_overlay[safe_border, :] = (0, 191, 0, 255) risk_focus = np.zeros_like(risk) risk_focus[risk > risk_threshold] = risk[risk > risk_threshold] hazard_intensity = np.where(hazard, np.maximum(risk_focus, 1.0), risk_focus) hazard_alpha = (np.clip(hazard_intensity, 0.0, 1.0) * 255).astype(np.uint8) hazard_overlay = np.zeros((h, w, 4), dtype=np.uint8) hazard_overlay[..., 0] = 255 # pure red for depth-based hazards hazard_overlay[..., 3] = hazard_alpha hazard_mask_bool = hazard hazard_border = _outline_mask(hazard_mask_bool) hazard_overlay[hazard_border, :] = (255, 0, 0, 255) safe_img = Image.fromarray(safe_overlay, mode="RGBA").resize(rgb.size, resample=Image.NEAREST) hazard_img = Image.fromarray(hazard_overlay, mode="RGBA").resize(rgb.size, resample=Image.NEAREST) score_gray = Image.fromarray((safe * 255).astype(np.uint8)).resize(rgb.size, resample=Image.NEAREST) return safe_img, hazard_img, score_gray def make_flatness_heatmap(std_map_vis: np.ndarray, target_size: tuple[int, int]) -> Image.Image: # Normalize and map to a simple turbo-like palette std_norm = std_map_vis if std_norm.max() > std_norm.min(): std_norm = (std_norm - std_norm.min()) / (np.ptp(std_norm) + 1e-6) cmap = np.array( [ [48, 18, 59], [65, 68, 135], [42, 120, 142], [34, 168, 132], [122, 209, 81], [253, 231, 36], ], dtype=np.float32, ) idx = np.clip((std_norm * (len(cmap) - 1)).astype(np.int32), 0, len(cmap) - 1) heat_rgb = cmap[idx] heat_overlay = np.zeros((std_norm.shape[0], std_norm.shape[1], 4), dtype=np.uint8) heat_overlay[..., :3] = heat_rgb.astype(np.uint8) heat_overlay[..., 3] = (np.clip(std_norm, 0.0, 1.0) * 160).astype(np.uint8) return Image.fromarray(heat_overlay, mode="RGBA").resize(target_size, resample=Image.BILINEAR) def build_result_layers( image: Image.Image, depth_raw: np.ndarray, std_map_vis: np.ndarray, grad_norm: np.ndarray, grad_thresh: float, safe_mask: np.ndarray, risk_map: np.ndarray, footprint_img_px: int, center_img: Tuple[int, int], water_mask: np.ndarray | None, road_mask: np.ndarray | None, roof_mask: np.ndarray | None, tree_mask: np.ndarray | None, hazard_mask: np.ndarray, ) -> Dict[str, Image.Image]: depth_vis = Image.fromarray(visualize_depth(depth_raw, cmap="Spectral")).resize( image.size, resample=Image.BILINEAR ) flatness_img = Image.fromarray((std_map_vis / (std_map_vis.max() + 1e-6) * 255).astype(np.uint8)).resize( image.size, resample=Image.NEAREST ) grad_img = Image.fromarray((grad_norm * 255).astype(np.uint8)).resize(image.size, resample=Image.BILINEAR) grad_mask_img = Image.fromarray(((grad_norm < grad_thresh).astype(np.uint8) * 255)).resize( image.size, resample=Image.NEAREST ) def _mask_to_image(mask: np.ndarray | None) -> Image.Image: if mask is None: return Image.new("L", image.size, 0) return Image.fromarray((mask.astype(np.uint8) * 255)).resize(image.size, resample=Image.NEAREST) water_mask_img = _mask_to_image(water_mask) road_mask_img = _mask_to_image(road_mask) roof_mask_img = _mask_to_image(roof_mask) tree_mask_img = _mask_to_image(tree_mask) def _hatched_overlay( mask: np.ndarray | None, color: tuple[int, int, int], alpha: int = 180, hatch_step: int = 8, hatch_thickness: int = 3, ) -> Image.Image: if mask is None: return Image.new("RGBA", image.size, (0, 0, 0, 0)) m = np.array( Image.fromarray((mask.astype(np.uint8) * 255)).resize(image.size, resample=Image.NEAREST) ).astype(bool) overlay = np.zeros((image.height, image.width, 4), dtype=np.uint8) overlay[..., :3] = np.array(color, dtype=np.uint8) overlay[..., 3] = np.where(m, alpha, 0).astype(np.uint8) if hatch_step > 0 and hatch_thickness > 0: xs = np.arange(image.width, dtype=np.int32) ys = np.arange(image.height, dtype=np.int32) grid_x, grid_y = np.meshgrid(xs, ys) stripe = ((grid_x + grid_y) % hatch_step) < hatch_thickness stripe_alpha = (alpha // 3) overlay[..., 3] = np.where(m & stripe, stripe_alpha, overlay[..., 3]) return Image.fromarray(overlay, mode="RGBA") # Segmentation hazards: black with hatch to separate from depth-risk red. hazard_color = (0, 0, 0) overlays = [] for mask, alpha in ( (water_mask, 130), (road_mask, 130), (tree_mask, 130), (roof_mask, 150), ): ov = _hatched_overlay(mask, hazard_color, alpha=alpha, hatch_step=24, hatch_thickness=2) overlays.append(ov) water_hazard_overlay, road_hazard_overlay, tree_hazard_overlay, roof_hazard_overlay = overlays for overlay in overlays: mask = np.array(overlay.getchannel("A")) > 0 if not mask.any(): continue border = _outline_mask(mask) if not border.any(): continue arr = np.array(overlay) arr[border, :] = (0, 0, 0, 255) arr[np.logical_and(border, arr[..., 3] > 0), 3] = 255 overlay.paste(Image.fromarray(arr, mode="RGBA")) safe_overlay, hazard_overlay, heat_gray = make_safety_heatmap(image, safe_mask, hazard_mask, risk_map) flat_heat_overlay = make_flatness_heatmap(std_map_vis, image.size) spot_overlay = Image.new("RGBA", image.size, (0, 0, 0, 0)) draw = ImageDraw.Draw(spot_overlay) cx_img, cy_img = center_img side_img = max(3, footprint_img_px | 1) half_img = side_img // 2 bx0 = cx_img - half_img by0 = cy_img - half_img bx1 = bx0 + side_img - 1 by1 = by0 + side_img - 1 clipped_x = False clipped_y = False if bx0 < 0: shift = -bx0 bx0 = 0 bx1 += shift clipped_x = True if bx1 >= image.width: shift = bx1 - (image.width - 1) bx1 = image.width - 1 bx0 = max(0, bx0 - shift) clipped_x = True if by0 < 0: shift = -by0 by0 = 0 by1 += shift clipped_y = True if by1 >= image.height: shift = by1 - (image.height - 1) by1 = image.height - 1 by0 = max(0, by0 - shift) clipped_y = True if clipped_x: cx_draw = int(round((bx0 + bx1) / 2.0)) else: cx_draw = int(round(min(max(cx_img, bx0), bx1))) if clipped_y: cy_draw = int(round((by0 + by1) / 2.0)) else: cy_draw = int(round(min(max(cy_img, by0), by1))) overlay_box = Image.new("RGBA", image.size, (0, 0, 0, 0)) box_draw = ImageDraw.Draw(overlay_box) fill = (255, 140, 0, 90) # translucent orange outline = (255, 140, 0, 255) # Crosshair sized 3x the landing box for clearer focus, thicker lines. cross_half = int(round(side_img * 1.5)) hx0 = max(0, cx_draw - cross_half) hx1 = min(image.width - 1, cx_draw + cross_half) hy0 = max(0, cy_draw - cross_half) hy1 = min(image.height - 1, cy_draw + cross_half) cross_width = 6 draw.line((hx0, cy_draw, hx1, cy_draw), fill=outline, width=cross_width) draw.line((cx_draw, hy0, cx_draw, hy1), fill=outline, width=cross_width) box_draw.rectangle((bx0, by0, bx1, by1), fill=fill, outline=outline, width=6) box_draw.line((cx_draw, by0, cx_draw, by1), fill=outline, width=3) box_draw.line((bx0, cy_draw, bx1, cy_draw), fill=outline, width=3) radius = 10 box_draw.ellipse((cx_draw - radius, cy_draw - radius, cx_draw + radius, cy_draw + radius), fill=outline) return { "RGB": image, "Depth": depth_vis, "Flatness map (std)": flatness_img, "Depth gradient": grad_img, "Gradient mask": grad_mask_img, "Water mask": water_mask_img, "Road mask": road_mask_img, "Roof mask": roof_mask_img, "Tree mask": tree_mask_img, "Safety heatmap overlay": safe_overlay, "Hazard overlay": hazard_overlay, "Water hazard overlay": water_hazard_overlay, "Road hazard overlay": road_hazard_overlay, "Tree hazard overlay": tree_hazard_overlay, "Roof hazard overlay": roof_hazard_overlay, "Flatness heatmap overlay": flat_heat_overlay, "Safety score": heat_gray, "Landing spot overlay": Image.alpha_composite(spot_overlay, overlay_box), } def compose_view( images_dict: dict, base_view: str, heat_on: bool, heat_alpha: float, risk_on: bool, risk_alpha: float, hazards_on: bool, grad_on: bool, flat_on: bool, flat_heat_on: bool, spot_on: bool, ) -> Image.Image: import gradio as gr if not images_dict: raise gr.Error("Run inference first, then select a view.") if base_view not in images_dict: raise gr.Error(f"Unknown view: {base_view}") base = images_dict.get(base_view) if base is None: raise gr.Error(f"No image for view: {base_view}") out = base.convert("RGBA") if heat_on and "Safety heatmap overlay" in images_dict: safe_overlay = images_dict["Safety heatmap overlay"] if safe_overlay is not None: safe_rgba = safe_overlay.convert("RGBA") alpha_factor = max(0.0, min(1.0, heat_alpha)) alpha_channel = np.array(safe_rgba.getchannel("A"), dtype=np.uint8) alpha_channel = (alpha_channel.astype(np.float32) * alpha_factor).astype(np.uint8) safe_rgba.putalpha(Image.fromarray(alpha_channel, mode="L")) out = Image.alpha_composite(out, safe_rgba) if risk_on and "Hazard overlay" in images_dict: hazard = images_dict.get("Hazard overlay") if hazard is not None: hazard_rgba = hazard.convert("RGBA") alpha_factor = max(0.0, min(1.0, risk_alpha)) alpha_channel = np.array(hazard_rgba.getchannel("A"), dtype=np.uint8) alpha_channel = (alpha_channel.astype(np.float32) * alpha_factor).astype(np.uint8) hazard_rgba.putalpha(Image.fromarray(alpha_channel, mode="L")) out = Image.alpha_composite(out, hazard_rgba) if hazards_on: if "Water hazard overlay" in images_dict: water = images_dict.get("Water hazard overlay") if water is not None: out = Image.alpha_composite(out, water.convert("RGBA")) if "Road hazard overlay" in images_dict: road = images_dict.get("Road hazard overlay") if road is not None: out = Image.alpha_composite(out, road.convert("RGBA")) if "Tree hazard overlay" in images_dict: tree = images_dict.get("Tree hazard overlay") if tree is not None: out = Image.alpha_composite(out, tree.convert("RGBA")) if "Roof hazard overlay" in images_dict: roof = images_dict.get("Roof hazard overlay") if roof is not None: out = Image.alpha_composite(out, roof.convert("RGBA")) if grad_on and "Depth gradient" in images_dict: grad_img = images_dict["Depth gradient"] if grad_img is not None: grad_rgba = grad_img.convert("RGBA") grad_rgba.putalpha(int(GRAD_ALPHA * 255)) out = Image.alpha_composite(out, grad_rgba) if flat_on and "Flatness map (std)" in images_dict: flat_img = images_dict["Flatness map (std)"] if flat_img is not None: flat_rgba = flat_img.convert("RGBA") flat_rgba.putalpha(int(FLAT_ALPHA * 255)) out = Image.alpha_composite(out, flat_rgba) if flat_heat_on and "Flatness heatmap overlay" in images_dict: flat_heat = images_dict["Flatness heatmap overlay"] if flat_heat is not None: flat_heat_rgba = flat_heat.convert("RGBA") out = Image.alpha_composite(out, flat_heat_rgba) if spot_on and "Landing spot overlay" in images_dict: spot = images_dict["Landing spot overlay"] if spot is not None: out = Image.alpha_composite(out, spot.convert("RGBA")) return out.convert("RGB") __all__ = ["build_result_layers", "compose_view", "make_safety_heatmap"]