fourm/data/modality_transforms.py

# Copyright 2024 EPFL and Apple Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import gzip
import json
import random
from pathlib import Path
from typing import Optional, Tuple, List, Dict
from abc import ABC, abstractmethod
import decord
from PIL import Image
import cv2

import albumentations as A
import numpy as np
import torch
import torchvision.transforms.functional as TF
import torchvision.transforms as T
from einops import rearrange, repeat, reduce

from fourm.utils import to_2tuple
from fourm.utils.data_constants import (IMAGENET_DEFAULT_MEAN,
                                  IMAGENET_DEFAULT_STD, IMAGENET_INCEPTION_MEAN,
                                  IMAGENET_SURFACE_NORMAL_STD, IMAGENET_SURFACE_NORMAL_MEAN,
                                  IMAGENET_INCEPTION_STD, SEG_IGNORE_INDEX, PAD_MASK_VALUE)
from fourm.data.vjepa.VIDTransforms import make_transforms
import itertools
from itertools import accumulate

# The @-symbol is used to specify the resolution of a modality. Syntax: modality@resolution
def get_transform_key(mod_name):
    return mod_name.split('@')[0]

def get_transform_resolution(mod_name, default_resolution, to_tuple=True):
    res = int(mod_name.split('@')[1]) if '@' in mod_name else default_resolution
    return to_2tuple(res) if to_tuple else res

def get_transform(mod_name, transforms_dict):
    return transforms_dict.get(get_transform_key(mod_name), IdentityTransform())

def get_pil_resample_mode(resample_mode: str):
    """
    Returns the PIL resampling mode for the given resample mode string.

    Args:
        resample_mode: Resampling mode string
    """
    if resample_mode is None:
        return None
    elif resample_mode == "bilinear":
        return Image.Resampling.BILINEAR if hasattr(Image, 'Resampling') else Image.BILINEAR
    elif resample_mode == "bicubic":
        return Image.Resampling.BICUBIC if hasattr(Image, 'Resampling') else Image.BICUBIC
    elif resample_mode == "nearest":
        return Image.Resampling.NEAREST if hasattr(Image, 'Resampling') else Image.NEAREST
    else:
        raise ValueError(f"Resample mode {resample_mode} is not supported.")

class UnifiedDataTransform(object):
    def __init__(self, transforms_dict, image_augmenter, resample_mode: str = None, add_sizes: bool = False, **kwargs):
        """Unified data augmentation for FourM

        Args:
            transforms_dict (dict): Dict of transforms for each modality
            image_augmenter (AbstractImageAugmenter): Image augmenter
            resample_mode (str, optional): Resampling mode for PIL images (default: None -> uses default resampling mode for data type)
                One out of ["bilinear", "bicubic", "nearest", None].
            add_sizes (bool, optional): Whether to add crop coordinates and original size to the output dict
        """

        self.transforms_dict = transforms_dict
        self.image_augmenter = image_augmenter
        self.resample_mode = resample_mode
        self.add_sizes = add_sizes
        self.resize_video_transform = None

    def unified_image_augment(self, mod_dict, crop_settings):
        """Apply the image augmenter to all modalities where it is applicable

        Args:
            mod_dict (dict): Dict of modalities
            crop_settings (dict): Crop settings

        Returns:
            dict: Transformed dict of modalities
        """

        crop_coords, flip, orig_size, target_size, rand_aug_idx = self.image_augmenter(mod_dict, crop_settings)

        mod_dict = {
            k: self.transforms_dict[get_transform_key(k)].image_augment(
                v, crop_coords=crop_coords, flip=flip, orig_size=orig_size, 
                target_size=get_transform_resolution(k, target_size), rand_aug_idx=rand_aug_idx,
                resample_mode=self.resample_mode,
            )
            for k, v in mod_dict.items()
        }

        if self.add_sizes:
            mod_dict["crop_coords"] = torch.tensor(crop_coords)
            mod_dict["orig_size"] = torch.tensor(orig_size)

        return mod_dict

    def __call__(self, mod_dict):
        """Apply the augmentation to a dict of modalities (both image based and sequence based modalities)

        Args:
            mod_dict (dict): Dict of modalities

        Returns:
            dict: Transformed dict of modalities
        """
        crop_settings = mod_dict.pop("crop_settings", None)

        mod_dict = {k: get_transform(k, self.transforms_dict).preprocess(v) for k, v in mod_dict.items()}

        mod_dict = self.unified_image_augment(mod_dict, crop_settings)

        mod_dict = {k: get_transform(k, self.transforms_dict).postprocess(v) for k, v in mod_dict.items()}

        return mod_dict

    def __repr__(self):
        repr = "(UnifiedDataAugmentation,\n"
        repr += ")"
        return repr


class AbstractTransform(ABC):

    @abstractmethod
    def load(self, sample):
        pass

    @abstractmethod
    def preprocess(self, sample):
        pass

    @abstractmethod
    def image_augment(self, v, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple,
                      rand_aug_idx: Optional[int], resample_mode: str = None):
        pass

    @abstractmethod
    def postprocess(self, v):
        pass


class ImageTransform(AbstractTransform):

    @staticmethod
    def pil_loader(path: str) -> Image.Image:
        # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
        # with open(path, 'rb') as f:
        #     img = Image.open(f)
        img = Image.open(path)
        return img


    @staticmethod
    def image_hflip(img: Image, flip: bool):
        """Crop and resize an image

        :param img: Image to crop and resize
        :param flip: Whether to flip the image
        :return: Flipped image (if flip = True)
        """
        if flip:
            img = TF.hflip(img)
        return img

    @staticmethod
    def image_crop_and_resize(img: Image, crop_coords: Tuple, target_size: Tuple, resample_mode: str = None):
        """Crop and resize an image

        :param img: Image to crop and resize
        :param crop_coords: Coordinates of the crop (top, left, h, w)
        :param target_size: Coordinates of the resize (height, width)
        :return: Cropped and resized image
        """

        top, left, h, w = crop_coords
        resize_height, resize_width = target_size
        img = TF.crop(img, top, left, h, w)
        resample_mode = get_pil_resample_mode(resample_mode)
        img = img.resize((resize_height, resize_width), resample=resample_mode)
        return img

    # @staticmethod
    # def image_crop_and_resize_for_vq_tokens(sampled_frames, only_spatial, target_size, resample_mode,
    #                                          resize_transform):
    #     # Uses raw video as input, applies the cropping, normalization and then the resize transforms
    #
    #     pass


class RGBTransform(ImageTransform):

    def __init__(self, imagenet_default_mean_and_std=True, color_jitter=False, color_jitter_strength=0.5):
        self.rgb_mean = IMAGENET_INCEPTION_MEAN if not imagenet_default_mean_and_std else IMAGENET_DEFAULT_MEAN
        self.rgb_std = IMAGENET_INCEPTION_STD if not imagenet_default_mean_and_std else IMAGENET_DEFAULT_STD
        self.color_jitter = color_jitter
        self.color_jitter_transform = self.random_color_jitter(color_jitter_strength)

    def random_color_jitter(self, strength=0.5):
        # Color Jitter from Pix2Seq and SimCLR
        # Source: https://github.com/google-research/pix2seq/blob/main/data/data_utils.py#L114
        t = T.Compose([
            T.RandomApply([T.ColorJitter(brightness=0.8 * strength, contrast=0.8 * strength, saturation=0.8 * strength, hue=0.2 * strength)], p=0.8),
            T.RandomApply([T.Grayscale(num_output_channels=3)], p=0.2),
        ])

        return t

    def rgb_to_tensor(self, img):
        img = TF.to_tensor(img)
        img = TF.normalize(img, mean=self.rgb_mean, std=self.rgb_std)
        return img

    def load(self, path):
        # TODO: Instead of converting to RGB here, do it either in the preprocess or the postprocess step. Makes it compatible with wds dataloading.
        sample = self.pil_loader(path)
        return sample

    def preprocess(self, sample):
        sample = sample.convert('RGB')
        
        if self.color_jitter:
            sample = self.color_jitter_transform(sample)

        return sample

    def image_augment(self, img, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple, 
                      rand_aug_idx: Optional[int], resample_mode: str = None):
        img = self.image_crop_and_resize(img, crop_coords, target_size, resample_mode=resample_mode)
        img = self.image_hflip(img, flip)
        return img


    def postprocess(self, sample):
        sample = self.rgb_to_tensor(sample)
        return sample

class JEPA_Video_Transform(ImageTransform):

    def __init__(self, total_frames_to_be_extracted=17):
        self.full_transforms = make_transforms(
        random_horizontal_flip=False, # no flipping, as it can change the class label too..
        random_resize_aspect_ratio=[0.75, 1.35],
        random_resize_scale=[0.3, 1.0],
        reprob=0.,
        auto_augment=False,
        motion_shift=False,
        crop_size=224)
        self.total_frames_to_be_extracted = total_frames_to_be_extracted - 1 # substract 1 to become consistant with video understanding encoders (16 frames)

    def random_color_jitter(self, strength=0.5):
        # Color Jitter from Pix2Seq and SimCLR
        # Source: https://github.com/google-research/pix2seq/blob/main/data/data_utils.py#L114
        t = T.Compose([
            T.RandomApply([T.ColorJitter(brightness=0.8 * strength, contrast=0.8 * strength,
                                         saturation=0.8 * strength, hue=0.2 * strength)], p=0.8),
            T.RandomApply([T.Grayscale(num_output_channels=3)], p=0.2),
        ])

        return t

    def preprocess(self, sample):
        # just return decord object as it is
        return sample

    def image_augment(self, img, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple,
                                rand_aug_idx: Optional[int], resample_mode: str = None, resize_transform=None):
        # First extract the given window...
        starting_time, ending_time = crop_coords[0], crop_coords[1]
        original_video_fps = img.get_avg_fps()
        selected_frame_indices = np.linspace(
            starting_time * original_video_fps,
            ending_time * original_video_fps,
            self.total_frames_to_be_extracted,
            dtype=np.int32
        )
        # remove the last one due to indexing rules
        selected_frame_indices[-1] = selected_frame_indices[-1] - 1
        # sample frames from the given window
        frames = img.get_batch(selected_frame_indices)
        # now lets remove the video_reader
        img.seek(0)
        del img
        sampled_frames = torch.tensor(frames.asnumpy()).permute(0, 3, 1, 2)
        # perform crop and resizing based on crop setting
        top, left, h, w = (crop_coords[2], crop_coords[3], crop_coords[4], crop_coords[5])
        # Let's crop the image,
        img = TF.crop(sampled_frames, int(top), int(left), int(h), int(w))
        # now use the v-jepa transforms
        processed_img = self.full_transforms(img)
        # return back in expected shape
        return processed_img

    def load(self, path):
        # TODO: Instead of converting to RGB here, do it either in the preprocess or the postprocess step. Makes it compatible with wds dataloading.
        # sample = self.pil_loader(path)
        sample = decord.VideoReader(path, num_threads=1)
        return sample
    def postprocess(self, sample):
        # just return processed sample as it is.
        return sample


