util_vis.py

from pathlib import Path
import cv2
import numpy as np

def vis_tensors_A(l_tensor_or_named_tensor, path_grid, vis_batch_size=4, layout='auto'):
    """Visualize a list of tensors in a grid layout.
    Args:
        l_tensor_or_named_tensor: [tensor | (name, tensor), ..]. each tensor: B,(C,)H,W is in [-1,1] range
        path_grid: Path object for saving the grid visualization
        vis_batch_size: number of samples to visualize
        layout: 'BxI' (batch x images) or 'IxB' (images x batch) or 'auto'
    """
    import torch
    from torchvision.utils import make_grid, save_image
    path_grid = Path(path_grid)
    path_grid.parent.mkdir(parents=0, exist_ok=True)
    # Helper function to unnormalize and prepare images for saving
    def prepare_for_vis(tensor, ):
        if tensor is None:
            return None
        shape = tensor.shape
        assert shape[1]<=3
        if len(shape)==3 or shape[1]==1:
            is_mask = True
        else:  is_mask = False
        if is_mask:
            return tensor.repeat(1, 3, 1, 1).cpu()  # Expand mask to 3 channels
        else:
            return (tensor * 0.5 + 0.5).cpu()  # Unnormalize from [-1, 1] to [0, 1]
    named_tensors = []
    for tensor_or_named_tensor in l_tensor_or_named_tensor:
        if isinstance(tensor_or_named_tensor, tuple):
            name, tensor = tensor_or_named_tensor
        else:
            name = ""
            tensor = tensor_or_named_tensor
        if tensor is not None:
            named_tensors.append((name, prepare_for_vis(tensor.detach()[:vis_batch_size], )))
    # Make sure all tensors have the same spatial dimensions  
    all_shapes = [img.shape[2:] for _, img in named_tensors if img is not None]
    if len(set(all_shapes)) > 1:  # Pad images to match the largest dimensions
        max_h = max(shape[0] for shape in all_shapes)
        max_w = max(shape[1] for shape in all_shapes)
        for i in range(len(named_tensors)):
            name, img = named_tensors[i]
            if img is None:
                continue
            if img.shape[2] == max_h and img.shape[3] == max_w:
                continue
            pad_h = max_h - img.shape[2]
            pad_w = max_w - img.shape[3]
            named_tensors[i] = (name, torch.nn.functional.pad(img, (0, pad_w, 0, pad_h), value=0))
    tensors = []
    for _, (name, tensor) in enumerate(named_tensors):
        tensor = tensor.detach()
        if name:
            for b in range(tensor.shape[0]):
                # Convert tensor to numpy for OpenCV
                img = tensor[b].permute(1, 2, 0).numpy()
                img = (img * 255).astype(np.uint8).copy()  # Make contiguous copy for OpenCV
                # Add text
                cv2.putText(img, name, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 
                           0.7, (0, 0, 0), 2, cv2.LINE_AA)
                cv2.putText(img, name, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 
                           0.7, (255, 255, 255), 1, cv2.LINE_AA)
                img_tensor = torch.from_numpy(img).permute(2, 0, 1) / 255.0 # Convert back to tensor
                tensors.append(img_tensor)
        else:
            for b in range(tensor.shape[0]):
                tensors.append(tensor[b])
    if tensors: # I*B,3,..
        all_images_flat = torch.stack(tensors) # I*B,3,..
        I = len(named_tensors)
        B = vis_batch_size
        if layout == 'auto':
            if B/I > 0.8:
                layout = 'IxB'
            else:
                layout = 'BxI'
        if layout == 'BxI':
            all_images_nonflat = all_images_flat.reshape(I, B, *all_images_flat.shape[1:])
            all_images_nonflat = all_images_nonflat.permute(1, 0, 2, 3, 4)
            all_images_flat = all_images_nonflat.reshape(-1, *all_images_flat.shape[1:])
            nrow = I
        else:  # 'IxB'
            nrow = B
        save_image(make_grid(all_images_flat, nrow=nrow), path_grid)
        print(f"{path_grid=}")

def visualize_landmarks(image, landmarks, save_path):
    """
    Draw landmarks on an image and save the result.
    
    Args:
        image: Input image as a numpy array (H,W,3) with values in [0,255]
        landmarks: Numpy array of shape (136,) or (68,2) containing 68 keypoint coordinates
        save_path: Path where the annotated image should be written
    """
    # Clone the image and ensure uint8 type
    image = image.copy().astype(np.uint8)
    
    # Ensure the image buffer is contiguous
    image = np.ascontiguousarray(image)
    
    # Reshape landmarks into (68,2) if needed
    if landmarks.shape[0] == 136:
        landmarks = landmarks.reshape(68, 2)
    
    # Draw each landmark point
    for (x, y) in landmarks:
        cv2.circle(image, (int(x), int(y)), 2, (0, 255, 0), -1)
    
    # Save the annotated image
    cv2.imwrite(save_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))

def visualize_headPose(img_path, yaw, pitch, roll, save_path):
    """Visualize pose angles on image using arrows
    Args:
        img_path: Path to input image
        yaw: Yaw angle in degrees
        pitch: Pitch angle in degrees
        roll: Roll angle in degrees
        save_path: Path to save visualization
    """
    import matplotlib.pyplot as plt
    img = cv2.imread(str(img_path))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    h, w = img.shape[:2]
    center = (w//2, h//2)
    
    plt.figure(figsize=(10, 10))
    plt.imshow(img)
    
    # Yaw (left-right)
    yaw_rad = np.radians(yaw)
    yaw_end = (center[0] + int(100 * np.sin(yaw_rad)), 
               center[1] - int(100 * np.cos(yaw_rad)))
    plt.arrow(center[0], center[1], yaw_end[0]-center[0], yaw_end[1]-center[1],
              color='r', width=2, head_width=20, label=f'Yaw: {yaw:.1f}°')
    
    # Pitch (up-down)
    pitch_rad = np.radians(pitch)
    pitch_end = (center[0] + int(100 * np.sin(pitch_rad)),
                 center[1] - int(100 * np.cos(pitch_rad)))
    plt.arrow(center[0], center[1], pitch_end[0]-center[0], pitch_end[1]-center[1],
              color='g', width=2, head_width=20, label=f'Pitch: {pitch:.1f}°')
    
    # Roll (tilt)
    roll_rad = np.radians(roll)
    roll_end = (center[0] + int(100 * np.cos(roll_rad)),
                center[1] + int(100 * np.sin(roll_rad)))
    plt.arrow(center[0], center[1], roll_end[0]-center[0], roll_end[1]-center[1],
              color='b', width=2, head_width=20, label=f'Roll: {roll:.1f}°')
    
    plt.legend()
    plt.axis('off')
    
    # Save visualization
    save_path = Path(save_path)
    save_path.parent.mkdir(parents=True, exist_ok=True)
    plt.savefig(save_path, bbox_inches='tight', pad_inches=0)
    plt.close()
    print(f"{save_path=}")