visualize_features.py

# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
import argparse
from collections import defaultdict
from functools import partial
import gc
import math
import os
from PIL import Image
import random
from tqdm import tqdm
from typing import Any, Dict, Iterable, List, Tuple

import cv2
import numpy as np
import torch
from torch import nn
import torch.distributed as dist
from torch.utils.data import DataLoader
import torch.nn.functional as F
from torchvision.utils import make_grid
from torchvision.datasets import ImageFolder

from einops import rearrange

from datasets import load_dataset_builder, load_dataset
from datasets.distributed import split_dataset_by_node

from common import rank_print, load_model, get_standard_transform, collate
from radio.input_conditioner import InputConditioner
from radio.radio_model import RadioOutput

try:
    import wandb
except ImportError:
    wandb = None


LAYER_STATS = dict()


def parse_int_list(string):
    """Parse a comma-separated list of integers and sort them."""
    try:
        return sorted([int(i) for i in string.split(',')])
    except ValueError:
        raise argparse.ArgumentTypeError(f"{string} is not a valid list of integers")


@torch.no_grad()
def main(rank: int = 0, world_size: int = 1):
    '''
    Computes the RankMe (http://arxiv.org/abs/2210.02885) and LiDAR (http://arxiv.org/abs/2312.04000)
    estimates of the rank of the produced embeddings. While RADIO doesn't train in a multi-view setting
    which is an assumption of LiDAR, the metric does integrate an important concept of the invariance of the
    summary features to different view/augmentations of the same image.
    '''

    local_rank = rank % torch.cuda.device_count()
    torch.cuda.set_device(local_rank)
    cv2.setNumThreads(1)

    device = torch.device('cuda', local_rank)
    parser = argparse.ArgumentParser(description='Compute SSL embedding rank estimates')
    parser.add_argument('-v', '--model-version', default='radio_v2',
                        help='Which radio model to load.'
    )
    parser.add_argument('-d', '--dataset', default='detection-datasets/coco',
                        help='The name of the dataset to classify'
    )
    parser.add_argument('--split', default='train',
                        help='The dataset split to use.'
    )
    parser.add_argument('-n', default=128, type=int, help='The number of samples to load')
    parser.add_argument('-r', '--resolution', nargs='+', type=int, default=None,
                        help='The input image resolution.'
                             ' If one value is specified, the shortest dimension is resized to this.'
                             ' If two, the image is center cropped.'
                             ' If not specified, center cropped 378px is used.'
                             ' Default: The RADIO model\'s preferred resolution.'
    )
    parser.add_argument('--max-dim', default=False, action='store_true', help='Resize the max dimension to the specified resolution')
    parser.add_argument('--resize-multiple', type=int, default=None,
                        help='Resize images with dimensions a multiple of this value.'
                             ' This should be equal to the patch size of a ViT (e.g. RADIOv1)'
    )
    parser.add_argument('--batch-size', type=int, default=16,
                        help='The batch size. If the input is variable sized, then this argument becomes a maximum.'
    )
    parser.add_argument('--workers', default=8, type=int, help='Number of loader workers to use')
    parser.add_argument('--vitdet-window-size', default=None, type=int, help='Enable ViTDet at the specific window size')
    parser.add_argument('--output-dir', default='vis_denoise', type=str)
    parser.add_argument('--adaptor-name', default=None, type=str, help='Generate features from a teacher adaptor')
    parser.add_argument('--patch-size', default=None, type=int, help='The model patch size')
    parser.add_argument('--torchhub-repo',
                        help="Path to the Torchhub repo", default="NVlabs/RADIO"
    )
    parser.add_argument('--intermediates',
                        type=parse_int_list,
                        help='Visualize intermediate layers, specified as a comma-separated list of integers (e.g., 1,2,3,4,5)'
    )
    parser.add_argument('--intermediate-aggregation',
                        default='sparse',
                        type=str,
                        help='How to aggregate intermediate layers',
                        choices=['sparse', 'dense'])
    parser.add_argument('--interpolation', default='bilinear', type=str, help='Interpolation mode')
    parser.add_argument('--skip', default=0, type=int, help='Skip the first N components')
    parser.add_argument('--neck', default=None, type=str, help='Generate features from specified neck')
    parser.add_argument('--animate-radio1d', action='store_true', help='Create animated PNGs sweeping radio1d_size')
    parser.add_argument('--radio1d-start', default=1, type=int, help='Starting radio1d_size for animation')
    parser.add_argument('--radio1d-end', default=512, type=int, help='Ending radio1d_size for animation')
    parser.add_argument('--radio1d-step', default=32, type=int, help='Step size for radio1d_size sweep')
    parser.add_argument('--animation-duration', default=100, type=int, help='Duration of each frame in ms')

    args, _ = parser.parse_known_args()

    torch.manual_seed(42 + rank)
    np.random.seed(42 + rank)
    random.seed(42 + rank)

    rank_print(f'Loading model: "{args.model_version}", ViTDet: {args.vitdet_window_size}, Adaptor: "{args.adaptor_name}", Resolution: {args.resolution}, Max: {args.max_dim}...')
    if os.path.isdir(args.model_version):
        for chk_name in ['last_release_half.pth.tar', 'last_release.pth.tar', 'last.pth.tar']:
            chk_path = os.path.join(args.model_version, chk_name)
            if os.path.exists(chk_path):
                args.model_version = chk_path
                print(f'Using "{chk_path}" as model version.')
                break

    model, preprocessor, info = load_model(args.model_version, vitdet_window_size=args.vitdet_window_size, adaptor_names=args.adaptor_name,
                                           torchhub_repo=args.torchhub_repo, neck_name=args.neck)
    model.to(device=device).eval()
    if isinstance(preprocessor, nn.Module):
        preprocessor.to(device).eval()
    rank_print('Done')

    rank_print('Loading dataset...')

    if args.resolution is None:
        args.resolution = (model.preferred_resolution.height, model.preferred_resolution.width)

    # if len(args.resolution) == 1:
    #     args.batch_size = 1
    args.batch_size = 1

    patch_size = getattr(model, 'patch_size', None) or args.patch_size

    if args.resize_multiple is None:
        args.resize_multiple = getattr(model, 'min_resolution_step', patch_size)

    transform = get_standard_transform(args.resolution, args.resize_multiple, max_dim=args.max_dim,
                                       pad_mean=preprocessor.norm_mean if isinstance(preprocessor, InputConditioner) else None,)

    if not os.path.isdir(args.dataset):
        try:
            ds_builder = load_dataset_builder(args.dataset, trust_remote_code=True)
        except:
            ds_builder = load_dataset_builder("imagefolder", data_dir=args.dataset)
        ds_builder.download_and_prepare()
        ds_builder = load_dataset_builder(args.dataset, trust_remote_code=True)
        dataset = ds_builder.as_dataset(split=args.split)
        dataset = dataset.to_iterable_dataset(num_shards=world_size * max(1, args.workers))
        dataset = split_dataset_by_node(dataset, rank=rank, world_size=world_size)
        dataset = dataset.map(lambda ex: dict(image=transform(ex['image']), label=torch.zeros(1, dtype=torch.int64)))
        rank_print(f'Description: {ds_builder.info.description}')
    else:
        dataset = ImageFolder(args.dataset, transform=transform)
        dataset.samples.sort(key=lambda s: s[0])

    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False,
                        num_workers=args.workers, collate_fn=partial(collate, group=False),
                        pin_memory=args.workers > 0,
                        drop_last=False,
    )
    rank_print('Done')

    dirs = dict(
        orig=os.path.join(args.output_dir, 'orig'),
        viz=os.path.join(args.output_dir, 'viz'),
        sbs=os.path.join(args.output_dir, 'sbs'),
        grid=os.path.join(args.output_dir, 'grid'),
    )

    for d in dirs.values():
        os.makedirs(d, exist_ok=True)

    if args.animate_radio1d:
        # Create animation directories
        anim_dirs = dict(
            viz=os.path.join(args.output_dir, 'anim_viz'),
            sbs=os.path.join(args.output_dir, 'anim_sbs'),
            grid=os.path.join(args.output_dir, 'anim_grid'),
        )
        for d in anim_dirs.values():
            os.makedirs(d, exist_ok=True)

        process_with_animation(
            loader, model, preprocessor, device, args, dirs, anim_dirs
        )
    else:
        # Original behavior
        ctr = 0

        for batches in loader:
            if ctr >= args.n:
                break

            for images, _ in batches:
                images = images.to(device=device, non_blocking=True)

                all_feat = []
                with torch.autocast(device.type, dtype=torch.bfloat16):
                    p_images = preprocessor(images)

                    if args.intermediates:
                        outputs = model.forward_intermediates(
                            p_images,
                            indices=args.intermediates,
                            return_prefix_tokens=False,
                            norm=False,
                            stop_early=True,
                            output_fmt='NCHW',
                            intermediates_only=True,
                            aggregation=args.intermediate_aggregation,
                            norm_alpha_scheme="none",
                        )
                        assert args.adaptor_name is None
                        all_feat = outputs
                    else:
                        kwargs = {}
                        if args.neck == "encoder":
                            all_tokens = math.prod(p_images.shape[-2:]) // (4 * model.patch_size ** 2)
                            # Assuming 1D encoder with 2x2 downsampling.
                            kwargs['num_tokens'] = all_tokens
                            output = model(p_images, feature_fmt='NLC', **kwargs)
                            features = output.features.reshape(
                                output.features.shape[0],
                                p_images.shape[-2] // model.patch_size // 2,
                                p_images.shape[-1] // model.patch_size // 2,
                                output.features.shape[2]).permute(0, 3, 1, 2)
                            output = RadioOutput(output.summary, features)
                        else:
                            output = model(p_images, feature_fmt='NCHW', **kwargs)
                        if args.adaptor_name:
                            all_feat = [
                                output['backbone'].features,
                                output[args.adaptor_name].features,
                            ]
                        else:
                            all_feat = [output[1]]

                all_feat = [rearrange(f, 'b c h w -> b h w c').float() for f in all_feat]

                # b m h w c
                all_feat = list(zip(*all_feat))

                for i, feats in enumerate(all_feat):
                    colored = []
                    for features in feats:
                        color = get_pca_map(features, images.shape[-2:], interpolation=args.interpolation, skip_components=args.skip)
                        colored.append(color)

                    orig = cv2.cvtColor(images[i].permute(1, 2, 0).cpu().numpy(), cv2.COLOR_RGB2BGR)

                    cv2.imwrite(f'{dirs["orig"]}/vis_{ctr}.jpg', orig * 255)

                    if args.intermediates:
                        annotations = [f'layer_{i}_{args.intermediate_aggregation}' for i in args.intermediates]
                    else:
                        annotations = ["backbone"]
                        if args.adaptor_name is not None:
                            annotations.append(args.adaptor_name)
                    for annotation, img in zip(annotations, colored):
                        rdir = f'{dirs["viz"]}/{annotation}'
                        os.makedirs(rdir, exist_ok=True)
                        cv2.imwrite(f'{rdir}/vis_{ctr}.jpg', img * 255)

                    grid_image = create_image_grid_with_annotations(
                        images=[[orig] + colored],
                        annotations=["Original"] + annotations)
                    if args.intermediates:
                        grid_filename = f'{dirs["grid"]}/vis_{ctr}_{args.intermediate_aggregation}.jpg'
                    else:
                        grid_filename = f'{dirs["grid"]}/vis_{ctr}.jpg'
                    cv2.imwrite(grid_filename, grid_image * 255)

                    sbs_images = [orig] + list(colored)
                    op = np.concatenate(sbs_images, axis=1) * 255

                    cv2.imwrite(f'{dirs["sbs"]}/vis_{ctr}.jpg', op)

                    ctr += 1