class RGB_Video_Transform(ImageTransform):

    def __init__(self, imagenet_default_mean_and_std=True, color_jitter=False, color_jitter_strength=0.5, total_frames_to_be_extracted=17, target_size=128):
        self.rgb_mean = IMAGENET_INCEPTION_MEAN if not imagenet_default_mean_and_std else IMAGENET_DEFAULT_MEAN
        self.rgb_std = IMAGENET_INCEPTION_STD if not imagenet_default_mean_and_std else IMAGENET_DEFAULT_STD
        self.color_jitter = color_jitter
        self.color_jitter_transform = self.random_color_jitter(color_jitter_strength)
        self.total_frames_to_be_extracted = total_frames_to_be_extracted
        self.resize_transform = T.Compose([
                T.Resize((target_size, target_size)),
            ])

    def random_color_jitter(self, strength=0.5):
        # Color Jitter from Pix2Seq and SimCLR
        # Source: https://github.com/google-research/pix2seq/blob/main/data/data_utils.py#L114
        t = T.Compose([
            T.RandomApply([T.ColorJitter(brightness=0.8 * strength, contrast=0.8 * strength,
                                         saturation=0.8 * strength, hue=0.2 * strength)], p=0.8),
            T.RandomApply([T.Grayscale(num_output_channels=3)], p=0.2),
        ])

        return t

    def rgb_to_tensor(self, img):
        # img = TF.to_tensor(img) # its already done for videos
        img = TF.normalize(img, mean=self.rgb_mean, std=self.rgb_std)
        return img

    def load(self, path):
        # TODO: Instead of converting to RGB here, do it either in the preprocess or the postprocess step. Makes it compatible with wds dataloading.
        sample = self.pil_loader(path)
        return sample

    def preprocess(self, sample):
        # just return decord object as it is
        return sample

    def image_augment(self, img, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple,
                                rand_aug_idx: Optional[int], resample_mode: str = None, resize_transform=None):
        # First extract the given window...
        starting_time, ending_time = crop_coords[0], crop_coords[1]
        original_video_fps = img.get_avg_fps()
        selected_frame_indices = np.linspace(
            starting_time * original_video_fps,
            ending_time * original_video_fps,
            self.total_frames_to_be_extracted,
            dtype=np.int32
        )
        # remove the last one due to indexing rules
        selected_frame_indices[-1] = selected_frame_indices[-1] - 1
        # sample frames from the given window
        frames = img.get_batch(selected_frame_indices)
        # now lets remove the video_reader
        img.seek(0)
        del img
        sampled_frames = torch.tensor(frames.asnumpy()).permute(0, 3, 1, 2)
        # sampled_frames = torch.tensor(img.get_batch(selected_frame_indices).asnumpy()).permute(0, 3, 1, 2)
        # perform crop and resizing based on crop setting
        only_spatial = (crop_coords[2], crop_coords[3], crop_coords[4], crop_coords[5])
        img = self.image_crop_and_resize_for_vq_tokens(sampled_frames, only_spatial, target_size, resample_mode=resample_mode,
                                                       resize_transform=self.resize_transform)
        # img = self.image_hflip(img, flip) # actually no flipping is required for video case..., as it can change the label
        return img

    def postprocess(self, sample):
        sample = self.rgb_to_tensor(sample)
        return sample


class DepthTransform(ImageTransform):

    def __init__(self, standardize_depth=True):
        self.standardize_depth = standardize_depth

    def depth_to_tensor(self, img):
        img = torch.Tensor( img / (2 ** 16 - 1.0) )
        img = img.unsqueeze(0)  # 1 x H x W
        if self.standardize_depth:
            img = self.truncated_depth_standardization(img)
        return img

    @staticmethod
    def truncated_depth_standardization(depth, thresh: float = 0.1):
        """Truncated depth standardization

        :param depth: Depth map
        :param thresh: Threshold
        :return: Robustly standardized depth map
        """
        # Flatten depth and remove bottom and top 10% of values
        trunc_depth = torch.sort(depth.reshape(-1), dim=0)[0]
        trunc_depth = trunc_depth[int(thresh * trunc_depth.shape[0]): int((1 - thresh) * trunc_depth.shape[0])]
        return (depth - trunc_depth.mean()) / torch.sqrt(trunc_depth.var() + 1e-6)

    def load(self, path):
        sample = self.pil_loader(path)
        return sample

    def preprocess(self, sample):
        return sample

    def image_augment(self, img, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple, 
                      rand_aug_idx: Optional[int], resample_mode: str = None):
        img = self.image_crop_and_resize(img, crop_coords, target_size, resample_mode=resample_mode)
        img = self.image_hflip(img, flip)
        return img

    def postprocess(self, sample):
        sample = np.array(sample)
        sample = self.depth_to_tensor(sample)
        return sample


class NormalTransform(ImageTransform):

    def __init__(self, standardize_surface_normals=False):
        self.normal_mean = (0.5, 0.5, 0.5) if not standardize_surface_normals else IMAGENET_SURFACE_NORMAL_MEAN
        self.normal_std = (0.5, 0.5, 0.5) if not standardize_surface_normals else IMAGENET_SURFACE_NORMAL_STD

    def normal_to_tensor(self, img):
        img = TF.to_tensor(img)
        img = TF.normalize(img, mean=self.normal_mean, std=self.normal_std)
        return img

    def load(self, path):
        sample = self.pil_loader(path)
        return sample

    def preprocess(self, sample):
        return sample

    def image_hflip(self, img: Image, flip: bool):
        if flip:
            img = TF.hflip(img)
            flipped_np = np.array(img)
            flipped_np[:, :, 0] = 255 - flipped_np[:, :, 0]
            img = Image.fromarray(flipped_np)

        return img

    def image_augment(self, img, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple,
                      rand_aug_idx: Optional[int], resample_mode: str = None):
        img = self.image_crop_and_resize(img, crop_coords, target_size, resample_mode=resample_mode)
        img = self.image_hflip(img, flip)
        return img

    def postprocess(self, sample):
        sample = self.normal_to_tensor(sample)
        return sample
    

class SemsegTransform(ImageTransform):

    def __init__(self, scale_factor=1.0, shift_idx_by_one=False, id_mapping: Optional[Dict] = None, select_channel=None):
        self.scale_factor = scale_factor
        self.shift_idx_by_one = shift_idx_by_one
        self.id_mapping = id_mapping
        self.select_channel = select_channel

    def map_semseg_values(self, sample):
        sample = np.asarray(sample)
        mapping_fn = lambda x: self.id_mapping.get(x, x)
        sample = np.vectorize(mapping_fn)(sample)
        sample = Image.fromarray(sample, mode='P')
        return sample

    def semseg_to_tensor(self, img):
        # Rescale to scale factor
        if self.scale_factor != 1.0:
            target_height, target_width = int(img.height * self.scale_factor), int(img.width * self.scale_factor)
            img = img.resize((target_width, target_height))
        # Using pil_to_tensor keeps it in uint8, to_tensor converts it to float (rescaled to [0, 1])
        img = TF.pil_to_tensor(img).to(torch.long).squeeze(0)
        # 255->0, 254->0, all else shifted up by one
        return img

    def load(self, path):
        sample = self.pil_loader(path)
        if self.select_channel is not None:
            sample = sample.split()[self.select_channel]
        return sample

    def preprocess(self, sample):
        sample = sample.convert('P')

        if self.id_mapping is not None:
            sample = self.map_semseg_values(sample)

        if self.shift_idx_by_one:
            sample = np.asarray(sample)
            sample = sample + 1
            sample = Image.fromarray(sample, mode='P')

        return sample

    def image_augment(self, img, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple,
                      rand_aug_idx: Optional[int], resample_mode: str = None):
        # Value for padding with TF.crop is always 0.
        # Override resampling mode to 'nearest' for semseg
        img = self.image_crop_and_resize(img, crop_coords, target_size, resample_mode='nearest')
        img = self.image_hflip(img, flip)
        return img

    def postprocess(self, sample):
        img = self.semseg_to_tensor(sample)
        return img


