util/visualization.py

import os, sys, math
from typing import List, Optional, Tuple, Dict
import numpy as np
from PIL import Image, ImageDraw, ImageFont




PALETTE: List[Tuple[int, int, int]] = [
    (255, 80, 80),
    (80, 255, 80),
    (80, 160, 255),
    (255, 200, 80),
    (200, 80, 255),
    (80, 255, 220),
]





def _frame_idx_to_mmss(frame_idx: int, fps: float) -> str:
    sec = int(frame_idx / fps)
    mm = sec // 60
    ss = sec % 60
    return f"{mm:02d}:{ss:02d}"




# -------------------------
# Visualization utilities
# -------------------------
def _to_uint8_rgb(frame: np.ndarray) -> np.ndarray:
    if frame.dtype != np.uint8:
        frame = np.clip(frame * 255.0, 0, 255).astype(np.uint8)

    if frame.ndim == 2:
        frame = np.stack([frame] * 3, axis=-1)
    elif frame.ndim == 3 and frame.shape[-1] == 4:
        frame = frame[..., :3]
    return frame


def _load_font(font_size: int):
    try:
        return ImageFont.truetype("DejaVuSans.ttf", font_size)
    except Exception:
        return ImageFont.load_default()


def _build_closed_boundary_color_maps(
    ranges_closed: List[Tuple[int, int]],
    T: int,
    palette: List[Tuple[int, int, int]],
    end_exclusive = False,
):
    start_map, end_map = {}, {}
    for k, (s, e) in enumerate(ranges_closed):

        if end_exclusive:
            e = e - 1

        if s > e:
            s, e = e, s
        if e < 0 or s >= T:
            continue
        s = max(0, s)
        e = min(T - 1, e)
        rgba = (*palette[k % len(palette)], 255)
        start_map[s] = rgba
        end_map[e] = rgba
    return start_map, end_map




def visualize_concated_frames(
    frames: np.ndarray,
    out_dir: str,
    highlight_ranges_closed: Optional[List[Tuple[int, int]]],
    max_frames_per_img: int = 600,
    cols: int = 12,
    pad: int = 3,
    bg_color: Tuple[int, int, int] = (0, 0, 0),
    resize_to: Optional[Tuple[int, int]] = None,
    start_index: int = 0,           # time 的对齐的位置
    text_color: Tuple[int, int, int] = (255, 0, 0),
    font_size: int = 18,
    text_pad: Tuple[int, int] = (4, 2),
    draw_text_bg: bool = True,
    text_bg_rgba: Tuple[int, int, int, int] = (0, 0, 0, 160),
    out_prefix: str = "concat_",
    out_ext: str = ".jpg",
    jpg_quality: int = 75,
    bar_thickness: int = 14,
    palette: Optional[List[Tuple[int, int, int]]] = PALETTE,
    fps: Optional[float] = 24,
    end_range_exclusive = False,             # Whether end range is inclusive or exclusive
    verbose = False,
):

    os.makedirs(out_dir, exist_ok=True)
    font = _load_font(font_size)

    f0 = _to_uint8_rgb(frames[0])
    tile_h, tile_w = f0.shape[:2]
    frames_per_page = max_frames_per_img
    rows = math.ceil(frames_per_page / cols)

    canvas_w = cols * tile_w + (cols - 1) * pad
    canvas_h = rows * tile_h + (rows - 1) * pad

    def new_canvas():
        im = Image.new("RGB", (canvas_w, canvas_h), color=bg_color)
        return im, ImageDraw.Draw(im, "RGBA")

    start_map, end_map = ({}, {})
    if highlight_ranges_closed:
        start_map, end_map = _build_closed_boundary_color_maps(highlight_ranges_closed, len(frames), palette, end_exclusive=end_range_exclusive)

    page = 0
    global_idx = start_index
    canvas, draw = new_canvas()
    saved_paths = []

    # Iterate
    for i, fr in enumerate(frames):
        local_i = i % frames_per_page

        # Store
        if local_i == 0 and i > 0:
            saved_path = os.path.join(out_dir, f"{out_prefix}{page:04d}{out_ext}")
            canvas.save(saved_path, quality=jpg_quality)
            saved_paths.append(saved_path)

            # Update
            page += 1
            canvas, draw = new_canvas()
            

        r = local_i // cols
        c = local_i % cols
        x = c * (tile_w + pad)
        y = r * (tile_h + pad)

        pil = Image.fromarray(_to_uint8_rgb(fr))
        canvas.paste(pil, (x, y))

        if i in start_map:
            draw.rectangle((x, y, x + bar_thickness, y + tile_h), fill=start_map[i])
        if i in end_map:
            draw.rectangle((x + tile_w - bar_thickness, y, x + tile_w, y + tile_h), fill=end_map[i])

        text1 = str(global_idx)
        text2 = _frame_idx_to_mmss(global_idx, fps) if fps is not None else ""

        tx, ty = x + text_pad[0], y + text_pad[1]
        if draw_text_bg:
            bbox = draw.textbbox((tx, ty), text1 + "\n" + text2, font=font)
            draw.rectangle(bbox, fill=text_bg_rgba)

        draw.text((tx, ty), f"{text1}\n{text2}", font=font, fill=text_color)

        global_idx += 1


    # Save the last one
    saved_path = os.path.join(out_dir, f"{out_prefix}{page:04d}{out_ext}")
    canvas.save(saved_path, quality=jpg_quality)
    saved_paths.append(saved_path)

    
    # if verbose:
    #     print(f"Done. Total frames: {len(frames)} | Pages: {page + 1}")


    return saved_paths