def process_with_animation(loader, model, preprocessor, device, args, dirs, anim_dirs):
    """
    Process images with radio1d_size animation.
    First pass: Extract PCA stats with max radio1d_size.
    Second pass: Generate frames for all radio1d_size values using those stats.
    """
    rank_print(f'Processing with radio1d animation: {args.radio1d_start} to {args.radio1d_end} step {args.radio1d_step}')

    # Collect all images first
    all_images = []
    ctr = 0
    for batches in loader:
        if ctr >= args.n:
            break
        for images, _ in batches:
            images = images.to(device=device, non_blocking=True)
            all_images.append(images)
            ctr += len(images)

    rank_print(f'Collected {len(all_images)} batches')

    # First pass: Get PCA stats from the highest radio1d_size
    rank_print(f'First pass: Computing PCA stats with radio1d_size={args.radio1d_end}')
    pca_stats_list = []

    for images in tqdm(all_images, desc='Computing PCA stats'):
        with torch.autocast(device.type, dtype=torch.bfloat16):
            p_images = preprocessor(images)

            if args.intermediates:
                outputs = model.forward_intermediates(
                    p_images,
                    indices=args.intermediates,
                    return_prefix_tokens=False,
                    norm=False,
                    stop_early=True,
                    output_fmt='NCHW',
                    intermediates_only=True,
                    aggregation=args.intermediate_aggregation,
                    norm_alpha_scheme="none",
                )
                all_feat = outputs
            else:
                kwargs = {}
                if args.neck:
                    kwargs['num_tokens'] = args.radio1d_end
                output = model(p_images, feature_fmt='NCHW', **kwargs)
                if args.adaptor_name:
                    all_feat = [
                        output['backbone'].features,
                        output[args.adaptor_name].features,
                    ]
                else:
                    all_feat = [output[1]]

        all_feat = [rearrange(f, 'b c h w -> b h w c').float() for f in all_feat]
        all_feat = list(zip(*all_feat))

        # Get PCA stats for each image
        for feats in all_feat:
            img_pca_stats = []
            for features in feats:
                _, pca_stats = get_pca_map(
                    features,
                    images.shape[-2:],
                    interpolation=args.interpolation,
                    return_pca_stats=True,
                    skip_components=args.skip
                )
                img_pca_stats.append(pca_stats)
            pca_stats_list.append(img_pca_stats)

    # Second pass: Generate frames for all radio1d_size values
    rank_print(f'Second pass: Generating frames for all radio1d_size values')
    radio1d_sizes = list(range(args.radio1d_start, args.radio1d_end + 1, args.radio1d_step))

    for img_idx, images in enumerate(tqdm(all_images, desc='Processing images')):
        # Storage for animation frames
        viz_frames = defaultdict(list)  # annotation -> list of frames
        sbs_frames = []
        grid_frames = []

        orig = cv2.cvtColor(images[0].permute(1, 2, 0).cpu().numpy(), cv2.COLOR_RGB2BGR)

        # Save original image
        cv2.imwrite(f'{dirs["orig"]}/vis_{img_idx}.jpg', orig * 255)

        for radio1d_size in tqdm(radio1d_sizes, desc=f'radio1d_size sweep (img {img_idx})', leave=False):
            with torch.autocast(device.type, dtype=torch.bfloat16):
                p_images = preprocessor(images)

                if args.intermediates:
                    outputs = model.forward_intermediates(
                        p_images,
                        indices=args.intermediates,
                        return_prefix_tokens=False,
                        norm=False,
                        stop_early=True,
                        output_fmt='NCHW',
                        intermediates_only=True,
                        aggregation=args.intermediate_aggregation,
                        norm_alpha_scheme="none",
                    )
                    all_feat = outputs
                else:
                    kwargs = {}
                    if args.neck:
                        kwargs['num_tokens'] = radio1d_size
                    output = model(p_images, feature_fmt='NCHW', **kwargs)
                    if args.adaptor_name:
                        all_feat = [
                            output['backbone'].features,
                            output[args.adaptor_name].features,
                        ]
                    else:
                        all_feat = [output[1]]

            all_feat = [rearrange(f, 'b c h w -> b h w c').float() for f in all_feat]
            all_feat = list(zip(*all_feat))

            # Render each feature map with its precomputed PCA basis (computed
            # once at radio1d_end). Reusing the basis keeps colors stable across
            # frames of the radio1d_size sweep.
            feats = all_feat[0]  # First image in batch
            colored = []
            for feat_idx, features in enumerate(feats):
                color = get_pca_map(
                    features,
                    images.shape[-2:],
                    interpolation=args.interpolation,
                    pca_stats=pca_stats_list[img_idx][feat_idx],
                    skip_components=args.skip,
                )
                colored.append(color)

            if args.intermediates:
                annotations = [f'layer_{i}_{args.intermediate_aggregation}' for i in args.intermediates]
            else:
                annotations = ["backbone"]
                if args.adaptor_name is not None:
                    annotations.append(args.adaptor_name)

            for annotation, img in zip(annotations, colored):
                frame = (img * 255).astype(np.uint8)
                frame = add_text_overlay(frame, f"radio1d_size={radio1d_size}")
                viz_frames[annotation].append(frame)

            sbs_images = [orig] + list(colored)
            sbs_frame = np.concatenate(sbs_images, axis=1) * 255
            sbs_frame = sbs_frame.astype(np.uint8)
            sbs_frame = add_text_overlay(sbs_frame, f"radio1d_size={radio1d_size}")
            sbs_frames.append(sbs_frame)

            # Grid frame
            grid_image = create_image_grid_with_annotations(
                images=[[orig] + colored],
                annotations=["Original"] + annotations
            )
            grid_image = (grid_image * 255).astype(np.uint8)
            grid_image = add_text_overlay(grid_image, f"radio1d_size={radio1d_size}", position='bottom')
            grid_frames.append(grid_image)

        # Save animated PNGs
        rank_print(f'Saving animated PNGs for image {img_idx}')

        # Save individual viz animations
        for annotation, frames in viz_frames.items():
            rdir = f'{anim_dirs["viz"]}/{annotation}'
            os.makedirs(rdir, exist_ok=True)
            save_animated_image(frames, f'{rdir}/vis_{img_idx}.png', duration=args.animation_duration)

        # Save side-by-side animation
        save_animated_image(sbs_frames, f'{anim_dirs["sbs"]}/vis_{img_idx}.gif', duration=args.animation_duration, format='GIF')

        # Save grid animation
        save_animated_image(grid_frames, f'{anim_dirs["grid"]}/vis_{img_idx}.png', duration=args.animation_duration)

    rank_print('Animation processing complete!')


