visualize.py

import numpy as np
import matplotlib.pyplot as plt
import imageio
import os
from PIL import Image
from tqdm import tqdm
import open3d as o3d
import shutil



def depthmap_viz(depth, min_d=0.75, max_d=2.5, output_filename="depth_data.png"):
    depth = np.clip(depth, min_d, max_d)
    depth = (depth - min_d) / (max_d - min_d)
    depth_image = depth[:, :, 0]

    fig, ax = plt.subplots(figsize=(8, 6))
    ax.imshow(depth_image, cmap='gray', vmin=min_d, vmax=max_d)

    if output_filename:
        plt.savefig(output_filename, bbox_inches='tight')
    else:
        print("No output filename provided. Please specify a filename to save the visualization.")

    plt.close(fig)

def viz_pts_3d(pts, xrange=None, yrange=None, zrange=None, title=None, elev=-135, azim_start=0, azim_end=360, output_filename="cam_pnts_3d.gif"):
    # Create a temporary directory to store frames
    temp_dir = "temp_frames"
    os.makedirs(temp_dir, exist_ok=True)

    # Define the range of azimuth angles
    azim_range = np.linspace(azim_start, azim_end, 60)  # 60 frames

    for i, azim in tqdm(enumerate(azim_range), desc = "Visualizing"):
        fig = plt.figure(figsize=(10, 10))
        ax = fig.add_subplot(111, projection='3d')

        ax.scatter(pts[0, :], pts[1, :], pts[2, :], s=1)

        ax.set_xlabel('X [m]')
        ax.set_ylabel('Y [m]')
        ax.set_zlabel('Z [m]')

        if xrange is not None:
            ax.set_xlim(xrange)
        if yrange is not None:
            ax.set_ylim(yrange)
        if zrange is not None:
            ax.set_zlim(zrange)
        if title is not None:
            ax.set_title(title)
        ax.view_init(elev, azim)

        # Save the frame
        plt.savefig(os.path.join(temp_dir, f"frame_{i:03d}.png"), bbox_inches='tight')
        plt.close(fig)

    # Load frames and ensure they have the same shape
    frames = []
    frame_paths = sorted(os.listdir(temp_dir))
    first_frame = imageio.imread(os.path.join(temp_dir, frame_paths[0]))
    for frame_path in frame_paths:
        frame = imageio.imread(os.path.join(temp_dir, frame_path))
        # Resize frame if necessary
        if frame.shape != first_frame.shape:
            frame = np.array(Image.fromarray(frame).resize(first_frame.shape[1::-1]))
        frames.append(frame)

    # Create a GIF from the frames
    imageio.mimsave(output_filename, frames, fps=10, loop=0)

    # Clean up temporary frames
    shutil.rmtree(temp_dir)


def optimized_viz_pts_3d(pts, output_filename):
    # Convert to Open3D point cloud
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(pts.T)
    pcd = pcd.voxel_down_sample(voxel_size=0.05)

    # Initialize renderer
    width, height = 800, 800
    renderer = o3d.visualization.rendering.OffscreenRenderer(width, height)
    material = o3d.visualization.rendering.MaterialRecord()
    material.shader = "defaultUnlit"

    # Add geometry
    renderer.scene.set_background([0, 0, 0, 1])
    renderer.scene.add_geometry("pcd", pcd, material)

    # Compute camera parameters
    bounds = pcd.get_axis_aligned_bounding_box()
    center = bounds.get_center()
    extent = np.linalg.norm(bounds.get_extent())
    eye = center + np.array([0, 0, extent * 2.0])   # camera is "in front" of object
    up = np.array([0, 1, 0])                        # y-up

    # Set camera
    renderer.setup_camera(60.0, center, eye, up)

    # Render rotating frames
    images = []
    for i in range(60):
        angle = np.deg2rad(i * 6)
        cam_pos = center + np.array([
            np.sin(angle) * extent * 2.0,
            0,
            np.cos(angle) * extent * 2.0
        ])
        renderer.setup_camera(60.0, center, cam_pos, up)

        img = renderer.render_to_image()
        img_np = np.asarray(img)[:, :, :3]  # drop alpha
        images.append(img_np)

    # Save GIF
    imageio.mimsave(output_filename, images, fps=20, loop=0)



def show_mask(mask, ax, obj_id=None, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        cmap = plt.get_cmap("tab10")
        cmap_idx = 0 if obj_id is None else obj_id
        color = np.array([*cmap(cmap_idx)[:3], 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)


def show_points(coords, labels, ax, marker_size=200):
    pos_points = coords[labels==1]
    neg_points = coords[labels==0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)


def show_box(box, ax):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2))