class SAMInstanceTransform(AbstractTransform):

    def __init__(self, mask_size=64, max_instance_n=20, bbox_area_threshold=0.0005):
        self.mask_size = mask_size
        self.max_instance_n = max_instance_n
        self.bbox_area_threshold = bbox_area_threshold

    def get_bbox(self, instance):
        """ Gets bounding box of the given instance
        """
        min_h, max_h =  instance[:,:,1].min(), instance[:,:,1].max()
        min_w, max_w =  instance[:,:,0].min(), instance[:,:,0].max()
        return [min_h, min_w, max_h, max_w]

    def extend_instance_points(self, instance, border_fn):
        """ Given an instance and a border function `border_fn`, extends the instance points with crossing points between the instance and 
        the crop borders. The crossing points are obtained using border_fn.
        """
        p = instance[:,0]
        p_next = np.roll(p, (-1), axis=(0))
        final_points = []
        for x, xn in zip(p, p_next):
            final_points.append(x)
            for r in border_fn(x, xn):
                final_points.append(r.astype(np.int32))
        p = np.stack(final_points)
        return p[:,None]

    def remove_redundant_lines(self, orig_instance, instance):
        """ Removes the redundant lines added during cropping.
        """
        final_points = []
        for p in instance:
            distance = cv2.pointPolygonTest(orig_instance, (p[0,0].item(), p[0,1].item()), measureDist=True)
            if distance >= 0:
                final_points.append(p[0])
        return np.stack(final_points)[:,None]

    def get_border_functions(self, crop_points):
        """ Creates and returns a function `fn` using crop region coordinates given in crop_points.
        `fn` receives two input points x and xn and returns all the crossing points between the line connecting
        x and xn, and the borders of the cropping rectangle.
        """
        p = crop_points[:,0]
        p_next = np.roll(p, (-1), axis=(0))
        def fn(x, xn):
            output = []
            c_diff = p_next - p
            x_diff = x - xn
            for diff, c in zip(c_diff, p):
                A = np.array([
                        [diff[0], x_diff[0]],
                        [diff[1], x_diff[1]]
                    ])
                b = x - c
                try:
                    lmbda = np.linalg.solve(A, b)
                    if 0 <= lmbda[0] <= 1 and 0 <= lmbda[1] <= 1:
                        output.append(lmbda[1] * xn + (1-lmbda[1]) * x)
                except:
                    continue
            return output
        return fn

    def crop_sample(self, sample, crop_coords):
        """ Crop the sample using crop coordinates.
        """
        top, left, h, w = crop_coords
        crop_region = (left, top, left + w, top + h)
        crop_points = np.array([
            [crop_region[0], crop_region[1]],
            [crop_region[2], crop_region[1]],
            [crop_region[2], crop_region[3]],
            [crop_region[0], crop_region[3]],
        ])[:,None]
        border_functions = self.get_border_functions(crop_points)
        cropped_sample = []
        for instance in sample:
            instance = self.extend_instance_points(instance, border_functions)
            filter_condition = (
                (instance[:, :, 0] > crop_region[0]) &
                (instance[:, :, 0] < crop_region[2]) &
                (instance[:, :, 1] > crop_region[1]) &
                (instance[:, :, 1] < crop_region[3])
            )
            if not np.any(filter_condition):
                continue
            
            instance_copy = instance.copy()
            instance_copy[:, :, 0] = np.clip(instance[:, :, 0], a_min=crop_region[0], a_max=crop_region[2])
            instance_copy[:, :, 1] = np.clip(instance[:, :, 1], a_min=crop_region[1], a_max=crop_region[3])
            instance_copy = self.remove_redundant_lines(instance, instance_copy)
            instance_copy[:, :, 0] -= crop_region[0]
            instance_copy[:, :, 1] -= crop_region[1]

            cropped_sample.append(instance_copy)
        return cropped_sample
    
    def resize_sample(self, sample, original_size, target_size):
        """ Resize the sample
        """
        width_scale = target_size[1] / original_size[1]
        height_scale = target_size[0] / original_size[0]
        resized_sample = []
        for instance in sample:
            instance_copy = instance.copy()
            instance_copy[:, :, 0] = np.round(width_scale * instance_copy[:, :, 0])
            instance_copy[:, :, 1] = np.round(height_scale * instance_copy[:, :, 1])
            resized_sample.append(instance_copy)
        return resized_sample
    
    def remove_tiny_instances(self, sample, image_size):
        """ Remove instances that have an area ratio smaller than `bbox_area_threshold`.
        """
        filtered_sample = []
        for instance in sample:
            min_h, min_w, max_h, max_w = self.get_bbox(instance)
            bbox_area_ratio = (max_h - min_h) * (max_w - min_w) / (image_size[0] * image_size[1])
            if bbox_area_ratio < self.bbox_area_threshold:
                continue
            filtered_sample.append(instance)
        return filtered_sample

    def hflip(self, sample, width):
        """ Horizontal flipping the instances in a sample.
        """
        flipped_sample = []
        for instance in sample:
            instance_copy = instance.copy()
            instance_copy[:, :, 0] = width - instance_copy[:, :, 0]
            flipped_sample.append(instance_copy)
        return flipped_sample
        
    def get_binary_masks(self, sample):
        """ Creates the binary mask of each instance in the sample.
        """
        if self.max_instance_n is None:
            max_instance_n = len(sample)
        else:
            max_instance_n = self.max_instance_n
        masks = np.zeros((max_instance_n, self.mask_size, self.mask_size)) 
        bboxes = np.zeros((max_instance_n, 4))
        valid = np.full(max_instance_n, False)
        for i, instance in enumerate(sample):
            bbox = self.get_bbox(instance)
            min_h, min_w, max_h, max_w = bbox
            instance_copy = instance.copy()
            mask = np.zeros((self.mask_size, self.mask_size), dtype=np.uint8)
            instance_copy[:,:,0] = (instance_copy[:,:,0] - min_w) / (max_w - min_w) * self.mask_size
            instance_copy[:,:,1] = (instance_copy[:,:,1] - min_h) / (max_h - min_h) * self.mask_size
            cv2.drawContours(mask, [instance_copy], 0, (255), thickness=cv2.FILLED)
            masks[i] = mask / 255.0
            bboxes[i] = np.array(bbox)
            valid[i] = True
        return masks, bboxes, valid

    def load(self, path):
        sample = np.load(path, allow_pickle=True)
        return sample

    def preprocess(self, sample):
        if self.max_instance_n is None or len(sample) <= self.max_instance_n:
            indecies = np.arange(len(sample))
        else:
            indecies = np.random.choice(len(sample), size=self.max_instance_n, replace=False)
        return [p['points'] for i, p in enumerate(sample) if i in indecies]

    def image_augment(self, v, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple,
                      rand_aug_idx: Optional[int], resample_mode: str = None):
        v = self.crop_sample(v, crop_coords)
        _, _, h, w = crop_coords
        v = self.resize_sample(v, (h, w), target_size)
        v = self.remove_tiny_instances(v, target_size)
        if flip:
            v = self.hflip(v, target_size[0])
        return v

    def postprocess(self, sample):
        sample, bboxes, valid = self.get_binary_masks(sample)
        return {
            'instance': torch.from_numpy(sample).to(torch.float32), 
            'bbox': torch.from_numpy(bboxes).to(torch.float32), 
            'valid': torch.from_numpy(valid)
        }


class MaskTransform(ImageTransform):

    def __init__(self, mask_pool_size=1):
        assert isinstance(mask_pool_size, int)
        self.mask_pool_size = mask_pool_size # Use to expand masks

    def mask_to_tensor(self, img):
        mask = TF.to_tensor(img)
        if self.mask_pool_size > 1:
            mask = reduce(mask, 'c (h1 h2) (w1 w2) -> c h1 w1', 'min', h2=self.mask_pool_size, w2=self.mask_pool_size)
            mask = repeat(mask, 'c h1 w1 -> c (h1 h2) (w1 w2)', h2=self.mask_pool_size, w2=self.mask_pool_size)
        return (mask == 1.0)

    def load(self, path):
        sample = self.pil_loader(path)
        return sample

    def preprocess(self, sample):
        return sample

    def image_augment(self, img, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple,
                      rand_aug_idx: Optional[int], resample_mode: str = None):
        # Override resampling mode to 'nearest' for masks
        img = self.image_crop_and_resize(img, crop_coords, target_size, resample_mode='nearest')
        img = self.image_hflip(img, flip)
        return img

    def postprocess(self, sample):
        sample = self.mask_to_tensor(sample)
        return sample


class TokTransform(AbstractTransform):

    def __init__(self):
        pass

    def load(self, path):
        sample = np.load(path).astype(int)
        return sample

    def preprocess(self, sample):
        return sample

    def image_augment(self, v, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple,
                      rand_aug_idx: Optional[int], resample_mode: str = None):
        if rand_aug_idx is None:
            raise ValueError("Crop settings / augmentation index are missing but a pre-tokenized modality is being used")
        v = torch.tensor(v[rand_aug_idx])
        return v

    def postprocess(self, sample):
        return sample