def add_text_overlay(image, text, position='top', font=cv2.FONT_HERSHEY_SIMPLEX,
                     font_scale=1.0, font_color=(255, 255, 255), thickness=2,
                     bg_color=(0, 0, 0), padding=10):
    """
    Add text overlay to an image.

    Args:
        image: numpy array (H, W, C) in BGR format with values 0-255
        text: Text to display
        position: 'top' or 'bottom' for text placement
        font: OpenCV font type
        font_scale: Font scale
        font_color: Text color (BGR)
        thickness: Text thickness
        bg_color: Background color for text (BGR)
        padding: Padding around text

    Returns:
        Image with text overlay
    """
    image = image.copy()

    # Get text size
    (text_width, text_height), baseline = cv2.getTextSize(text, font, font_scale, thickness)

    # Calculate position
    img_height, img_width = image.shape[:2]

    if position == 'top':
        text_x = padding
        text_y = padding + text_height
        bg_y1 = 0
        bg_y2 = text_height + 2 * padding
    else:  # bottom
        text_x = padding
        text_y = img_height - padding - baseline
        bg_y1 = img_height - text_height - 2 * padding - baseline
        bg_y2 = img_height

    bg_x1 = 0
    bg_x2 = text_width + 2 * padding

    # Draw semi-transparent background
    overlay = image.copy()
    cv2.rectangle(overlay, (bg_x1, bg_y1), (bg_x2, bg_y2), bg_color, -1)
    alpha = 0.7
    image = cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0)

    # Draw text
    cv2.putText(image, text, (text_x, text_y), font, font_scale, font_color, thickness, cv2.LINE_AA)

    return image


