tools/visualization.py

import sys
import os
from os import path as osp
import argparse
import warnings
import torch
import numpy as np
from PIL import Image
from detectron2.config import instantiate, LazyConfig
from torchvision import transforms

from typing import List, Tuple, Optional, Union

sys.path.append(osp.dirname(osp.dirname(__file__)))
from utils import *


warnings.simplefilter(action="ignore", category=FutureWarning)
Color = Tuple[int, int, int]  # BGR (OpenCV)


def overlay_heatmap_on_image(image_bgr: np.ndarray,
                             heatmap_01: np.ndarray,
                             boxes: Tuple[int, int, int, int],
                             gaze_xy: Tuple[float, float],            # NEW: (x, y)
                             inout: bool = True,                     # NEW: whether in‐out prediction
                             color: Tuple[int, int, int] = (0, 255, 0),  # BGR
                             color_inout: Tuple[int, int, int] = (255, 0, 0),  # BGR for in-out
                             thickness: int = 4,
                             alpha: float = 0.2,                       # small blend
                             colormap: int = cv2.COLORMAP_JET,
                             mask_thresh: float = 0.10                  # ignore tiny heat
                             ) -> Tuple[np.ndarray, np.ndarray]:
    """
    Draws:
      1) Heatmap overlay (masked by mask_thresh)
      2) Head bounding box
      3) A green line from head-box center to gaze point
      4) A thick green dot at the gaze point

    Returns:
      overlay_bgr: image with overlays
      colored_heat_bgr: colorized heatmap (BGR)
    """
    H, W = image_bgr.shape[:2]

    # # --- heatmap -> resized -> [0,1] -> uint8 for colormap ---
    heat_resized = cv2.resize(heatmap_01.astype(np.float32), (W, H), interpolation=cv2.INTER_LINEAR)
    heat_resized = np.clip(heat_resized, 0.0, 1.0)
    heat_uint8 = (heat_resized * 255).round().astype(np.uint8)
    colored_heat_bgr = cv2.applyColorMap(heat_uint8, colormap)

    # --- inputs to float32 for blending ---
    img_u8 = image_bgr if image_bgr.dtype == np.uint8 else np.clip(image_bgr, 0, 255).astype(np.uint8)
    img_f  = img_u8.astype(np.float32)
    heat_f = colored_heat_bgr.astype(np.float32)

    # # --- standard blended image ---
    blended = cv2.addWeighted(heat_f, alpha, img_f, 1 - alpha, 0.0)

    # # --- only apply where heat is meaningful ---
    mask = (heat_resized >= mask_thresh).astype(np.float32)[..., None]  # (H,W,1)
    overlay_f = np.where(mask > 0, blended, img_f)
    overlay_bgr = np.clip(overlay_f, 0, 255).astype(np.uint8)
    
    # overlay_bgr = img_u8

    # --- robust single-box handling ---
    b = np.array(boxes, dtype=np.float32).reshape(-1)
    if b.size != 4:
        raise ValueError(f"'boxes' must be 4 numbers, got shape {np.array(boxes).shape} / size {b.size}")
    x1, y1, x2, y2 = map(int, b)
    x1 = int(max(0, min(W - 1, x1)));  y1 = int(max(0, min(H - 1, y1)))
    x2 = int(max(0, min(W - 1, x2)));  y2 = int(max(0, min(H - 1, y2)))
    if x2 < x1: x1, x2 = x2, x1
    if y2 < y1: y1, y2 = y2, y1
    
    if inout:
        # Draw head bounding box
        cv2.rectangle(overlay_bgr, (x1, y1), (x2, y2), color, thickness)

        # --- compute head center ---
        cx = int(round((x1 + x2) * 0.5))
        cy = int(round((y1 + y2) * 0.5))

        # --- handle gaze point (clip to image and draw) ---
        gx, gy = gaze_xy
        gx_i = int(round(max(0, min(W - 1, gx))))
        gy_i = int(round(max(0, min(H - 1, gy))))

        # Draw line from head center to gaze point
        line_thickness = max(2, thickness)  # slightly robust
        cv2.line(overlay_bgr, (cx, cy), (gx_i, gy_i), color, line_thickness, lineType=cv2.LINE_AA)

        # Draw a thicker filled dot at the gaze point
        gaze_radius = max(12, thickness * 2)
        cv2.circle(overlay_bgr, (gx_i, gy_i), gaze_radius, color, thickness=-1, lineType=cv2.LINE_AA)
    else:
        # Draw head bounding box in different color for out-of-frame
        cv2.rectangle(overlay_bgr, (x1, y1), (x2, y2), color_inout, thickness)

    return overlay_bgr

            
            