class DetectionTransform(AbstractTransform):

    def __init__(self, det_threshold=0.6, det_max_instances=None, bbox_order='dist_to_orig', coord_bins=1000, min_visibility=0.0, return_raw=False):
        self.det_threshold = det_threshold
        self.det_max_instances = det_max_instances
        self.coord_bins = coord_bins
        self.min_visibility = min_visibility
        self.return_raw = return_raw

        if bbox_order == 'area':
            self.bbox_order = self.order_bboxes_by_area
        elif bbox_order == 'score':
            self.bbox_order = self.order_bboxes_by_score
        elif bbox_order == 'random':
            self.bbox_order = self.shuffle_bboxes
        else:
            self.bbox_order = self.order_bboxes_by_dist_to_orig

    @staticmethod
    def order_bboxes_by_area(bboxes):
        return sorted(bboxes, key=lambda x: (x[2] - x[0]) * (x[3] - x[1]), reverse=True)

    @staticmethod
    def order_bboxes_by_dist_to_orig(bboxes):
        return sorted(bboxes, key=lambda x: x[0] ** 2 + x[1] ** 2)

    @staticmethod
    def order_bboxes_by_score(bboxes):
        return sorted(bboxes, key=lambda x: x[5], reverse=True)

    @staticmethod
    def shuffle_bboxes(bboxes):
        return sorted(bboxes, key=lambda x: random.random())

    def convert_detection_instance(self, instances):
        """Convert instances dict to list of lists where each list takes the form:
        [xmin, ymin, xmax, ymax, class_name, score]
        """

        instances = [inst['boxes'] + [inst['class_name'], inst['score']] for inst in instances if inst['score'] >= self.det_threshold]
        return instances

    def bboxes_hflip(self, bboxes: List[Tuple], image_size: Tuple, flip: bool):
        image_height, image_width = image_size
        if flip:
            bboxes = [tuple(A.bbox_hflip(bbox[:4], rows=image_height, cols=image_width)) + tuple(bbox[4:])
                      for bbox in bboxes]

        return bboxes

    def bboxes_crop_and_resize(self, bboxes: List[Tuple], crop_coords: Tuple, orig_size: Tuple):
        """Crop and resize bounding boxes

        Args:
            bboxes: Bounding boxes to crop and resize
            crop_coords: Coordinates of the crop (top, left, h, w)
            orig_size: Size of the original image

        Returns:
            Cropped and resized bounding boxes
        """
        orig_height, orig_width = orig_size
        top, left, h, w = crop_coords
        xmin, ymin, xmax, ymax = left, top, left + w, top + h
        bboxes = [tuple(A.bbox_crop(bbox[:4], x_min=xmin, y_min=ymin, x_max=xmax, y_max=ymax, rows=orig_height,
                                    cols=orig_width)) + tuple(bbox[4:])
                  for bbox in bboxes]
        bboxes = A.core.bbox_utils.filter_bboxes(bboxes, rows=h, cols=w, min_visibility=self.min_visibility)
        # No need to resize, bounding boxes in albumentations format are scale invariant

        return bboxes

    def order_and_filter_bboxes(self, bboxes):
        if self.det_max_instances is not None and len(bboxes) > self.det_max_instances:
            bboxes = self.order_bboxes_by_score(bboxes)[:self.det_max_instances]

        return self.bbox_order(bboxes)

    def convert_bboxes_to_string(self, bboxes: List[Tuple]):
        """Convert bounding boxes to a string. 
        xmin, ymin, xmax, ymax are mapped to v0, v1, v2, v3 special tokens.

        Args:
            bboxes: Bounding boxes

        Returns:
            String representation of the bounding boxes
        """
        # Remove score, quantize coordinates
        bins = self.coord_bins

        bboxes = [
            [
                f"v0={round(xmin * (bins - 1))}",
                f"v1={round(ymin * (bins - 1))}",
                f"v2={round(xmax * (bins - 1))}",
                f"v3={round(ymax * (bins - 1))}",
                cls,
            ]
            for (xmin, ymin, xmax, ymax, cls, score) in bboxes
        ]
        # Convert each bounding box to a string
        bboxes = [' '.join(b) for b in bboxes]
        # Convert the list to a str
        return ' '.join(bboxes)

    def load(self, path):
        with open(path, 'r') as f:
            sample = json.load(f)

        return sample

    def preprocess(self, sample):
        instances = sample['instances']
        return self.convert_detection_instance(instances)

    def image_augment(self, bboxes: List[Tuple], crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple, 
                      rand_aug_idx=None, resample_mode: str = None):
        bboxes = self.bboxes_crop_and_resize(bboxes, crop_coords, orig_size)
        bboxes = self.bboxes_hflip(bboxes, target_size, flip)
        bboxes = self.order_and_filter_bboxes(bboxes)
        return bboxes

    def postprocess(self, bboxes):
        if self.return_raw:
            return bboxes
        bboxes = self.convert_bboxes_to_string(bboxes)
        return bboxes


class VideoDetectionTransform(AbstractTransform):

    def __init__(self, det_threshold=0.6, det_max_instances=None, max_main_instances=4, max_passive_instances=2, bbox_order='dist_to_orig', coord_bins=1000, min_visibility=0.0, return_raw=False):
        self.det_threshold = det_threshold
        self.det_max_instances = det_max_instances
        self.coord_bins = coord_bins
        self.min_visibility = min_visibility
        self.return_raw = return_raw
        self.max_main_instances = max_main_instances
        self.max_passive_instances = max_passive_instances
        self.total_max_instances = self.max_main_instances + self.max_passive_instances
        self.special_detection_tokens_video = ['[FRAME_1]','[FRAME_2]','[FRAME_3]','[FRAME_4]','[FRAME_5]','[FRAME_6]','[FRAME_7]','[FRAME_8]','[FRAME_9]','[FRAME_10]','[FRAME_11]','[FRAME_12]','[FRAME_13]','[FRAME_14]','[FRAME_15]','[FRAME_16]','[FRAME_17]']
        if bbox_order == 'area':
            self.bbox_order = self.order_bboxes_by_area
        elif bbox_order == 'score':
            self.bbox_order = self.order_bboxes_by_score
        elif bbox_order == 'random':
            self.bbox_order = self.shuffle_bboxes
        else:
            self.bbox_order = self.order_bboxes_by_dist_to_orig

    @staticmethod
    def order_bboxes_by_area(bboxes):
        return sorted(bboxes, key=lambda x: (x[2] - x[0]) * (x[3] - x[1]), reverse=True)

    @staticmethod
    def order_bboxes_by_dist_to_orig(bboxes):
        return sorted(bboxes, key=lambda x: x[0] ** 2 + x[1] ** 2)

    @staticmethod
    def order_bboxes_by_score(bboxes):
        return sorted(bboxes, key=lambda x: x[5], reverse=True)

    @staticmethod
    def shuffle_bboxes(bboxes):
        return sorted(bboxes, key=lambda x: random.random())

    def create_chunks(self, chunk_sizes, data_list):
        """Using list comprehension with accumulate"""
        indices = [0] + list(accumulate(chunk_sizes))
        all_detections = [data_list[indices[i]:indices[i + 1]][:self.total_max_instances] for i in range(len(chunk_sizes))]
        updated_chunk_sizes = [min(self.total_max_instances, current_frame_instances) for current_frame_instances in chunk_sizes]
        # Lets only retain self.total_max_instances for each frame
        # modified_list = [sublist[:self.total_max_instances] for sublist in all_detections]
        return all_detections, updated_chunk_sizes

    def reformat_video_detection_instances(self, instances_list, main_instance_masks):
        """
        Reformat video detection instances with improved efficiency.

        Args:
            instances_list: List of frame instances
            main_instance_masks: Corresponding main instance masks

        Returns:
            Tuple of (frame_detections, per_frame_num_instances)
        """
        frame_detections = []
        per_frame_num_instances = []

        for frame_idx, (frame_instances, frame_mask) in enumerate(zip(instances_list, main_instance_masks)):


            # # Efficiently slice and combine main and passive instances
            # total_main_instances = sum(frame_mask)
            # main_instances = frame_instances[:min(total_main_instances, self.max_main_instances)]
            # passive_instances = frame_instances[total_main_instances:][:self.max_passive_instances]
            # final_instances = main_instances + passive_instances
            # Lets only use the main instances
            final_instances = frame_instances

            # Process instances in a single comprehension
            frame_detections_current = [
                inst['bbox'] + [inst['label'], inst['confidence'], [frame_idx]]
                for inst in final_instances
            ]

            frame_detections.extend(frame_detections_current)
            per_frame_num_instances.append(len(frame_detections_current))

        return frame_detections, per_frame_num_instances

    def bboxes_hflip(self, bboxes: List[Tuple], image_size: Tuple, flip: bool):
        image_height, image_width = image_size
        if flip:
            bboxes = [tuple(A.bbox_hflip(bbox[:4], rows=image_height, cols=image_width)) + tuple(bbox[4:])
                      for bbox in bboxes]

        return bboxes

    def bboxes_crop_and_resize(self, bboxes: List[Tuple], crop_coords: Tuple, orig_size: Tuple, per_frame_num_instances):
        """Crop and resize bounding boxes

        Args:
            bboxes: Bounding boxes to crop and resize
            crop_coords: Coordinates of the crop (top, left, h, w)
            orig_size: Size of the original image

        Returns:
            Cropped and resized bounding boxes
        """
        orig_width, orig_height = orig_size
        _, _, top, left, h, w = crop_coords
        xmin, ymin, xmax, ymax = left, top, left + w, top + h
        bboxes_scaled = [tuple(A.bbox_crop(bbox[:4], x_min=xmin, y_min=ymin, x_max=xmax, y_max=ymax, rows=orig_height,
                                    cols=orig_width)) + tuple(bbox[4:])
                  for bbox in bboxes]
        # Also update the per_frame_num_instances, given that some will be filtered out
        filtered_bbox_indices = np.array([bboxes_scaled.index(x) for x in bboxes_scaled if (x[2] <= 0 or x[0] >= 1 or x[3] <= 0) or x[1] >= 1])
        bbox_to_frame = np.repeat(np.arange(len(per_frame_num_instances)), per_frame_num_instances)
        if len(filtered_bbox_indices) != 0:
            # only update per_frame_num_instances if there are any filtered bboxes
            per_frame_num_instances = per_frame_num_instances - np.bincount(bbox_to_frame[filtered_bbox_indices],
                                                                  minlength=len(per_frame_num_instances))
        bboxes_filtered = A.core.bbox_utils.filter_bboxes(bboxes_scaled, rows=h, cols=w, min_visibility=self.min_visibility)
        assert len(bboxes_filtered) == sum(per_frame_num_instances), "something wrong in the code, history of boxes per frames not tracked correctly"
        # if np.any(per_frame_num_instances == 0):
        #     print("Warning: some frames have no instances left after cropping and resizing")
        # print(per_frame_num_instances)
        assert  not np.any(per_frame_num_instances < 0), f"there cannot be negative instances left after cropping and resizing {per_frame_num_instances}"
        # No need to resize, bounding boxes in albumentations format are scale invariant
        return bboxes_filtered, per_frame_num_instances.tolist()

    def merge_tokens_with_frames(self, tokens, frames, counts):
        """Merge tokens with frame identifiers efficiently."""
        if not tokens or not frames or not counts:
            return ""

        output = []
        indices = [0] + list(accumulate(counts))

        for i in range(len(counts)):
            frame_tokens = ' '.join(tokens[indices[i]:indices[i + 1]])
            merged = f"{frames[i]} {frame_tokens} [EOS]" if frame_tokens else frames[i]
            output.append(merged)

        return ' '.join(output)

    def order_and_filter_bboxes(self, all_frame_bboxes):
        """Apply ordering to each frame's bboxes and flatten the result."""

        processed_bboxes = []
        for frame_bboxes in all_frame_bboxes:
            processed_bboxes.extend(self.bbox_order(frame_bboxes))
        return processed_bboxes

    def convert_bboxes_to_string(self, bboxes_tuple):
        """Convert bounding boxes to a string.
        xmin, ymin, xmax, ymax are mapped to v0, v1, v2, v3 special tokens.

        Args:
            bboxes: Bounding boxes and instance_mask per frame (list)

        Returns:
            String representation of the bounding boxes
        """
        bboxes, per_frame_num_instances = bboxes_tuple
        # Remove score, quantize coordinates
        bins = self.coord_bins
        # let's also add special detection tags for each frame
        bboxes = [
            [
                f"v0={round(xmin * (bins - 1))}",
                f"v1={round(ymin * (bins - 1))}",
                f"v2={round(xmax * (bins - 1))}",
                f"v3={round(ymax * (bins - 1))}",
                cls,
            ]
            for (xmin, ymin, xmax, ymax, cls, score, indices) in bboxes
        ]
        # Now,
        bboxes = [' '.join(b) for b in bboxes]
        # Convert each bounding box to a string
        # single_processed_string = self.merge_tokens_with_frames(bboxes, self.special_detection_tokens_video, per_frame_num_instances) # this will be now used directly in masking function
        return (bboxes, per_frame_num_instances)

    def load(self, path):
        with open(path, 'r') as f:
            # this will contain per-frame detections
            sample = json.load(f)
        return sample

    def preprocess(self, sample):
        """Preprocess the input sample (pass-through for now)."""
        return sample


    def image_augment(self, frame_detections, crop_coords: Tuple, flip: bool,
                      orig_size: Tuple, target_size: Tuple, rand_aug_idx=None, resample_mode: str = None):
        """Apply image augmentations to frame detections."""

        # Let's first sample the frames required for this specific crop:

        starting_time, ending_time, i, j, h, w = crop_coords # e.g., 3, 7, 124, 0, 320, 320, 0
        video_info = frame_detections['video_info']
        fps = video_info['fps']

        selected_frame_indices = np.linspace(
            starting_time * fps,
            ending_time * fps,
            17,
            dtype=np.int32
        )
        selected_frame_indices[-1] = selected_frame_indices[-1] - 1

        # Update original size from video info
        orig_size = (video_info['width'], video_info['height'])

        # now let's retrieve only those frames relevant for this given crop setting:
        detections = frame_detections['detections']
        main_masks = frame_detections['main_instance_masks']

        selected_frame_detections = [detections[f"frame{idx}"] for idx in selected_frame_indices]
        mask_instances_list = [main_masks[f"frame{idx}"] for idx in selected_frame_indices]

        # Process detections
        preprocessed_detections_combined_frames, per_frame_num_instances = self.reformat_video_detection_instances(
            selected_frame_detections, mask_instances_list
        )


        # Apply transformations

        detections, per_frame_num_instances = self.bboxes_crop_and_resize(preprocessed_detections_combined_frames, crop_coords, orig_size, np.array(per_frame_num_instances))
        detections = self.bboxes_hflip(detections, target_size, flip)

        # Reorganize into per-frame chunks and apply ordering
        per_frame_detections, num_instances_updated = self.create_chunks(per_frame_num_instances, detections)
        detections = self.order_and_filter_bboxes(per_frame_detections)
        assert len(detections) == sum(num_instances_updated), "something wrong in the code, history of boxes per frames not tracked correctly"
        return (detections, num_instances_updated)

    def postprocess(self, bboxes):
        if self.return_raw:
            return bboxes
        bboxes = self.convert_bboxes_to_string(bboxes)
        return bboxes