def concat_image_lists_horizontal(
    list1: List[str],
    list2: List[str],
    out_dir: str,
    bar_width: int = 40,
    bar_color: Tuple[int, int, int] = (255, 0, 0),
    out_prefix: str = "merged_",
    out_ext: str = ".jpg",
    jpg_quality: int = 90,
    resize_mode: str = "match_height",  # ["match_height", "match_width", "none"]
    verbose: bool = True,
) -> List[str]:
    """
    Horizontally concatenate images from two path lists (same index),
    with a thick visual bar in between.

    Args:
        list1, list2: list of image file paths (must have same length)
        out_dir: directory to save merged images
        bar_width: thickness of the separator bar
        bar_color: RGB color of separator bar
        resize_mode:
            - "match_height": resize second image to match height
            - "match_width": resize second image to match width
            - "none": no resize (heights must match)
    Returns:
        List of saved image paths
    """

    assert len(list1) == len(list2), "Two lists must have same length"

    os.makedirs(out_dir, exist_ok=True)

    saved_paths = []

    for idx, (p1, p2) in enumerate(zip(list1, list2)):

        img1 = Image.open(p1).convert("RGB")
        img2 = Image.open(p2).convert("RGB")

        # ---------- Resize logic ----------
        if resize_mode == "match_height":
            if img1.height != img2.height:
                new_w = int(img2.width * img1.height / img2.height)
                img2 = img2.resize((new_w, img1.height), Image.BILINEAR)

        elif resize_mode == "match_width":
            if img1.width != img2.width:
                new_h = int(img2.height * img1.width / img2.width)
                img2 = img2.resize((img1.width, new_h), Image.BILINEAR)

        elif resize_mode == "none":
            assert img1.height == img2.height, \
                "Heights must match when resize_mode='none'"

        else:
            raise ValueError("resize_mode must be one of ['match_height', 'match_width', 'none']")

        # ---------- Create bar ----------
        bar = Image.new("RGB", (bar_width, img1.height), color=bar_color)

        # ---------- Create canvas ----------
        total_width = img1.width + bar_width + img2.width
        canvas = Image.new("RGB", (total_width, img1.height))

        canvas.paste(img1, (0, 0))
        canvas.paste(bar, (img1.width, 0))
        canvas.paste(img2, (img1.width + bar_width, 0))

        # ---------- Save ----------
        out_path = os.path.join(out_dir, f"{out_prefix}{idx:04d}{out_ext}")
        canvas.save(out_path, quality=jpg_quality)
        saved_paths.append(out_path)

    if verbose:
        print(f"Done. Saved {len(saved_paths)} merged images to {out_dir}")


    return saved_paths