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

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)


# 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

    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


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 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 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:
                val = val[0] if rand_aug_idx is None else val[rand_aug_idx]
            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 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:
                    emb, mask_valid = emb[rand_aug_idx], mask_valid[rand_aug_idx]
            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 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