class CaptionTransform(AbstractTransform):

    def __init__(self, aligned_captions=True, no_aug=False):
        self.aligned_captions = aligned_captions
        self.no_aug = no_aug

    def load(self, path):
        # Caption can either be stored as .txt or .json.gz (in which case it's a list of dicts)
        if path.endswith('.txt'):
            sample = Path(path).read_text()
        elif path.endswith('.json'):
            with open(path, 'r') as f:
                sample = json.load(f)
        elif path.endswith('.json.gz'):
            with gzip.open(path, 'rb') as f:
                sample = json.load(f)
        return sample

    def preprocess(self, sample):
        return sample

    def image_augment(self, val, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple, 
                      rand_aug_idx: Optional[int], resample_mode: str = None):

        if isinstance(val, list) or isinstance(val, tuple):
            if self.aligned_captions:
                if rand_aug_idx is not None and rand_aug_idx < len(val):
                    val = val[rand_aug_idx]
                else:
                    # this means the video is not that lengthy, so we will use information from the first crop, as already decided during the dataset creation process
                    val = val[0]  # this will be the same for captions and transcriptions
            else:
                val = random.choice(val) if not self.no_aug else val[0]

        if isinstance(val, dict):
            # If each caption is saved as a dict, extract the string
            val = val["caption"]
        assert isinstance(val, str)

        return val

    def postprocess(self, sample):
        return sample

class TranscriptionTransform(AbstractTransform):

    def __init__(self, aligned_captions=True, no_aug=False):
        self.aligned_captions = aligned_captions
        self.no_aug = no_aug
        self.special_tokens = ["[SEC_1]", "[SEC_2]", "[SEC_3]", "[SEC_4]"]

    def load(self, path):
        # Caption can either be stored as .txt or .json.gz (in which case it's a list of dicts)
        if path.endswith('.txt'):
            sample = Path(path).read_text()
        elif path.endswith('.json'):
            with open(path, 'r') as f:
                sample = json.load(f)
        elif path.endswith('.json.gz'):
            with gzip.open(path, 'rb') as f:
                sample = json.load(f)
        return sample

    def finalize_transcription(self, sample):
        # it is a list, and
        return sample  # eos will be decided directly in masking function
        # return " ".join(f"{token} {text} [EOS]" for token, text in zip(self.special_tokens, sample))

    def preprocess(self, sample):
        return sample

    def image_augment(self, val, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple,
                      rand_aug_idx: Optional[int], resample_mode: str = None):

        if isinstance(val, list) or isinstance(val, tuple):
            if self.aligned_captions:
                if rand_aug_idx is not None and rand_aug_idx < len(val):
                    val = val[rand_aug_idx]
                else:
                    # this means the video is not that lengthy, so we will use information from the first crop, as already decided during the dataset creation process
                    val = val[0]  # this will be the same for captions and transcriptions
            else:
                val = random.choice(val) if not self.no_aug else val[0]

        if isinstance(val, dict):
            # If each caption is saved as a dict, extract the string
            val = val["caption"]
        # assert isinstance(val, str)

        return val

    def postprocess(self, sample):
        # now, let's convert it into a single sentence with special tokens
        sample = self.finalize_transcription(sample)
        return sample


class CaptionEmbTransform(AbstractTransform):

    def __init__(self, aligned_captions=True, no_aug=False):
        self.aligned_captions = aligned_captions
        self.no_aug = no_aug

    def load(self, path):
        if path.endswith('.npz'):
            sample = np.load(path)
            sample = {'emb': sample['emb'], 'mask_valid': sample['mask_valid']}
        else:
            raise ValueError(f"Invalid file format for caption embedding: {path}")
        return sample

    def preprocess(self, sample):
        return sample

    def image_augment(self, val, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple, 
                      rand_aug_idx: Optional[int], resample_mode: str = None):
        
        emb = val['emb']
        mask_valid = val['mask_valid'].astype(bool)
        num_sequences = emb.shape[0]

        if num_sequences > 1:
            if self.aligned_captions:
                if rand_aug_idx is None:
                    emb, mask_valid = emb[0], mask_valid[0]
                else:
                    if rand_aug_idx is not None and rand_aug_idx < num_sequences:
                        emb, mask_valid = emb[rand_aug_idx], mask_valid[rand_aug_idx]
                    else:
                        # this means the video is not that lengthy, so we will use information from the first crop, as already decided during the dataset creation process
                        # val = val[0]  # this will be the same for captions and transcriptions
                        emb, mask_valid = emb[0], mask_valid[0]
            else:
                if self.no_aug:
                    emb, mask_valid = emb[0], mask_valid[0]
                else:
                    rand_idx = random.randint(0, num_sequences - 1)
                    emb, mask_valid = emb[rand_idx], mask_valid[rand_idx]
        else:
            emb, mask_valid = emb[0], mask_valid[0]

        emb = emb[mask_valid] # Keep only valid embeddings

        return emb

    def postprocess(self, sample):
        return torch.tensor(sample)
      

