visualization.py

"""
visualization.py - Colour assignment, PLY export, and optional 3-D viewer.
"""

import numpy as np
import open3d as o3d

from utils import LABEL_COLORS, LABEL_NOISE


def assign_colors_by_label(semantic_labels: np.ndarray) -> np.ndarray:
    """Map per-point semantic label strings to RGB colours.

    Colour scheme (matches common colour conventions for indoor scenes):
      FLOOR     → blue   (  0,   0, 255)
      CEILING   → white  (255, 255, 255)
      WALL      → red    (255,   0,   0)
      FURNITURE → green  (  0, 255,   0)
      NOISE     → grey   (128, 128, 128)

    Args:
        semantic_labels: 1-D object array of per-point label strings (shape N,).

    Returns:
        Uint8 RGB array of shape (N, 3).
    """
    colors = np.zeros((semantic_labels.shape[0], 3), dtype=np.uint8)
    for label, rgb in LABEL_COLORS.items():
        mask = semantic_labels == label
        colors[mask] = rgb
    return colors


def save_segmented_ply(
    points: np.ndarray,
    colors: np.ndarray,
    output_path: str,
) -> None:
    """Write a colour-coded point cloud to a PLY file.

    The colours are stored as per-vertex ``red``, ``green``, ``blue`` uint8
    properties, which is the format expected by CloudCompare, MeshLab, and
    most other standard viewers.

    Args:
        points:      Float array of shape (N, 3).
        colors:      Uint8 RGB array of shape (N, 3).
        output_path: Destination file path (overwritten if it exists).
    """
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points.astype(np.float64))

    # Open3D expects colours in [0, 1] float range
    pcd.colors = o3d.utility.Vector3dVector(colors.astype(np.float64) / 255.0)

    success = o3d.io.write_point_cloud(output_path, pcd, write_ascii=False)
    if not success:
        raise RuntimeError(f"Open3D failed to write PLY to: {output_path}")
    print(f"  Segmented PLY saved → {output_path}")


def visualize_point_cloud(
    points: np.ndarray,
    colors: np.ndarray,
    window_title: str = "3D Room Semantic Segmentation",
) -> None:
    """Display a colour-coded point cloud in an interactive Open3D window.

    Controls (shown in the viewer's help overlay):
      Left-drag  – rotate
      Right-drag – pan
      Scroll     – zoom
      Q / Esc    – close

    Args:
        points:       Float array of shape (N, 3).
        colors:       Uint8 RGB array of shape (N, 3).
        window_title: Title bar text for the viewer window.
    """
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points.astype(np.float64))
    pcd.colors = o3d.utility.Vector3dVector(colors.astype(np.float64) / 255.0)

    print("  Launching 3-D viewer … close the window to continue.")
    o3d.visualization.draw_geometries(
        [pcd],
        window_name=window_title,
        width=1280,
        height=800,
        point_show_normal=False,
    )


def save_topdown_png(
    points: np.ndarray,
    colors: np.ndarray,
    output_path: str,
    resolution: float = 0.02,
) -> None:
    """Render a 2-D top-down (floor-plan) view and save it as PNG.

    Each pixel represents a *resolution*-metre square cell.  The brightest
    colour among all points that fall into a cell is used for that pixel,
    which gives a clean projection even for dense clouds.

    Args:
        points:      Float array of shape (N, 3).
        colors:      Uint8 RGB array of shape (N, 3).
        output_path: Destination PNG path.
        resolution:  Metres per pixel (smaller = higher resolution).
    """
    try:
        from PIL import Image
    except ImportError:
        print("  Pillow not installed – skipping floor-plan PNG export.")
        return

    x, y = points[:, 0], points[:, 1]
    x_min, y_min = x.min(), y.min()

    cols = np.floor((x - x_min) / resolution).astype(np.int32)
    rows = np.floor((y - y_min) / resolution).astype(np.int32)

    img_width = cols.max() + 1
    img_height = rows.max() + 1

    canvas = np.zeros((img_height, img_width, 3), dtype=np.uint8)
    # Paint in Z order so higher points overwrite lower ones (walls over floor)
    z_order = np.argsort(points[:, 2])
    for idx in z_order:
        canvas[rows[idx], cols[idx]] = colors[idx]

    img = Image.fromarray(canvas, mode="RGB")
    img = img.transpose(Image.FLIP_TOP_BOTTOM)  # Y axis convention
    img.save(output_path)
    print(f"  Floor-plan PNG saved → {output_path}")