def save_animated_image(frames, output_path, duration=100, format='PNG'):
    """
    Save a list of numpy arrays as an animated image.

    Args:
        frames: List of numpy arrays (H, W, C) in BGR format with values 0-255
        output_path: Path to save the APNG file
        duration: Duration of each frame in milliseconds
    """
    # Convert BGR frames to RGB PIL Images
    pil_frames = []
    for frame in frames:
        # Convert BGR to RGB
        frame_rgb = cv2.cvtColor(frame.astype(np.uint8), cv2.COLOR_BGR2RGB)
        pil_frames.append(Image.fromarray(frame_rgb))

    # Save as animated PNG
    pil_frames[0].save(
        output_path,
        save_all=True,
        append_images=pil_frames[1:],
        duration=duration,
        loop=0,  # 0 means infinite loop
        format=format
    )


def get_robust_pca(features: torch.Tensor, m: float = 2, remove_first_component=False, skip: int = 0):
    # features: (N, C)
    # m: a hyperparam controlling how many std dev outside for outliers
    assert len(features.shape) == 2, "features should be (N, C)"
    reduction_mat = torch.pca_lowrank(features, q=3 + skip, niter=20)[2]
    reduction_mat = reduction_mat[:, skip:]
    colors = features @ reduction_mat
    if remove_first_component:
        colors_min = colors.min(dim=0).values
        colors_max = colors.max(dim=0).values
        tmp_colors = (colors - colors_min) / (colors_max - colors_min)
        fg_mask = tmp_colors[..., 0] < 0.2
        reduction_mat = torch.pca_lowrank(features[fg_mask], q=3, niter=20)[2]
        colors = features @ reduction_mat
    else:
        fg_mask = torch.ones_like(colors[:, 0]).bool()
    d = torch.abs(colors[fg_mask] - torch.median(colors[fg_mask], dim=0).values)
    mdev = torch.median(d, dim=0).values
    s = d / mdev
    try:
        rins = colors[fg_mask][s[:, 0] < m, 0]
        gins = colors[fg_mask][s[:, 1] < m, 1]
        bins = colors[fg_mask][s[:, 2] < m, 2]
        rgb_min = torch.tensor([rins.min(), gins.min(), bins.min()])
        rgb_max = torch.tensor([rins.max(), gins.max(), bins.max()])
    except:
        rins = colors
        gins = colors
        bins = colors
        rgb_min = torch.tensor([rins.min(), gins.min(), bins.min()])
        rgb_max = torch.tensor([rins.max(), gins.max(), bins.max()])

    return reduction_mat, rgb_min.to(reduction_mat), rgb_max.to(reduction_mat)


