# 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