class MetadataTransform(AbstractTransform):

    def __init__(self, 
                 special_vmin: int = 0, 
                 special_vmax: int = 999, 
                 shuffle: bool = True, 
                 random_trunc: bool = False, 
                 return_chunks: bool = True, 
                 return_raw: bool = False,
                 image_dim_bin_size: int = 32,):
        """Metadata transform that takes in a metadata dictionary and converts 
        it into a string, or list of strings (for chunked span masking).
        Uses special tokens v1 to denote metadata types, and v0 for their values.

        Args:
            special_vmin: Minimum value for special tokens
            special_vmax: Maximum value for special tokens
            shuffle: Whether to shuffle the metadata order
            random_trunc: Whether to randomly truncate the returned metadata
            return_chunks: Whether to return a list of strings (for chunked span masking),
                or a single string with all metadata concatenated
            return_raw: Whether to return the raw metadata dictionary
        """
        self.special_vmin = special_vmin
        self.special_vmax = special_vmax
        self.shuffle = shuffle
        self.random_trunc = random_trunc
        self.return_chunks = return_chunks
        self.return_raw = return_raw
        self.image_dim_bin_size = image_dim_bin_size

        # Explicit map to make sure that additional entries do not change existing IDs
        # TODO: Make this work with other text tokenizers
        self.metadata_id_map = {
            'original_width': 'v1=0',
            'original_height': 'v1=1',
            'caption_n_chars': 'v1=2',
            'caption_n_words': 'v1=3',
            'caption_n_sentences': 'v1=4',
            'n_humans': 'v1=5',
            'n_sam_instances': 'v1=6',
            'n_coco_instances': 'v1=7',
            'coco_instance_diversity': 'v1=8',
            'colorfulness': 'v1=9',
            'brightness': 'v1=10',
            'contrast': 'v1=11',
            'saturation': 'v1=12',
            'entropy': 'v1=13',
            'walkability': 'v1=14',
            'objectness': 'v1=15',
            'semantic_diversity': 'v1=16',
            'geometric_complexity': 'v1=17',
            'occlusion_score': 'v1=18',
            'watermark_score': 'v1=19',
            'aesthetic_score': 'v1=20',
        }
        self.id_metadata_map = {v: k for k, v in self.metadata_id_map.items()}

        # Image-dimension modalities are binned into 32 bins
        self.image_dim_modalities = ['original_height', 'original_width']

        # Integer modalities that don't undergo any scaling (except for truncation)
        self.metadata_int_modalities = [
            'caption_n_chars', 'caption_n_words', 'caption_n_sentences', 
            'n_humans', 'n_sam_instances', 'n_coco_instances', 
            'coco_instance_diversity', 'semantic_diversity', 
        ]

        # Bin boundaries for manually defined metadata modalities.
        # Lowest and highest bin boundaries are implicitly set to -inf and +inf
        self.metadata_manual_bins = {
            'watermark_score': [0.5],
            'aesthetic_score': [4.5, 5.5],
        }

        # All other float or integer modalities that are binned into a defined number of bins
        # Dictionary entries are (vmin, vmax, num_bins)
        self.metadata_min_max_bins = {
            'colorfulness': (0, 150, 50),
            'brightness': (0, 255, 50),
            'contrast': (0, 127, 50),
            'saturation': (0, 255, 50),
            'entropy': (0, 10, 50),
            'walkability': (0, 1, 50),
            'objectness': (0, 1, 50),
            'geometric_complexity': (0, 0.75, 50),
            'occlusion_score': (0, 0.25, 50),
        }

    def image_dim_to_string(self, metadata, key, bin_size=32):
        value = metadata[key] // bin_size
        value = max(self.special_vmin, min(value, self.special_vmax))
        return f"{self.metadata_id_map[key]} v0={value}"

    def int_metadata_to_string(self, metadata, key):
        value = max(self.special_vmin, min(metadata[key], self.special_vmax))
        return f"{self.metadata_id_map[key]} v0={value}"

    def float_metadata_to_string(self, metadata, key, vmin, vmax, bins):
        value = max(vmin, min(metadata[key], vmax))
        value = (value - vmin) / (vmax - vmin)
        value = int(value * (bins-1))
        return f"{self.metadata_id_map[key]} v0={value}"
    
    def manual_bin_metadata_to_string(self, metadata, key):
        value = metadata[key]
        bin_idx = 0
        for bin_value in self.metadata_manual_bins[key]:
            if value < bin_value:
                break
            bin_idx += 1
        return f"{self.metadata_id_map[key]} v0={bin_idx}"
    
    def metadata_to_string(self, metadata, keys: List[str] = None):
        keys = list(metadata.keys()) if keys is None else keys

        if self.shuffle:
            # Randomly shuffle
            random.shuffle(keys)
        if self.random_trunc:
            # Randomly truncate
            keys = keys[:random.randint(1,len(keys))]

        metadata_strings = []
        
        for key in keys:
            if key in self.image_dim_modalities:
                # Image dimension modalities
                metadata_str = self.image_dim_to_string(metadata, key, bin_size=self.image_dim_bin_size)
            elif key in self.metadata_int_modalities:
                # Integer modalities that don't undergo any scaling
                metadata_str = self.int_metadata_to_string(metadata, key)
            elif key in self.metadata_manual_bins:
                # Metadata modalities for which bin boundaries are manually defined
                metadata_str = self.manual_bin_metadata_to_string(metadata, key)
            else:
                # All other modalities
                vmin, vmax, bins = self.metadata_min_max_bins[key]
                metadata_str = self.float_metadata_to_string(metadata, key, vmin, vmax, bins)

            metadata_strings.append(metadata_str)

        if self.return_chunks:
            return metadata_strings
        else:
            return ' '.join(metadata_strings)

    def load(self, path):
        with open(path, 'r') as f:
            sample = json.load(f)
        return sample

    def preprocess(self, sample):
        return sample

    def image_augment(self, val, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple, 
                      rand_aug_idx=None, resample_mode: str = None):
        return val

    def postprocess(self, metadata):
        if self.return_raw:
            return metadata
        metadata_str = self.metadata_to_string(metadata)
        return metadata_str
    

class HumanPoseTransform(AbstractTransform):

    def __init__(self, coord_bins=1000, only_pose=False, return_raw=False):
        self.coord_bins = coord_bins
        self.return_raw = return_raw
        self.only_pose = only_pose

    def convert_humanpose_instance(self, instances, only_pose=False):
        """Convert instances dict to list of lists where each list takes the form:
        [human, xmin xmax ymin ymax global val1 val2 ... val10 pose val1 val2 ... val 207 shape val1 val2 ... val10 camera val1 val2 val3 val4] 
        Like for bounding boxes, xmin, ymin, xmax, and ymax map to v0, v1, v2, and v3 respectively.
        """
        if only_pose: # used for tokenizer training for pose
            if len(instances) == 0:
                return torch.zeros(207)
            else:
                return torch.from_numpy(np.array(instances['pred_smpl_params']['body_pose'][0]).flatten()).float()
        if len(instances) == 0: #empty, i.e. there are no humans
            return 'none'
        
        for k in instances:
            if k!='pred_smpl_params':
                instances[k] = torch.from_numpy(np.array(instances[k]))

        smpl_params = (instances['pred_smpl_params'])

        for k in smpl_params:
            smpl_params[k] = torch.from_numpy(np.array(smpl_params[k]))

        total_num_instances = len(instances['bbox_xyxy'])
        instances_converted = []
        for ii in range(total_num_instances):
            instances_converted.append(['human'] + (np.array(instances['bbox_xyxy'][ii]).flatten().tolist()) + ['global'] + (np.array(instances['pred_smpl_params']['global_orient'][ii]).flatten().tolist()) + ['pose'] + (instances['pose_tokenized'][ii].flatten().tolist()) + ['shape'] + (instances['pred_smpl_params']['betas'][ii].flatten().tolist()) + ['camera'] + (instances['pred_cam'][ii].flatten().tolist()))
        return instances_converted

    def humanposes_crop_and_resize(self, humanposes: List[Tuple], crop_coords: Tuple, orig_size: Tuple,):
        """Crop and resize human poses (and their bounding boxes)
        """
        orig_height, orig_width = orig_size
        top, left, h, w = crop_coords

        humanposes_converted_resized = []
        for instance in humanposes:
            bbox_curr = instance[1:5]
            bbox_curr = np.array(bbox_curr)
            bbox_curr[0::2] = bbox_curr[0::2] / orig_width
            bbox_curr[1::2] = bbox_curr[1::2] / orig_height

            xmin, ymin, xmax, ymax = left, top, left + w, top + h
            bbox_curr = A.bbox_crop(bbox_curr, x_min=xmin, y_min=ymin, x_max=xmax, y_max=ymax, rows=orig_height,
                                        cols=orig_width) 
            bbox_curr = np.array(bbox_curr)
            if np.all(bbox_curr[1::2]<0) or np.all(bbox_curr[0::2]<0): #bbox is out of range, remove it 
                continue
            if np.all(bbox_curr[1::2]>1.0) or np.all(bbox_curr[0::2]>1.0): #bbox is out of range, remove it 
                continue
            bbox_curr = np.clip(bbox_curr, a_min=0, a_max=1.)

            instance[1:5] = bbox_curr
            humanposes_converted_resized.append(instance)

        # now return all instances, or none if there is no instance
        if len(humanposes_converted_resized)>0:
            pass
        else: #no valid masks remains
            return 'none'

        humanpose_returned = humanposes_converted_resized

        return humanpose_returned

    def convert_humanposes_to_string(self, all_humanposes: List[Tuple]):
        """Convert humanposes to a string
           range of global orientation: [-1, 1]
           range of object pose: [-1, 1]
           range of shape (betas): [-3, 3] 
           range of camera: [-1, 19]
        """
        bins = self.coord_bins

        instance_final_all = ''

        for humanposes in all_humanposes:
            human = humanposes[0]
            bboxes = humanposes[1:5]
            glob = humanposes[5]
            global_orient = np.array(humanposes[6:15])
            pose = humanposes[15]
            pose_params = np.array(humanposes[16:24]) 
            shape = humanposes[24] 
            shape_params = np.array(humanposes[25:35]) 
            camera = humanposes[35] 
            camera_params = np.clip(np.array(humanposes[36:]), a_min=-1., a_max=19.) 

            bboxes_new = [
                    f"v0={round(bboxes[0] * (bins - 1))}",
                    f"v1={round(bboxes[1] * (bins - 1))}",
                    f"v2={round(bboxes[2] * (bins - 1))}",
                    f"v3={round(bboxes[3] * (bins - 1))}"]

            global_orient = 499.5*global_orient
            global_orient_new = []
            for ii in range(len(global_orient)):
                global_orient_curr =  f"v0={round(global_orient[ii]+499.5)}"
                global_orient_new.append(global_orient_curr)

            pose_params_new = []
            for ii in range(len(pose_params)):
                if pose_params[ii]<512: 
                    pose_params_curr =  f"v0={round(pose_params[ii])}"
                else: 
                    pose_params_curr =  f"v1={round(pose_params[ii] - 512)}"
                pose_params_new.append(pose_params_curr)

            shape_params = 166.5*shape_params
            shape_params_new = []
            for ii in range(len(shape_params)):
                shape_params_curr =  f"v0={round(shape_params[ii]+499.5)}"
                shape_params_new.append(shape_params_curr)

            camera_params = 49.95*camera_params
            camera_params_new = []
            for ii in range(len(camera_params)):
                camera_params_curr =  f"v0={round(camera_params[ii]+49.95)}"
                camera_params_new.append(camera_params_curr)
            
            #randomly shuffle everything except bbox part of the sequence
            all_strings = [[pose]+pose_params_new, [glob] + global_orient_new, [camera] + camera_params_new, [shape] + shape_params_new ]
            rand_perm = torch.randperm(4)
            instance_final = [human] + bboxes_new + all_strings[rand_perm[0]] + all_strings[rand_perm[1]] + all_strings[rand_perm[2]] + all_strings[rand_perm[3]]
            
        
            instance_final = ', '.join(instance_final)
            instance_final = instance_final.replace(",", "")
            instance_final_all = instance_final_all + instance_final + ' '

        return instance_final_all 

    def load(self, path):
        with open(path, 'r') as f:
            sample = json.load(f)

        return sample

    def preprocess(self, sample):
        instances = sample 
        instances = self.convert_humanpose_instance(instances, only_pose=self.only_pose)
        return instances

    def image_augment(self, humanposes: List[Tuple], crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple, 
                      rand_aug_idx=None, resample_mode: str = None):
        if humanposes=='none' or self.only_pose:
            return humanposes
        humanposes = self.humanposes_crop_and_resize(humanposes, crop_coords, orig_size)
        return humanposes

    def postprocess(self, humanposes):
        if humanposes=='none' or self.only_pose:
            return humanposes if not self.return_raw else []
        if self.return_raw:
            return humanposes
        humanposes = self.convert_humanposes_to_string(humanposes)
        return humanposes