def inference_gaze(image_paths, texts, model, visualization_dir, use_dark_inference=True):
    img = [Image.open(image_path).convert("RGB") for image_path in image_paths]
    save_paths = [image_path.split("/")[-1] for image_path in image_paths]
    img_size = [img.size for img in img]
    
    if not isinstance(texts, List):
        texts = [texts] * len(image_paths)
            
    image_transform = transforms.Compose(
        [
            transforms.Resize((512, 512)),
            transforms.ToTensor(),
            transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
        ]
    )
    
    for i in range(len(img)):
        img[i] = image_transform(img[i]).unsqueeze(0).cuda()
    
    img = torch.cat(img, dim=0)
    
    with torch.no_grad():
        gaze_heatmap_pred, inout_pred, bbox_pred = model.inference(img, texts)
    
    gaze_heatmap_pred = (
        gaze_heatmap_pred.squeeze(1).cpu().detach().numpy()
    )
    
    inout_pred = (
        inout_pred.cpu().detach().numpy()
    )
    
    bbox_pred = (
        bbox_pred.cpu().detach().numpy()
    )
    
    for b_i in range(len(gaze_heatmap_pred)):
        visualization_save_path = osp.join(
            visualization_dir,
            save_paths[b_i]
        )
        
        inout = (inout_pred[b_i] > 0.5)
        
        head_bbox = [0.0, 0.0, img_size[b_i][0], img_size[b_i][1]]
        
        head_bbox[0] = (bbox_pred[b_i][0] - bbox_pred[b_i][2] / 2) * img_size[b_i][0]  # x1
        head_bbox[1] = (bbox_pred[b_i][1] - bbox_pred[b_i][3] / 2) * img_size[b_i][1]  # y1
        head_bbox[2] = (bbox_pred[b_i][0] + bbox_pred[b_i][2] / 2) * img_size[b_i][0]  # x2
        head_bbox[3] = (bbox_pred[b_i][1] + bbox_pred[b_i][3] / 2) * img_size[b_i][1]  # y2   
        
        os.makedirs(osp.dirname(visualization_save_path), exist_ok=True)
        
        # AUC: area under curve of ROC
        if use_dark_inference:
            pred_x, pred_y = dark_inference(gaze_heatmap_pred[b_i])
        else:
            pred_x, pred_y = argmax_pts(gaze_heatmap_pred[b_i])
        norm_p = [
            pred_x / gaze_heatmap_pred[b_i].shape[-2],
            pred_y / gaze_heatmap_pred[b_i].shape[-1],
        ]
        
        scaled_heatmap = np.array(
            Image.fromarray(gaze_heatmap_pred[b_i]).resize(
                img_size[b_i],
                resample=Image.BILINEAR,
            )
        )
        
        
        flat_index = np.argmax(scaled_heatmap)
        gaze_x, gaze_y = np.unravel_index(flat_index, scaled_heatmap.shape)
        
        image = cv2.imread(image_paths[b_i])
        overlay_bgr = overlay_heatmap_on_image(image, scaled_heatmap, head_bbox, gaze_xy=(gaze_y, gaze_x), inout=inout)
        # overlay_bgr, colored_heat_bgr = overlay_heatmap_on_image(image, scaled_heatmap, head_bbox, alpha=0.1)
        visualization_pred = overlay_bgr
        cv2.imwrite(visualization_save_path, visualization_pred)


def do_test(image_paths, text, model, visualization_dir, use_dark_inference=False):
    model.train(False)
    
    for i in range(0, len(image_paths), 16):
        batch_image_paths = image_paths[i:i+16]
        inference_gaze(batch_image_paths, text, model, visualization_dir, use_dark_inference=use_dark_inference)
    

def main(args):
    cfg = LazyConfig.load(args.config_file)
    image_paths = os.listdir(args.input_path)
    image_paths = [osp.join(args.input_path, image_path) for image_path in image_paths]
    text = args.text
    batch_size = 16
    visualization_dir = args.output_path # osp.join(args.output_path, "visualization")
    if not osp.exists(visualization_dir):
        os.makedirs(visualization_dir)
    model: torch.Module = instantiate(cfg.model)
    model.load_state_dict(torch.load(args.model_weights, weights_only=False)["model"])
    model = model.cuda()
    do_test(image_paths, text, model, visualization_dir, use_dark_inference=args.use_dark_inference)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--config_file", type=str, help="config file")
    parser.add_argument("--input_path", type=str, help="input path")
    parser.add_argument("--text", type=str, help="input text")
    parser.add_argument("--output_path", type=str, help="output path")
    parser.add_argument(
        "--model_weights",
        type=str,
        help="model weights",
    )
    parser.add_argument("--use_dark_inference", action="store_true")
    args = parser.parse_args()
    main(args)