data/utils/image_utils.py

#!/usr/bin/env python3
"""
Image processing utilities for dataset loading.
Common functions shared across different datasets (AC-One, RobotWin, ALOHA, etc.)
"""

import numpy as np
import cv2
import torch
from PIL import Image
from typing import Tuple, List
import random
from decord import VideoReader, cpu


def resize_with_padding(frame: np.ndarray, target_size: Tuple[int, int]) -> np.ndarray:
    """
    Resize image with aspect ratio preservation and padding to target size.
    
    This function ensures no image distortion by:
    1. Calculating the minimum scale ratio to fit the image within target size
    2. Resizing the image with this ratio to preserve aspect ratio
    3. Padding with black borders to reach exact target size
    4. Centering the resized image within the padded frame
    
    Args:
        frame: Input image [H, W, C]
        target_size: Target size (height, width)
        
    Returns:
        Processed image [target_height, target_width, C]
        
    Example:
        >>> frame = np.random.randint(0, 255, (720, 640, 3), dtype=np.uint8)
        >>> resized = resize_with_padding(frame, (384, 320))
        >>> print(resized.shape)  # (384, 320, 3)
    """
    target_height, target_width = target_size
    original_height, original_width = frame.shape[:2]
    
    # Calculate scaling ratio, use the smaller ratio to ensure image fits completely
    scale_height = target_height / original_height
    scale_width = target_width / original_width
    scale = min(scale_height, scale_width)
    
    # Calculate new dimensions after scaling
    new_height = int(original_height * scale)
    new_width = int(original_width * scale)
    
    # Resize with aspect ratio preservation
    resized_frame = cv2.resize(frame, (new_width, new_height))
    
    # Create black background with target size
    padded_frame = np.zeros((target_height, target_width, frame.shape[2]), dtype=frame.dtype)
    
    # Calculate center placement position
    y_offset = (target_height - new_height) // 2
    x_offset = (target_width - new_width) // 2
    
    # Place resized image at center
    padded_frame[y_offset:y_offset + new_height, x_offset:x_offset + new_width] = resized_frame
    
    return padded_frame


def load_video_frames(video_path: str, frame_indices: List[int], target_size: Tuple[int, int] = None) -> torch.Tensor:
    """
    Load random frames from a video using decord, with optional aspect-ratio-preserving resize and padding.

    - Decoder: decord.VideoReader
    - Access: get_batch(frame_indices) to fetch arbitrary frames in one call
    - Color: decord returns frames in RGB with HWC layout

    Args:
        video_path: Path to the video file.
        frame_indices: Frame indices to read (can be unordered and may repeat).
        target_size: Optional target size (height, width). If provided, resize with aspect ratio preserved and
            center-pad with black borders to the exact target size.

    Returns:
        torch.Tensor of shape [T, C, H, W] with values in [0, 1].
    """

    # Open video (CPU decoding)
    vr = VideoReader(video_path, ctx=cpu(0), num_threads=4)
    total_frames = len(vr)

    if any(idx < 0 or idx >= total_frames for idx in frame_indices):
        raise ValueError(
            f"Some frame indices are out of bounds for video {video_path} (total frames: {total_frames})"
        )

    # Batch-read frames; returns decord NDArray with shape [T, H, W, 3] in RGB
    batch = vr.get_batch(frame_indices)
    frames_np = batch.asnumpy()  # uint8, [T, H, W, C]

    # Optional: aspect-ratio-preserving resize with padding to the target size
    if target_size is not None:
        th, tw = target_size
        _, h, w, _ = frames_np.shape
        if (h, w) != (th, tw):
            # Apply resize_with_padding to each frame (OpenCV implementation; no distortion)
            resized = [resize_with_padding(frames_np[i], target_size) for i in range(frames_np.shape[0])]
            frames_np = np.stack(resized, axis=0)

    # Convert to [T, C, H, W] and normalize to [0, 1]
    video_tensor = torch.from_numpy(frames_np).permute(0, 3, 1, 2).float() / 255.0
    return video_tensor