def get_pca_map(
    feature_map: torch.Tensor,
    img_size,
    interpolation="bicubic",
    return_pca_stats=False,
    pca_stats=None,
    skip_components: int = 0,
):
    """
    feature_map: (1, h, w, C) is the feature map of a single image.
    """
    if feature_map.shape[0] != 1:
        # make it (1, h, w, C)
        feature_map = feature_map[None]
    if pca_stats is None:
        reduct_mat, color_min, color_max = get_robust_pca(
            feature_map.reshape(-1, feature_map.shape[-1]), skip=skip_components,
        )
    else:
        reduct_mat, color_min, color_max = pca_stats
    pca_color = feature_map @ reduct_mat
    pca_color = (pca_color - color_min) / (color_max - color_min)
    pca_color = pca_color.clamp(0, 1)
    pca_color = F.interpolate(
        pca_color.permute(0, 3, 1, 2),
        size=img_size,
        mode=interpolation,
    ).permute(0, 2, 3, 1)
    pca_color = pca_color.cpu().numpy().squeeze(0)
    if return_pca_stats:
        return pca_color, (reduct_mat, color_min, color_max)
    return pca_color


def get_scale_map(
    scalar_map: torch.Tensor,
    img_size,
    interpolation="nearest",
):
    """
    scalar_map: (1, h, w, C) is the feature map of a single image.
    """
    if scalar_map.shape[0] != 1:
        scalar_map = scalar_map[None]
    scalar_map = (scalar_map - scalar_map.min()) / (
        scalar_map.max() - scalar_map.min() + 1e-6
    )
    scalar_map = F.interpolate(
        scalar_map.permute(0, 3, 1, 2),
        size=img_size,
        mode=interpolation,
    ).permute(0, 2, 3, 1)
    # cmap = plt.get_cmap("viridis")
    # scalar_map = cmap(scalar_map)[..., :3]
    # make it 3 channels
    scalar_map = torch.cat([scalar_map] * 3, dim=-1)
    scalar_map = scalar_map.cpu().numpy().squeeze(0)
    return scalar_map