class VideoHumanPoseTransform(AbstractTransform):

    def __init__(self, coord_bins=1000, only_pose=False, return_raw=False):
        self.coord_bins = coord_bins
        self.return_raw = return_raw
        self.only_pose = only_pose
        self.total_max_instances = 7 # per frame poses

    def create_chunks(self, chunk_sizes, data_list):
        """Using list comprehension with accumulate"""
        indices = [0] + list(accumulate(chunk_sizes))
        all_poses = []
        for i in range(len(chunk_sizes)):
            start, end = indices[i], indices[i + 1]
            all_poses.extend(data_list[start:min(start + self.total_max_instances, end)])
        updated_chunk_sizes = [min(self.total_max_instances, current_frame_instances) for current_frame_instances in chunk_sizes]
        return all_poses, updated_chunk_sizes

    def convert_humanpose_instance(self, instances):
        """Convert instances dict to list of lists where each list takes the form:
        [human, xmin xmax ymin ymax global val1 val2 ... val10 pose val1 val2 ... val 207 shape val1 val2 ... val10 camera val1 val2 val3 val4]
        Like for bounding boxes, xmin, ymin, xmax, and ymax map to v0, v1, v2, and v3 respectively.

        update: instances will be now dict of frames and each frame will be dict of lists
        """
        all_frame_poses = []
        per_frame_num_instances = []

        for single_frame_instances in instances:
            # single_frame_instances = instances[single_frame_key]

            if single_frame_instances == {}:  # empty, i.e. there are no humans in the frame
                per_frame_num_instances.append(0)
                continue
            # convert the dict of lists to dict of tensors
            for k in single_frame_instances:
                if k != 'pred_smpl_params':
                    single_frame_instances[k] = torch.from_numpy(np.array(single_frame_instances[k]))

            smpl_params = (single_frame_instances['pred_smpl_params'])

            for k in smpl_params:
                smpl_params[k] = torch.from_numpy(np.array(smpl_params[k]))

            total_num_instances = len(single_frame_instances['bbox_xyxy'])

            # now convert them into strings
            instances_converted = []
            for ii in range(total_num_instances):
                instances_converted.append(
                    ['human'] + (np.array(single_frame_instances['bbox_xyxy'][ii]).flatten().tolist()) + ['global'] + (
                        np.array(single_frame_instances['pred_smpl_params']['global_orient'][ii]).flatten().tolist()) + ['pose'] + (
                        single_frame_instances['pose_tokenized'][ii].flatten().tolist()) + ['shape'] + (
                        single_frame_instances['pred_smpl_params']['betas'][ii].flatten().tolist()) + ['camera'] + (
                        single_frame_instances['pred_cam'][ii].flatten().tolist()))
            all_frame_poses.extend(instances_converted)
            per_frame_num_instances.append(total_num_instances)

        return all_frame_poses, per_frame_num_instances

    def humanposes_crop_and_resize(self, humanposes: List[Tuple], per_frame_pose_count, crop_coords: Tuple, orig_size: Tuple):
        """Crop and resize human poses (and their bounding boxes)
        """
        orig_height, orig_width = orig_size
        _, _, top, left, h, w = crop_coords

        humanposes_converted_resized = []
        filtered_bbox_indices = []
        for idx, instance in enumerate(humanposes):
            bbox_curr = instance[1:5]
            bbox_curr = np.array(bbox_curr)
            bbox_curr[0::2] = bbox_curr[0::2] / orig_width
            bbox_curr[1::2] = bbox_curr[1::2] / orig_height

            xmin, ymin, xmax, ymax = left, top, left + w, top + h
            bbox_curr = A.bbox_crop(bbox_curr, x_min=xmin, y_min=ymin, x_max=xmax, y_max=ymax, rows=orig_height,
                                    cols=orig_width)
            bbox_curr = np.array(bbox_curr)
            if np.all(bbox_curr[1::2] < 0) or np.all(bbox_curr[0::2] < 0):  # bbox is out of range, remove it
                filtered_bbox_indices.append(idx)
                continue
            if np.all(bbox_curr[1::2] > 1.0) or np.all(bbox_curr[0::2] > 1.0):  # bbox is out of range, remove it
                filtered_bbox_indices.append(idx)
                continue
            bbox_curr = np.clip(bbox_curr, a_min=0, a_max=1.) # TODO: let's perform explicit ablation on this.

            instance[1:5] = bbox_curr
            humanposes_converted_resized.append(instance)
        filtered_bbox_indices = np.array(filtered_bbox_indices)
        if len(filtered_bbox_indices) != 0:
            pose_to_frame = np.repeat(np.arange(len(per_frame_pose_count)), per_frame_pose_count)
            per_frame_num_instances = per_frame_pose_count - np.bincount(pose_to_frame[filtered_bbox_indices],
                                                                            minlength=len(per_frame_pose_count))
        else:
            per_frame_num_instances = per_frame_pose_count


        assert len(humanposes_converted_resized) == sum(per_frame_num_instances), "something wrong in the code, history of boxes per frames not tracked correctly"
        assert  not np.any(per_frame_num_instances < 0), f"there cannot be negative instances left after cropping and resizing {per_frame_num_instances}"
        # No need to resize, bounding boxes in albumentations format are scale invariant
        return humanposes_converted_resized, per_frame_num_instances.tolist()

    def convert_humanposes_to_string(self, all_humanposes_with_instance_count,
                                     global_orient_bins=1000,
                                     shape_bins=1000,
                                     camera_bins=1000):
        """Convert humanposes to a string with configurable bin resolution

        Args:
            all_humanposes: List of tuples containing human pose data
            global_orient_bins: Number of bins for global orientation quantization (default: 1000)
            shape_bins: Number of bins for shape parameters quantization (default: 1000)
            camera_bins: Number of bins for camera parameters quantization (default: 1000)

        Parameter ranges:
           range of global orientation: [-1, 1]
           range of object pose: [-1, 1]
           range of shape (betas): [-3, 3] -> now it is [-6.94, 5.54]
           range of camera: [-1, 19] -> it is actually [-0.88, 10.5]
        """
        all_humanposes, per_frame_num_instances = all_humanposes_with_instance_count
        bbox_bins = self.coord_bins  # Keep bbox bins unchanged for ablation
        # print("shape bins: k600 ", shape_bins)
        instance_final_all_list = []

        for humanposes in all_humanposes:
            human = humanposes[0]
            bboxes = humanposes[1:5]
            glob = humanposes[5]
            global_orient = np.array(humanposes[6:15])
            pose = humanposes[15]
            pose_params = np.array(humanposes[16:24])
            shape = humanposes[24]
            shape_params = np.array(humanposes[25:35])
            camera = humanposes[35]
            camera_params = np.clip(np.array(humanposes[36:]), a_min=-1., a_max=13.)

            # Bounding boxes (unchanged)
            bboxes_new = [
                f"v0={round(bboxes[0] * (bbox_bins - 1))}",
                f"v1={round(bboxes[1] * (bbox_bins - 1))}",
                f"v2={round(bboxes[2] * (bbox_bins - 1))}",
                f"v3={round(bboxes[3] * (bbox_bins - 1))}"]

            # Global orientation: [-1, 1] -> [0, global_orient_bins-1]
            global_orient_scale = (global_orient_bins - 1) / 2.0
            global_orient_offset = global_orient_scale
            global_orient_quantized = global_orient * global_orient_scale + global_orient_offset
            global_orient_new = []
            for ii in range(len(global_orient_quantized)):
                global_orient_curr = f"v0={round(global_orient_quantized[ii])}"
                global_orient_new.append(global_orient_curr)

            # Pose parameters (unchanged - keeping original logic)
            pose_params_new = []
            for ii in range(len(pose_params)):
                if pose_params[ii] < 512:
                    pose_params_curr = f"v0={round(pose_params[ii])}"
                else:
                    pose_params_curr = f"v1={round(pose_params[ii] - 512)}"
                pose_params_new.append(pose_params_curr)

            # Shape parameters: [-6.94, 5.54] -> [0, shape_bins-1] (13 unit range)
            shape_range = 12.48  # from -6.94 to 5.54
            shape_scale = (shape_bins - 1) / shape_range
            shape_offset = 6.94 * shape_scale  # offset for -6.94 minimum
            shape_quantized = shape_params * shape_scale + shape_offset
            shape_params_new = []
            for ii in range(len(shape_quantized)):
                shape_params_curr = f"v0={round(shape_quantized[ii])}"
                shape_params_new.append(shape_params_curr)

            # Camera parameters: [-0.88, 10.5] -> [0, camera_bins-1] (11.38 unit range)
            camera_range = 11.38  # from -0.88 to 10.5
            camera_scale = (camera_bins - 1) / camera_range
            camera_offset = 0.88 * camera_scale  # offset for -0.88 minimum
            camera_quantized = camera_params * camera_scale + camera_offset
            camera_params_new = []
            for ii in range(len(camera_quantized)):
                camera_params_curr = f"v0={round(camera_quantized[ii])}"
                camera_params_new.append(camera_params_curr)

            # Randomly shuffle everything except bbox part of the sequence
            all_strings = [[pose] + pose_params_new, [glob] + global_orient_new,
                           [camera] + camera_params_new, [shape] + shape_params_new]
            rand_perm = torch.randperm(4)
            instance_final = ([human] + bboxes_new +
                              all_strings[rand_perm[0]] + all_strings[rand_perm[1]] +
                              all_strings[rand_perm[2]] + all_strings[rand_perm[3]])
            instance_final = ', '.join(instance_final)
            instance_final = instance_final.replace(",", "")
            instance_final_all_list.append(instance_final)

        return (instance_final_all_list, per_frame_num_instances)

    def load(self, path):
        with open(path, 'r') as f:
            sample = json.load(f)

        return sample

    def preprocess(self, sample):
        instances = sample
        # instances = self.convert_humanpose_instance(instances, only_pose=self.only_pose) # will do this once we filter the frames for the crop
        return instances

    def image_augment(self, humanposes_full_dict, crop_coords: Tuple, flip: bool, orig_size: Tuple,
                      target_size: Tuple,
                      rand_aug_idx=None, resample_mode: str = None):

        # STEP 1: extract pose info for frames that are required in that crop
        starting_time, ending_time, i, j, h, w = crop_coords  # e.g., 3, 7, 124, 0, 320, 320, 0
        video_info = humanposes_full_dict['video_info']
        fps = video_info['fps']

        selected_frame_indices = np.linspace(
            starting_time * fps,
            ending_time * fps,
            17,
            dtype=np.int32
        )
        selected_frame_indices[-1] = selected_frame_indices[-1] - 1

        # Update original size from video info
        orig_size = (video_info['height'], video_info['width'])

        # now let's retrieve only those frames relevant for this given crop setting:
        human_poses = humanposes_full_dict['frames']

        single_crop_poses = [human_poses[f"{idx}"] for idx in selected_frame_indices]
        if single_crop_poses == [{}] * len(single_crop_poses):
            return 'none'
        # Now, let's convert each of them into a string, as done by 4M by default

        string_conveted_humanposes, per_frame_num_instances = self.convert_humanpose_instance(
            single_crop_poses
        )

        # now, lets crop, all poses, so that poses lying outsides are skipped
        humanposes, per_frame_num_instances= self.humanposes_crop_and_resize(string_conveted_humanposes, np.array(per_frame_num_instances), crop_coords, orig_size)
        if sum(per_frame_num_instances) == 0:
            return 'none'
        # Reorganize into per-frame chunks and apply ordering
        per_frame_detections, num_instances_updated = self.create_chunks(per_frame_num_instances, humanposes)
        assert len(per_frame_detections) == sum(num_instances_updated), "something wrong in the code, history of boxes per frames not tracked correctly"
        return (per_frame_detections, num_instances_updated)

    def postprocess(self, humanposes):
        if humanposes == 'none' or self.only_pose:
            return (humanposes, [0] * 17) if not self.return_raw else []
        if self.return_raw:
            return humanposes
        humanposes = self.convert_humanposes_to_string(humanposes)
        return humanposes