def load_first_frame(video_path: str, frame_idx: int, target_size: Tuple[int, int]) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Load a single video frame in both resized and original resolution using decord only.
    
    Args:
        video_path: Path to video file
        frame_idx: Frame index to load
        target_size: Target size (height, width) for resizing with padding
        
    Returns:
        Tuple of (resized_frame, original_frame):
        - resized_frame: Frame tensor [C, H, W] in range [0, 1] with padding applied
        - original_frame: Frame tensor [C, H, W] in range [0, 1] with original resolution
    """
    vr = VideoReader(video_path, ctx=cpu(0), num_threads=1)
    total_frames = len(vr)
    if frame_idx >= total_frames:
        raise ValueError(f"Frame index {frame_idx} out of bounds for video {video_path} (total frames: {total_frames})")
    frame_rgb = vr[frame_idx].asnumpy()  # [H, W, 3], RGB uint8
    frame_original = torch.from_numpy(frame_rgb).permute(2, 0, 1).float() / 255.0
    frame_resized_np = resize_with_padding(frame_rgb, target_size)
    frame_resized = torch.from_numpy(frame_resized_np).permute(2, 0, 1).float() / 255.0
    return frame_resized, frame_original


def get_video_frame_count(video_path: str) -> int:
    """Get total frame count of a video using decord only."""
    vr = VideoReader(video_path, ctx=cpu(0), num_threads=1)
    return len(vr)


def tensor_to_pil(tensor: torch.Tensor) -> Image.Image:
    """
    Convert tensor [C, H, W] to PIL Image.
    
    Args:
        tensor: Input tensor in format [C, H, W]
        
    Returns:
        PIL Image in RGB mode
    """
    # Convert from [C, H, W] to [H, W, C] and to numpy
    if tensor.shape[0] == 3:  # RGB
        image_np = tensor.permute(1, 2, 0).numpy()
        # Convert from [0, 1] to [0, 255] if needed
        if image_np.max() <= 1.0:
            image_np = (image_np * 255).astype(np.uint8)
        else:
            image_np = image_np.astype(np.uint8)
        return Image.fromarray(image_np, mode='RGB')
    else:
        raise ValueError(f"Unsupported tensor shape: {tensor.shape}")


def apply_image_augmentation(frame: np.ndarray, 
                           brightness_prob: float = 0.5,
                           brightness_range: Tuple[float, float] = (0.8, 1.2),
                           flip_prob: float = 0.3) -> np.ndarray:
    """
    Apply common image augmentations to a frame.
    
    Args:
        frame: Input image [H, W, C]
        brightness_prob: Probability of applying brightness adjustment
        brightness_range: Range of brightness factors (min, max)
        flip_prob: Probability of applying horizontal flip
        
    Returns:
        Augmented image [H, W, C]
    """
    # Random brightness adjustment
    if random.random() < brightness_prob:
        brightness_factor = random.uniform(*brightness_range)
        frame = np.clip(frame * brightness_factor, 0, 255)
    
    # Random horizontal flip
    if random.random() < flip_prob:
        frame = np.fliplr(frame)
    
    return frame


# Test functions for validation
def test_resize_with_padding():
    """Test the resize_with_padding function with visual output."""
    import os
    
    print("=== Testing resize_with_padding ===")
    
    # Create test image with WHITE background and colored patterns
    test_frame = np.full((720, 640, 3), 255, dtype=np.uint8)  # White background
    test_frame[100:200, 100:200] = [255, 0, 0]    # Red square
    test_frame[300:400, 300:400] = [0, 255, 0]    # Green square  
    test_frame[500:600, 450:550] = [0, 0, 255]    # Blue square
    test_frame[50:670, 50:70] = [255, 255, 0]     # Yellow left border
    test_frame[50:670, 570:590] = [255, 0, 255]   # Magenta right border
    test_frame[50:70, 50:590] = [0, 255, 255]     # Cyan top border
    test_frame[650:670, 50:590] = [128, 128, 128] # Gray bottom border
    
    # Test resize
    target_size = (384, 320)
    resized_frame = resize_with_padding(test_frame, target_size)
    
    print(f"Original image size: {test_frame.shape}")
    print(f"Target size: {target_size}")
    print(f"Resized image size: {resized_frame.shape}")
    
    # Calculate expected values
    original_h, original_w = 720, 640
    target_h, target_w = 384, 320
    scale = min(target_h / original_h, target_w / original_w)  # 0.5
    expected_new_h = int(original_h * scale)  # 360
    expected_new_w = int(original_w * scale)  # 320
    
    print(f"Scaling factor: {scale}")
    print(f"Expected scaled size: {expected_new_h}x{expected_new_w}")
    print(f"Padding: top/bottom {(target_h - expected_new_h) // 2} pixels each")
    print(f"Padding: left/right {(target_w - expected_new_w) // 2} pixels each")
    
    # Verify padding
    y_offset = (target_h - expected_new_h) // 2
    x_offset = (target_w - expected_new_w) // 2
    
    if y_offset > 0:
        top_black = np.all(resized_frame[:y_offset, :, :] == 0)
        bottom_black = np.all(resized_frame[y_offset + expected_new_h:, :, :] == 0)
        print(f"Top padding check: {'✓' if top_black else '✗'}")
        print(f"Bottom padding check: {'✓' if bottom_black else '✗'}")
    
    if x_offset > 0:
        left_black = np.all(resized_frame[:, :x_offset, :] == 0)
        right_black = np.all(resized_frame[:, x_offset + expected_new_w:, :] == 0)
        print(f"Left padding check: {'✓' if left_black else '✗'}")
        print(f"Right padding check: {'✓' if right_black else '✗'}")
    
    # Save images if PIL is available
    try:
        from PIL import Image
        
        # Create output directory
        output_dir = "image_utils_test_output"
        os.makedirs(output_dir, exist_ok=True)
        
        # Save original image
        original_pil = Image.fromarray(test_frame)
        original_pil.save(os.path.join(output_dir, "original_720x640.png"))
        print(f"✓ Saved original image: {output_dir}/original_720x640.png")
        
        # Save resized image
        resized_pil = Image.fromarray(resized_frame)
        resized_pil.save(os.path.join(output_dir, "resized_384x320.png"))
        print(f"✓ Saved resized image: {output_dir}/resized_384x320.png")
        
        print("\n=== Test Complete ===")
        print("✓ Aspect ratio preserved")
        print("✓ No image distortion") 
        print("✓ Proper padding applied")
        print("✓ Images saved successfully")
        
    except ImportError as e:
        print(f"Could not save images due to missing dependencies: {e}")
    except Exception as e:
        print(f"Error saving images: {e}")


if __name__ == "__main__":
    test_resize_with_padding()