def create_image_grid_with_annotations(images, annotations, font=cv2.FONT_HERSHEY_SIMPLEX, font_scale=1, font_color=(255, 255, 255), thickness=2):
    """
    Arrange a nested list of MxN numpy arrays into a single image grid with column annotations.

    Args:
    images (list of list of np.array): Nested list containing the image arrays.
    annotations (list of str): List of annotations for each column.
    font (int): Font type for the annotations.
    font_scale (float): Font scale for the annotations.
    font_color (tuple): Font color for the annotations.
    thickness (int): Thickness of the font.

    Returns:
    np.array: The resulting image grid with annotations.
    """
    # Check if the input is valid
    if not images or not isinstance(images, list) or not all(isinstance(row, list) for row in images):
        raise ValueError("Input must be a nested list of numpy arrays.")

    # Determine the number of rows and columns
    num_rows = len(images)
    num_cols = len(images[0])

    # Check if all rows have the same number of columns
    if not all(len(row) == num_cols for row in images):
        raise ValueError("All rows must have the same number of columns.")

    # Get the dimensions of each image
    first_image_shape = images[0][0].shape
    img_height, img_width = first_image_shape[:2]

    # Check if all images have the same dimensions
    for row in images:
        for img in row:
            if img.shape[:2] != (img_height, img_width):
                raise ValueError("All images must have the same dimensions.")

    # Determine the number of channels
    num_channels = first_image_shape[2] if len(first_image_shape) == 3 else 1

    # Create an empty array for the resulting grid
    annotation_height = 50  # Height reserved for annotations
    grid_height = num_rows * img_height + annotation_height
    grid_width = num_cols * img_width

    if num_channels == 1:
        grid = np.zeros((grid_height, grid_width), dtype=images[0][0].dtype)
    else:
        grid = np.zeros((grid_height, grid_width, num_channels), dtype=images[0][0].dtype)

    # Add annotations at the top of each column
    for col_idx, annotation in enumerate(annotations):
        text_size = cv2.getTextSize(annotation, font, font_scale, thickness)[0]
        text_x = col_idx * img_width + (img_width - text_size[0]) // 2
        text_y = (annotation_height + text_size[1]) // 2
        cv2.putText(grid, annotation, (text_x, text_y), font, font_scale, font_color, thickness)

    # Populate the grid with the images
    for row_idx, row in enumerate(images):
        for col_idx, img in enumerate(row):
            start_y = annotation_height + row_idx * img_height
            start_x = col_idx * img_width
            grid[start_y:start_y+img_height, start_x:start_x+img_width] = img

    return grid


if __name__ == '__main__':
    rank = 0
    world_size = 1

    # if 'WORLD_SIZE' in os.environ:
    #     dist.init_process_group(backend='nccl')
    #     rank = dist.get_rank()
    #     world_size = dist.get_world_size()

    main(rank, world_size)