class ColorPaletteTransform(AbstractTransform):

    def __init__(self, coord_bins=1000, return_raw=False):
        self.coord_bins = coord_bins
        self.return_raw = return_raw

    def convert_palette_instance(self, instances):
        """Convert colors to v0= v0= ...
        """
        length = random.randint(1,7)
        instances_converted = np.array(instances[0][str(length)]).flatten().tolist()
        return instances_converted

    def palette_hflip(self, palettes: List[Tuple], image_size: Tuple, flip: bool):

        return palettes

    def convert_palettes_to_string(self, all_palettes: List[Tuple]):
        """Convert palettes to a string
        """

        colors = []
        len_palettes = len(all_palettes)
        colors.append(f"v1={round(len_palettes/3)}") # start with the length of the color palette to avoid confusion
        for ii in range(len(all_palettes)):
            color_new = f"v0={round(all_palettes[ii])}"
            colors.append(color_new)
        
        instance_final_all = colors
        instance_final_all = ', '.join(instance_final_all)
        instance_final_all = instance_final_all.replace(",", "")

        return instance_final_all 

    def load(self, path):
        with open(path, 'r') as f:
            sample = json.load(f)
        return sample

    def preprocess(self, sample):
        if self.return_raw:
            return sample
        instances = sample 
        instances = self.convert_palette_instance(instances)
        return instances

    def image_augment(self, palettes: List[Tuple], crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple, 
                      rand_aug_idx=None, resample_mode: str = None):
        return palettes

    def postprocess(self, palettes):
        if self.return_raw:
            return palettes
        palettes = self.convert_palettes_to_string(palettes)
        return palettes
    

class SAMInstanceTokTransform(AbstractTransform):

    def __init__(self, image_size=224, points_per_side=7, point_order='random'):
        self.H, self.W = to_2tuple(image_size)
        self.points_per_h, self.points_per_w = to_2tuple(points_per_side)
        assert point_order in ['random', 'grid']
        self.point_order = point_order

    def get_query_points(self):
        if self.point_order == 'grid':
            # Create and cache grid query points
            if not hasattr(self, 'grid_query_points'):
                y, x = np.meshgrid(np.linspace(0, self.H, self.points_per_h + 2)[1:-1], np.linspace(0, self.W, self.points_per_w + 2)[1:-1])
                grid = np.stack((x, y), axis=2).astype(np.int32)
                self.grid_query_points = grid.reshape(-1, 2)
            return self.grid_query_points
        elif self.point_order == 'random':
            # Randomly sample query points
            y = np.random.randint(0, self.H, self.points_per_h)
            x = np.random.randint(0, self.W, self.points_per_w)
            return np.concatenate((x[:,None], y[:,None]), axis=1)
        else:
            raise ValueError(f"Query point order mode {self.point_order} is not supported.")

    def get_target_tokens(self, sample, query_points):
        instances_coords = [coords[0] for coords in sample['points']]
        tokens = sample['token_ids']
        bboxes = sample['bbox']
        
        instance_tokens_per_qpoint = dict()
        for point in query_points:
            point = (int(point[0].item()), int(point[1].item()))
            instance_tokens_per_qpoint[point] = []
            for i, (coords, tok, bbox) in enumerate(zip(instances_coords, tokens, bboxes)):
                # Calculate the distance from the query point to the instance
                distance = cv2.pointPolygonTest(coords, point, measureDist=True)
                # If the query point is inside the instance, add its corresponding token
                if distance >= 0:
                    instance_tokens_per_qpoint[point].append((tok, bbox))
        
        return instance_tokens_per_qpoint

    def convert_target_tokens_to_string(self, target_tokens):
        result_text = []
        query_points = list(target_tokens.keys())
        # Randomly shuffle query points order (mainly for grid order)
        random.shuffle(query_points)
        for point in query_points:
            
            # Add query point coordinates to the string
            result_text.append('point')
            result_text.append(f'v0={point[1]}')
            result_text.append(f'v1={point[0]}')
            
            # Randomly shuffle the order of instance tokens per query point
            random.shuffle(target_tokens[point])
            if len(target_tokens[point]) == 0:
                # If no instances tokens are found, add 'none' to the string
                result_text.append('none')
            else:
                for tok, bbox in target_tokens[point]:
                    result_text.append(f'polygon')
                    
                    # Add bounding box coordinates to the string
                    ymin, xmin, ymax, xmax = bbox.astype(np.int32)
                    result_text.extend([
                        f'v0={xmin}',
                        f'v1={ymin}',
                        f'v2={xmax}',
                        f'v3={ymax}',
                    ])
                    
                    # Add instance tokens ids to the string
                    for idx in tok.tolist():
                        if idx < 512:
                            result_text.append(f'v0={idx}')
                        else:
                            result_text.append(f'v1={idx - 512}')
        
        return " ".join(result_text)

    def load(self, path):
        sample = np.load(path, allow_pickle=True)
        return sample

    def preprocess(self, sample):
        for s in sample:
            s['token_ids'] = s['token_ids'].astype(np.int32)
        return sample

    def image_augment(self, v, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple,
                      rand_aug_idx: Optional[int], resample_mode: str = None):
        if rand_aug_idx is None:
            raise ValueError("Crop settings / augmentation index are missing but a pre-tokenized modality is being used")
        v = v[rand_aug_idx]
        return v

    def postprocess(self, sample):
        query_points = self.get_query_points()
        target_tokens = self.get_target_tokens(sample, query_points)
        final_string = self.convert_target_tokens_to_string(target_tokens)
        return final_string


class CropSettingsTransform(AbstractTransform):

    def load(self, path):
        sample = np.load(path)
        return sample

    def preprocess(self, sample):
        raise NotImplementedError("CropSettingsTransform does not support preprocessing")

    def image_augment(self, val, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple, 
                      rand_aug_idx: Optional[int], resample_mode: str = None):
        raise NotImplementedError("CropSettingsTransform is not meant to be used for image augmentation")

    def postprocess(self, sample):
        raise NotImplementedError("CropSettingsTransform does not support postprocessing")


class IdentityTransform(AbstractTransform):

    def load(self, path):
        raise NotImplementedError("IdentityTransform does not support loading")

    def preprocess(self, sample):
        return sample

    def image_augment(self, val, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple, 
                      rand_aug_idx: Optional[int], resample_mode: str = None):
        return val

    def postprocess(self, sample):
        return sample


class JSONTransform(AbstractTransform):

    def load(self, path):
        if path.endswith('.json'):
            with open(path, 'r') as f:
                sample = json.load(f)
        elif path.endswith('.json.gz'):
            with gzip.open(path, 'rb') as f:
                sample = json.load(f)
        return sample

    def preprocess(self, sample):
        return sample

    def image_augment(self, val, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple, 
                      rand_aug_idx: Optional[int], resample_mode: str = None):
        return val

    def postprocess(self, sample):
        return sample