04_segmentation.py

"""
04_segmentation.py — Object detection + instance segmentation on recording2.mp4

Model  : YOLOv8x-seg (Ultralytics)
Output : outputs/segmentation.mp4
         Shows bounding boxes + filled instance masks + class labels + confidence scores
"""

import os
import sys
import numpy as np
import cv2
from tqdm import tqdm

# ── Paths ──────────────────────────────────────────────────────────────────────
BASE       = os.path.dirname(os.path.abspath(__file__))
VIDEO_PATH = os.path.join(BASE, 'data', 'recording2.mp4')
OUT_PATH   = os.path.join(BASE, 'outputs', 'segmentation.mp4')
os.makedirs(os.path.join(BASE, 'outputs'), exist_ok=True)

# ── Config ─────────────────────────────────────────────────────────────────────
MODEL_NAME   = 'yolov8x-seg.pt'
CONF_THRESH  = 0.35
IOU_THRESH   = 0.45
MASK_ALPHA   = 0.40
OUT_W, OUT_H = 1920, 1080

# Palette (BGR) — cycling over detected instances
PALETTE = [
    (255, 100, 71), (135, 206, 235), (127, 255, 0), (255, 215, 0),
    (219,  112, 147), (100, 149, 237), (144, 238, 144), (255, 165, 0),
    (186, 85, 211), (0, 206, 209), (240, 230, 140), (32, 178, 170),
]


def get_color(idx: int):
    return PALETTE[idx % len(PALETTE)]


def draw_mask_and_box(frame: np.ndarray, mask_bool: np.ndarray,
                      box_xyxy, cls_name: str, conf: float,
                      color, mask_alpha: float):
    """Draw filled mask + bounding box + label on frame (in-place)."""
    h, w = frame.shape[:2]

    # Resize mask to frame size if needed
    mh, mw = mask_bool.shape
    if (mh, mw) != (h, w):
        mask_uint8 = cv2.resize(mask_bool.astype(np.uint8) * 255, (w, h))
        mask_bool  = mask_uint8 > 127

    # Filled semi-transparent mask
    overlay = frame.copy()
    overlay[mask_bool] = color
    cv2.addWeighted(overlay, mask_alpha, frame, 1 - mask_alpha, 0, frame)

    # Bounding box
    x1, y1, x2, y2 = [int(v) for v in box_xyxy]
    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)

    # Label background
    label = f"{cls_name} {conf:.2f}"
    (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 1)
    label_y = max(y1, th + 6)
    cv2.rectangle(frame, (x1, label_y - th - 6), (x1 + tw + 4, label_y), color, -1)
    cv2.putText(frame, label, (x1 + 2, label_y - 3),
                cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 0, 0), 1, cv2.LINE_AA)


def draw_frame_hud(frame: np.ndarray, frame_idx: int, total: int, n_obj: int):
    """HUD overlay: frame counter + object count."""
    h, w = frame.shape[:2]
    cv2.rectangle(frame, (0, 0), (340, 36), (10, 10, 20), -1)
    cv2.putText(frame, f"Frame {frame_idx:04d}/{total}  |  Objects: {n_obj}",
                (8, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (79, 195, 247), 1, cv2.LINE_AA)
    cv2.putText(frame, "YOLOv8x-seg", (w - 145, 24),
                cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 195, 79), 1, cv2.LINE_AA)


def main():
    try:
        from ultralytics import YOLO
    except ImportError:
        print("[04] ultralytics not installed. Run: pip install ultralytics")
        sys.exit(1)

    print(f"[04] Loading {MODEL_NAME} …")
    model = YOLO(MODEL_NAME)

    cap = cv2.VideoCapture(VIDEO_PATH)
    if not cap.isOpened():
        raise RuntimeError(f"Cannot open video: {VIDEO_PATH}")

    fps    = cap.get(cv2.CAP_PROP_FPS)
    orig_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    orig_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # CAP_PROP_FRAME_COUNT returns 30 due to bad MP4 metadata (actual = 1316).
    # Use the VTS record count as the real frame total for the progress bar.
    import sys, os
    sys.path.insert(0, BASE)
    from parse_imu_vts_lib import parse_vts
    vts_records = parse_vts(os.path.join(BASE, 'data', 'recording2.vts'))
    total = len(vts_records)          # 1316 — the true frame count

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out    = cv2.VideoWriter(OUT_PATH, fourcc, fps, (orig_w, orig_h))
    print(f"[04] Processing {total} frames {orig_w}×{orig_h} @ {fps:.1f} fps → {OUT_PATH}")

    pbar = tqdm(total=total, unit='frame', desc='Segment')

    frame_idx = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        # Run YOLO inference
        results = model(
            frame,
            conf=CONF_THRESH,
            iou=IOU_THRESH,
            verbose=False,
        )[0]

        n_obj = 0
        if results.boxes is not None and len(results.boxes) > 0:
            boxes  = results.boxes.xyxy.cpu().numpy()
            confs  = results.boxes.conf.cpu().numpy()
            clsids = results.boxes.cls.cpu().numpy().astype(int)
            names  = results.names

            has_masks = results.masks is not None
            masks_data = results.masks.data.cpu().numpy() if has_masks else None

            n_obj = len(boxes)
            for i, (box, conf, cls_id) in enumerate(zip(boxes, confs, clsids)):
                color    = get_color(i)
                cls_name = names[cls_id]

                if has_masks and i < len(masks_data):
                    mask_bool = masks_data[i] > 0.5
                    draw_mask_and_box(frame, mask_bool, box, cls_name, conf, color, MASK_ALPHA)
                else:
                    # No mask — draw box only
                    x1, y1, x2, y2 = [int(v) for v in box]
                    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                    label = f"{cls_name} {conf:.2f}"
                    cv2.putText(frame, label, (x1, max(y1-6, 12)),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.55, color, 1, cv2.LINE_AA)

        draw_frame_hud(frame, frame_idx, total, n_obj)
        out.write(frame)
        pbar.update(1)
        frame_idx += 1

    pbar.close()
    cap.release()
    out.release()
    print(f"[04] Done → {OUT_PATH}")


if __name__ == '__main__':
    main()