# 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 os import numpy as np import torch import torchvision.transforms.functional as TF from einops import rearrange import textwrap import matplotlib.pyplot as plt from PIL import Image import cv2 from itertools import groupby # For visualizing CLIP feature maps from sklearn.decomposition import PCA # Detectron2 for semantic segmentation visualizations try: from detectron2.utils.visualizer import ColorMode, Visualizer from detectron2.data import MetadataCatalog coco_metadata = MetadataCatalog.get("coco_2017_val_panoptic") USE_DETECTRON = True except Exception as e: print(e) print("Detectron2 can be used for semseg visualizations. Please install detectron2 to use this feature, or plotting will fall back to matplotlib.") USE_DETECTRON = False from fourm.data.modality_transforms import get_transform_key, get_transform_resolution, MetadataTransform from fourm.utils.data_constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, COCO_SEMSEG_NUM_CLASSES from fourm.utils import denormalize, get_sentinel_to_id_mapping, merge_span_masking from fourm.utils.generation import unbatch device = 'cuda' if torch.cuda.is_available() else 'cpu' def tensor_to_images(tensor): """ Converts a (B C H W) tensor to numpy arrays. If B = 1, the tensor is unbatched and converted to a single image. If C = 1, the channel dimension is removed. Args: tensor (torch.Tensor): Tensor to convert to images. """ B, C, H, W = tensor.shape if B == 1: img = rearrange(unbatch(tensor), "c h w -> h w c") else: img = rearrange(tensor, "b c h w -> b h w c") if C == 1: img = img[..., 0] return img.detach().cpu().numpy() def pca_visualize(features, n_components=3): """ Visualizes a feature map using PCA. Args: features (torch.Tensor): CxHxW feature map to visualize. n_components (int): Number of PCA components to use. """ C, H, W = features.shape features_flat = rearrange(features.float(), 'c h w -> (h w) c').detach().cpu().numpy() pca = PCA(n_components=n_components) img_pca = rearrange(pca.fit_transform(features_flat), '(h w) c -> h w c', h=H, w=W) img_pca = (img_pca - img_pca.min()) / (img_pca.max() - img_pca.min()) return img_pca def np_squeeze(array, axis=0): """ Squeeses a numpy array along a given axis if that axis is one-dimensional. Otherwise, it returns the same array. Args: array (numpy.ndarray): Array to squeeze. axis (int): Axis to squeeze. """ if array.shape[axis] == 1: return np.squeeze(array, axis=axis) else: return array def decode_input_rgb(mod_dict, key='rgb'): """ Decodes (denormalizes) an RGB image from a model dictionary. Args: mod_dict (dict): Model output dictionary. key (str): Key of the RGB modality to decode. """ img = denormalize(mod_dict[key]['tensor']) return tensor_to_images(img) def decode_tok_rgb(mod_dict, tokenizers, key='tok_rgb', image_size=224, patch_size=16, t=25, verbose=False): """ Decodes a sequence of RGB tokens from a model dictionary into an RGB image. Args: mod_dict (dict): Model output dictionary. tokenizers (dict): Dictionary of tokenizers. key (str): Key of the tokenized RGB modality to decode. image_size (int): Size of the image. patch_size (int): Size of the patches. t (int): Number of timesteps to decode using the tokenizer diffusion model (if applicable). verbose (bool): Whether to print the decoding progress. """ img_tok = rearrange(mod_dict[key]['tensor'], "b (nh nw) -> b nh nw", nh=image_size//patch_size, nw=image_size//patch_size) rec = tokenizers[get_transform_key(key)].decode_tokens(img_tok, timesteps=t, image_size=image_size, verbose=verbose) rec = denormalize(rec, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)).clamp(0, 1) return tensor_to_images(rec) def decode_tok_rgb_controlnet(mod_dict, tokenizers, key='tok_rgb', image_size=224, patch_size=16, t=25, guidance_scale=2.5, cond_scale=0.8, verbose=False): """ Decodes a sequence of RGB tokens from a model dictionary into an RGB image using a ControlNet. Args: mod_dict (dict): Model output dictionary. tokenizers (dict): Dictionary of tokenizers. Needs to contain the key 'controlnet'. key (str): Key of the tokenized RGB modality to decode. image_size (int): Size of the image. patch_size (int): Size of the patches. t (int): Number of timesteps to decode using the ControlNet. guidance_scale (float): Classifier-free guidance scale. cond_scale (float): ControlNet conditioning scale. verbose (bool): Whether to print the decoding progress. """ img_tok = rearrange(mod_dict[key]['tensor'], "b (nh nw) -> b nh nw", nh=image_size//patch_size, nw=image_size//patch_size) rec = tokenizers['controlnet'].decode_tokens( img_tok, timesteps=t, guidance_scale=guidance_scale, cond_scale=cond_scale, verbose=verbose ) rec = tokenizers['controlnet'].vae_decode(rec) rec = denormalize(rec, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)).clamp(0, 1) return tensor_to_images(rec) def decode_tok_normal(mod_dict, tokenizers, key='tok_normal', image_size=224, patch_size=16, t=25, verbose=False): """ Decodes a sequence of surface normal tokens from a model dictionary into an RGB image. Args: mod_dict (dict): Model output dictionary. tokenizers (dict): Dictionary of tokenizers. key (str): Key of the tokenized normal modality to decode. image_size (int): Size of the image. patch_size (int): Size of the patches. t (int): Number of timesteps to decode using the tokenizer diffusion model (if applicable). verbose (bool): Whether to print the decoding progress. """ img_tok = rearrange(mod_dict[key]['tensor'], "b (nh nw) -> b nh nw", nh=image_size//patch_size, nw=image_size//patch_size) rec = tokenizers[get_transform_key(key)].decode_tokens(img_tok, timesteps=t, image_size=image_size, verbose=verbose) rec = denormalize(rec, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)).clamp(0, 1) return tensor_to_images(rec) def decode_tok_canny_edge(mod_dict, tokenizers, key='tok_canny_edge', image_size=224, patch_size=16, t=10, verbose=False): """ Decodes a sequence of Canny edges tokens from a model dictionary into an RGB image. Args: mod_dict (dict): Model output dictionary. tokenizers (dict): Dictionary of tokenizers. key (str): Key of the tokenized Canny edges modality to decode. image_size (int): Size of the image. patch_size (int): Size of the patches. t (int): Number of timesteps to decode using the tokenizer diffusion model (if applicable). verbose (bool): Whether to print the decoding progress. """ img_tok = rearrange(mod_dict[key]['tensor'], "b (nh nw) -> b nh nw", nh=image_size//patch_size, nw=image_size//patch_size) rec = tokenizers[get_transform_key(key)].decode_tokens(img_tok, timesteps=t, image_size=image_size, verbose=verbose) rec = (0.5*(rec+1)).clamp(0, 1) return tensor_to_images(rec) def decode_tok_sam_edge(mod_dict, tokenizers, key='tok_sam_edge', image_size=224, patch_size=16, t=10, verbose=False): """ Decodes a sequence of SAM edges from a model dictionary into an RGB image. Args: mod_dict (dict): Model output dictionary. tokenizers (dict): Dictionary of tokenizers. key (str): Key of the tokenized SAM edges modality to decode. image_size (int): Size of the image. patch_size (int): Size of the patches. t (int): Number of timesteps to decode using the tokenizer diffusion model (if applicable). verbose (bool): Whether to print the decoding progress. """ img_tok = rearrange(mod_dict[key]['tensor'], "b (nh nw) -> b nh nw", nh=image_size//patch_size, nw=image_size//patch_size) rec = tokenizers[get_transform_key(key)].decode_tokens(img_tok, timesteps=t, image_size=image_size, verbose=verbose) rec = (0.5*(rec+1)).clamp(0, 1) return tensor_to_images(rec) def decode_tok_depth(mod_dict, tokenizers, key='tok_depth', image_size=224, patch_size=16, t=25, verbose=False, cmap='turbo'): """ Decodes a sequence of depth tokens from a model dictionary into an RGB image. Args: mod_dict (dict): Model output dictionary. tokenizers (dict): Dictionary of tokenizers. key (str): Key of the tokenized depth modality to decode. image_size (int): Size of the image. patch_size (int): Size of the patches. t (int): Number of timesteps to decode using the tokenizer diffusion model (if applicable). verbose (bool): Whether to print the decoding progress. cmap (str): Colormap to use for the depth image. """ img_tok = rearrange(mod_dict[key]['tensor'], "b (nh nw) -> b nh nw", nh=image_size//patch_size, nw=image_size//patch_size) rec = tokenizers[get_transform_key(key)].decode_tokens(img_tok, timesteps=t, image_size=image_size, verbose=verbose) rec = rec.detach().cpu().numpy()[:,0] if cmap is None: return rec colormap = plt.get_cmap('turbo') imgs = [] for img in rec: img_norm = (img - np.min(img)) / (np.max(img) - np.min(img)) rgb_image = colormap(img_norm)[..., :3] imgs.append(rgb_image) rgb_image = np_squeeze(np.stack(imgs), axis=0) return rgb_image def decode_tok_semseg(rgb_img, mod_dict, tokenizers, key='tok_semseg', image_size=224, patch_size=16, use_detectron=True, return_logits=False): """ Decodes a sequence of semantic segmentation tokens from a model dictionary into an RGB image. Args: rgb_img (torch.Tensor): RGB image to overlay the semantic segmentation on. mod_dict (dict): Model output dictionary. tokenizers (dict): Dictionary of tokenizers. key (str): Key of the tokenized semantic segmentation modality to decode. image_size (int): Size of the image. patch_size (int): Size of the patches. use_detectron (bool): Uses detectron2's visualization for the semseg output. """ tokens = mod_dict[key]['tensor'] tokens = tokens.unsqueeze(0) if tokens.ndim == 1 else tokens img_tok = rearrange(tokens, "b (nh nw) -> b nh nw", nh=image_size//patch_size, nw=image_size//patch_size) rec = tokenizers[get_transform_key(key)].decode_tokens(img_tok).detach().cpu() if return_logits: return rec semsegs = rec.argmax(1) B, H, W = semsegs.shape if not use_detectron: return semsegs if B > 1 else semsegs[0] else: rgb_imgs = [rgb_img] * B imgs = [] for rgb, semseg in zip(rgb_imgs, semsegs): if USE_DETECTRON: v = Visualizer(255*rgb, coco_metadata, scale=1.2, instance_mode=ColorMode.IMAGE_BW) img = v.draw_sem_seg((semseg-1).cpu()).get_image() / 255.0 else: colormap = plt.get_cmap('viridis') img = colormap(semseg.cpu())[..., :3] imgs.append(img) imgs = np_squeeze(np.stack(imgs), axis=0) return imgs def decode_tok_clip(mod_dict, tokenizers, key='tok_clip', image_size=224, patch_size=16): """ Decodes a sequence of CLIP tokens from a model dictionary into an PCA representation. Args: mod_dict (dict): Model output dictionary. key (str): Key of the tokenized CLIP modality to decode. tokenizers (dict): Dictionary of tokenizers. image_size (int): Size of the image. patch_size (int): Size of the patches. """ n_patches = image_size // patch_size img_tok = rearrange(mod_dict[key]['tensor'], "b (nh nw) -> b nh nw", nh=n_patches, nw=n_patches) rec = tokenizers[get_transform_key(key)].decode_tokens(img_tok) pca_viz = [pca_visualize(feat) for feat in rec] pca_viz = np_squeeze(np.stack(pca_viz), axis=0) return pca_viz def decode_tok_dinov2(mod_dict, tokenizers, key='tok_dinov2', image_size=224, patch_size=14): """ Decodes a sequence of DINOv2 spatial tokens from a model dictionary into an PCA representation. Args: mod_dict (dict): Model output dictionary. key (str): Key of the tokenized CLIP modality to decode. tokenizers (dict): Dictionary of tokenizers. image_size (int): Size of the image. patch_size (int): Size of the patches. """ patch_size = 14 n_patches = image_size // patch_size img_tok = rearrange(mod_dict[key]['tensor'], "b (nh nw) -> b nh nw", nh=n_patches, nw=n_patches) rec = tokenizers[get_transform_key(key)].decode_tokens(img_tok) pca_viz = [pca_visualize(feat) for feat in rec] pca_viz = np_squeeze(np.stack(pca_viz), axis=0) return pca_viz def decode_tok_imagebind(mod_dict, tokenizers, key='tok_imagebind', image_size=224, patch_size=14): """ Decodes a sequence of ImageBind spatial tokens from a model dictionary into an PCA representation. Args: mod_dict (dict): Model output dictionary. key (str): Key of the tokenized CLIP modality to decode. tokenizers (dict): Dictionary of tokenizers. image_size (int): Size of the image. patch_size (int): Size of the patches. """ patch_size = 14 n_patches = image_size // patch_size img_tok = rearrange(mod_dict[key]['tensor'], "b (nh nw) -> b nh nw", nh=n_patches, nw=n_patches) rec = tokenizers[get_transform_key(key)].decode_tokens(img_tok) pca_viz = [pca_visualize(feat) for feat in rec] pca_viz = np_squeeze(np.stack(pca_viz), axis=0) return pca_viz def decode_tok_dinov2_global(mod_dict, tokenizers, key='tok_dinov2_global'): """ Decodes a sequence of DINOv2 global tokens from a model dictionary. Args: mod_dict (dict): Model output dictionary. key (str): Key of the tokenized DINOv2 global token modality to decode. tokenizers (dict): Dictionary of tokenizers. image_size (int): Size of the image. patch_size (int): Size of the patches. """ toks = rearrange(mod_dict[key]['tensor'].long(), 'b n -> b n 1 1') rec = tokenizers[get_transform_key(key)].decode_tokens(toks) return rec.squeeze() def decode_tok_imagebind_global(mod_dict, tokenizers, key='tok_imagebind_global'): """ Decodes a sequence of ImageBind global tokens from a model dictionary. Args: mod_dict (dict): Model output dictionary. key (str): Key of the tokenized ImageBind global token modality to decode. tokenizers (dict): Dictionary of tokenizers. image_size (int): Size of the image. patch_size (int): Size of the patches. """ toks = rearrange(mod_dict[key]['tensor'].long(), 'b n -> b n 1 1') rec = tokenizers[get_transform_key(key)].decode_tokens(toks) return rec.squeeze() def decode_color_palette(mod_dict, text_tokenizer, key='color_palette'): """ Decodes a sequence of color palettes from a model dictionary. Args: mod_dict (dict): Model output dictionary. key (str): Key of the tokenized ImageBind modality to decode. tokenizers (dict): Dictionary of tokenizers. image_size (int): Size of the image. patch_size (int): Size of the patches. """ decoded = decode_text(mod_dict, key, text_tokenizer)[2] all_decoded = decoded if isinstance(decoded, list) else [decoded] all_decoded = [d.replace(' [EOS]', '') for d in all_decoded] all_decoded = [visualize_palettes_multi(d) for d in all_decoded] all_decoded = all_decoded[0] if len(all_decoded) == 1 else all_decoded return all_decoded def decode_human_poses(mod_dict, tokenizers, text_tokenizer, key='human_poses'): """ Decodes human poses tokenized with text + BMLP """ decoded = decode_text(mod_dict, key, text_tokenizer)[2] all_decoded = decoded if isinstance(decoded, list) else [decoded] all_decoded = [d.replace(' [EOS]', '') for d in all_decoded] imgs = [] for decoded in all_decoded: img = np.ones((224,224,4)) if decoded != 'none': try: img = visualize_human_poses(decoded, tokenizers[key], mod_dict) except Exception as e: print('Error in decoding human poses. Packages required for plotting may not be installed. Trace:') print(e) imgs.append(img) imgs = np_squeeze(np.stack(imgs), axis=0) return imgs metadata_transform = MetadataTransform(shuffle=False, random_trunc=False, return_chunks=False) def _split_metadata_string(input_string): result = [] current_subseq = [] for part in input_string.split(): # If we encounter a "v1" and there's already a subsequence being built, # we add it to the result and start a new one if 'v1' in part and current_subseq: result.append(current_subseq) current_subseq = [] current_subseq.append(part) # Append any remaining subsequence to the result if current_subseq: result.append(current_subseq) return result def decode_metadata(mod_dict, text_tokenizer, key='metadata'): """ Decodes a sequence of metadata tokens into a dictionary of metadata. Args: mod_dict (dict): Model output dictionary. key (str): Key of the metadata modality to decode. text_tokenizer (tokenizers.Tokenizer): Text tokenizer. """ decoded = decode_text(mod_dict, key, text_tokenizer)[2] all_decoded = decoded if isinstance(decoded, list) else [decoded] all_decoded = [d.replace(' [EOS]', '').replace(' [PAD]', '') for d in all_decoded] all_metadata = [] for decoded in all_decoded: parts = _split_metadata_string(decoded) invalid_parts = [] metadata_dict = {} for part in parts: # Check if part has been parsed correctly if len(part) != 2: invalid_parts.append(str(part)) continue metadata_id, metadata_value = part if (not metadata_id.startswith('v1=') or not metadata_value.startswith('v0=') or metadata_id not in metadata_transform.id_metadata_map): invalid_parts.append(str(part)) # Parse metadata type and value metadata_type = metadata_transform.id_metadata_map[metadata_id] metadata_value = int(metadata_value.split('=')[1]) if metadata_type in metadata_transform.image_dim_modalities: metadata_value *= metadata_transform.image_dim_bin_size elif metadata_type in metadata_transform.metadata_min_max_bins: vmin, vmax, bins = metadata_transform.metadata_min_max_bins[metadata_type] metadata_value = (vmax - vmin) * (metadata_value / bins) + vmin metadata_dict[metadata_type] = metadata_value metadata_dict = {k: metadata_dict[k] for k in metadata_transform.metadata_id_map if k in metadata_dict} all_metadata.append(metadata_dict) all_metadata = all_metadata[0] if len(all_metadata) == 1 else all_metadata return all_metadata def decode_text(mod_dict, key, text_tokenizer): """ Decodes a text sequence from a model dictionary. Args: mod_dict (dict): Model output dictionary. key (str): Key of the text modality to decode. text_tokenizer (tokenizers.Tokenizer): Text tokenizer. """ input_texts, target_texts, merged_texts = [], [], [] sentinel_ids = set(get_sentinel_to_id_mapping(text_tokenizer).values()) B = mod_dict[key]['tensor'].shape[0] for i in range(B): input_seq = mod_dict[key]['tensor'][i] input_seq = input_seq[mod_dict[key]['input_mask'][i] == 0] input_seq = input_seq.tolist() target_seq = mod_dict[key]['tensor'][i] target_seq = target_seq[mod_dict[key]['target_mask'][i] == 0] target_seq = target_seq.tolist() merged_seq = merge_span_masking(input_seq, target_seq, sentinel_ids=sentinel_ids) input_text = text_tokenizer.decode(input_seq, skip_special_tokens=False) target_text = text_tokenizer.decode(target_seq, skip_special_tokens=False) merged_text = text_tokenizer.decode(merged_seq, skip_special_tokens=False) input_texts.append(input_text) target_texts.append(target_text) merged_texts.append(merged_text) if B == 1: input_texts, target_texts, merged_texts = input_texts[0], target_texts[0], merged_texts[0] return input_texts, target_texts, merged_texts def decode_sam_instances(mod_dict, tokenizers, text_tokenizer, key='sam_instance', image_size=224, token_len=16): ''' Decodes a sequence of SAM instance tokens into the instance representation. Args: mod_dict (dict): Model output dictionary. key (str): Key of the tokenized ImageBind modality to decode. tokenizers (dict): Dictionary of tokenizers. text_tokenizer (tokenizers.Tokenizer): Text tokenizer. image_size (int): Size of the image. token_len (int): Tokenized SAM instance token length. ''' assert image_size == 224, 'SAM instance decoding only supports 224x224 images' decoded = decode_text(mod_dict, key, text_tokenizer)[2] all_decoded = decoded if isinstance(decoded, list) else [decoded] all_decoded = [d.replace(' [EOS]', '') for d in all_decoded] # Generate deterministic SAM color palette rng = np.random.default_rng(seed=0) sam_palette = [rng.integers(0, 255, size=3) for i in range(1000)] def group_by_identifier(input_list, identifier): ''' Groups the input_list [a,b,c,a,d,d,c,..] using the identifier a, in the following format: [[b,c], [d,d,c], ...] ''' return [list(group) for key, group in groupby(input_list, lambda x: x == identifier) if not key] def map_locations(inp, tokens=False): ''' Converts v0, v1, v2, v3 textual representation into int. When tokens=True, inp is mapped to its corresponding token id. ''' if '=' not in inp: return None axis, position = inp.split("=") try: position = int(position) except: return None if tokens: if axis == 'v0': return position else: return position + 512 return position def iou(box1, box2): ''' Calculates iou of the input bounding boxes ''' # Calculate the coordinates of the intersection rectangle x1 = max(box1[0], box2[0]) y1 = max(box1[1], box2[1]) x2 = min(box1[2], box2[2]) y2 = min(box1[3], box2[3]) # Calculate the area of the intersection intersection_area = max(0, x2 - x1) * max(0, y2 - y1) # Calculate the areas of the individual bounding boxes area_box1 = (box1[2] - box1[0]) * (box1[3] - box1[1]) area_box2 = (box2[2] - box2[0]) * (box2[3] - box2[1]) # Calculate the union area union_area = area_box1 + area_box2 - intersection_area # Calculate and return the IoU return intersection_area / union_area all_sam_instances = [] for decoded in all_decoded: tokens_per_sample = [] bboxes_per_sample = [] areas_per_sample = [] parts = decoded.split() for part in group_by_identifier(parts, identifier='point'): instances = part[2:] # Ignore 'none' cases if len(instances) <= 1: continue for positions in group_by_identifier(part, identifier='polygon'): # Ignore incomplete polygons if len(positions) != token_len + 4: continue bbox, tokens = positions[:4], positions[4:] min_w, min_h, max_w, max_h = map(map_locations, bbox) # Ignore the cases where the bounding box is prediction is in incorrect format if None in [min_w, max_w, min_h, max_h] or (min_w >= max_w or min_h >= max_h): continue bbox = np.array([min_h, min_w, max_h, max_w]) tokens = list(map(lambda x: map_locations(x, tokens=True), tokens)) if None in tokens: continue tokens = np.array(tokens) tokens_per_sample.append(tokens) bboxes_per_sample.append(bbox) areas_per_sample.append((max_w - min_w) * (max_h - min_h)) final_instances = np.zeros((image_size, image_size, 3), dtype=np.uint8) if len(areas_per_sample) == 0: return final_instances # Sort the instance masks by area areas_per_sample = np.array(areas_per_sample) sorted_idx = np.argsort(-areas_per_sample) tokens_per_sample = np.stack(tokens_per_sample)[sorted_idx] bboxes_per_sample = np.stack(bboxes_per_sample)[sorted_idx] # Decoded tokens tokens_per_sample = torch.LongTensor(tokens_per_sample).reshape(-1, 4, 4).to(device) decoded_tokens = tokenizers[key].decode_tokens(tokens_per_sample) instances = torch.sigmoid(decoded_tokens).squeeze(1).cpu().detach().numpy() # Filter and group instances representive_masks = [] representive_bboxes = [] for (mask, bbox) in zip(instances, bboxes_per_sample): # Filter out unusual masks if (mask.max() - mask.min()) < 0.9: continue # Groups the duplicated instance masks duplicated_flag = False for rms, rbs in zip(representive_masks, representive_bboxes): rm, rb = rms[0], rbs[0] sim_score = 2 * ((rm * mask).sum() + 0.01) / (rm.sum() + mask.sum() + 0.01) box_iou = iou(rb, bbox) # If the similarity and IoU are high, consider them as the same instance and group them if sim_score > 0.8 and box_iou > 0.9: # Add the mask to its corresponding group rms.append(mask) rbs.append(bbox) duplicated_flag = True break if not duplicated_flag: representive_masks.append([mask]) representive_bboxes.append([bbox]) # Plot the instances for i, (rms, rbs) in enumerate(zip(representive_masks, representive_bboxes)): mask = np.mean(rms, axis=0) bbox = np.mean(rbs, axis=0).astype(np.int32) min_h, min_w, max_h, max_w = bbox.tolist() mask = cv2.resize(mask, (max_w - min_w, max_h - min_h), interpolation=cv2.INTER_CUBIC) max_w, max_h = min(max_w, final_instances.shape[1]), min(max_h, final_instances.shape[0]) mask = mask[:max_h - min_h,:max_w - min_w] > 0.5 final_instances[min_h:max_h, min_w:max_w, :][mask] = sam_palette[i] all_sam_instances.append(final_instances) all_sam_instances = all_sam_instances[0] if len(all_sam_instances) == 1 else np.stack(all_sam_instances) return all_sam_instances def decode_dict(mod_dict, tokenizers, text_tokenizer, image_size=224, patch_size=16, decoding_steps=25, activate_controlnet=False, controlnet_guidance_scale=2.5, controlnet_cond_scale=0.8, to_rgb=True, seed=None): """ Decodes the model output dictionary into a dictionary of images and text. Args: mod_dict (dict): Model output dictionary. tokenizers (dict): Dictionary of tokenizers. text_tokenizer (tokenizers.Tokenizer): Text tokenizer. image_size (int): Image size. patch_size (int): Patch size. decoding_steps (int): Number of diffusion decoding steps (if applicable). activate_controlnet (bool): Whether to activate the RGB ControlNet and override the RGB detokenizer. controlnet_guidance_scale (float): Classifier-free guidance scale for the ControlNet. controlnet_cond_scale (float): ControlNet conditioning scale. """ dec_dict = {} for key in mod_dict: k, res = get_transform_key(key), get_transform_resolution(key, image_size, to_tuple=False) if k == 'rgb': decoded = decode_input_rgb(mod_dict, key=key) elif k == 'tok_rgb': if not activate_controlnet or 'controlnet' not in tokenizers: decoded = decode_tok_rgb( mod_dict, tokenizers, key=key, image_size=res, patch_size=patch_size, t=decoding_steps, verbose=False ) else: decoded = decode_tok_rgb_controlnet( mod_dict, tokenizers, key=key, image_size=res, patch_size=patch_size, t=decoding_steps, guidance_scale=controlnet_guidance_scale, cond_scale=controlnet_cond_scale, verbose=False ) elif k == 'tok_canny_edge': decoded = decode_tok_canny_edge( mod_dict, tokenizers, key=key, image_size=res, patch_size=patch_size, t=decoding_steps, verbose=False ) elif k == 'tok_sam_edge': decoded = decode_tok_sam_edge( mod_dict, tokenizers, key=key, image_size=res, patch_size=patch_size, t=decoding_steps, verbose=False ) elif k == 'tok_normal': decoded = decode_tok_normal( mod_dict, tokenizers, key=key, image_size=res, patch_size=patch_size, t=decoding_steps, verbose=False ) elif k == 'tok_depth': decoded = decode_tok_depth( mod_dict, tokenizers, key=key, image_size=res, patch_size=patch_size, t=decoding_steps, verbose=False, cmap='turbo' if to_rgb else None ) elif k == 'tok_semseg': decoded = decode_tok_semseg( np.ones((res, res, 3)), mod_dict, tokenizers, key=key, image_size=res, patch_size=patch_size, return_logits=not to_rgb ) elif k == 'tok_clip': decoded = decode_tok_clip( mod_dict, tokenizers, key=key, image_size=res, patch_size=patch_size ) elif k == 'tok_dinov2': decoded = decode_tok_dinov2( mod_dict, tokenizers, key=key, image_size=res, patch_size=patch_size ) elif k == 'tok_dinov2_global': decoded = decode_tok_dinov2_global( mod_dict, tokenizers, key=key ) elif k == 'tok_imagebind': decoded = decode_tok_imagebind( mod_dict, tokenizers, key=key, image_size=res, patch_size=patch_size ) elif k == 'tok_imagebind_global': decoded = decode_tok_imagebind_global( mod_dict, tokenizers, key=key ) elif k == 'color_palette': decoded = decode_color_palette( mod_dict, text_tokenizer, key=key ) elif k == 'human_poses': decoded = decode_human_poses( mod_dict, tokenizers, text_tokenizer, key=key ) elif k in ['caption', 'det']: decoded = decode_text(mod_dict, key, text_tokenizer)[2] decoded = decoded if isinstance(decoded, list) else [decoded] decoded = [d.replace(' [EOS]', '') for d in decoded] elif k in ['metadata']: decoded = decode_metadata( mod_dict, text_tokenizer, key=key ) elif k == 'sam_instance': decoded = decode_sam_instances( mod_dict, tokenizers, text_tokenizer, key=key, image_size=224, ) elif k in ['t5_caption']: if 'ascii_tensor' in mod_dict[key]: decoded = [] for ascii_tensor in mod_dict[key]['ascii_tensor']: ascii_values = ascii_tensor.flatten().tolist() decoded_text = ''.join(chr(val) for val in ascii_values if val != 0) decoded.append(f"T5-XXL embedding of: {decoded_text}") decoded = decoded[0] if len(decoded) == 1 else decoded else: decoded = "T5-XXL embedding" dec_dict[key] = decoded return dec_dict # Plotting utils MOD_PRINT_NAMES = { 'rgb': 'RGB', 'tok_rgb': 'RGB (tok)', 'tok_normal': 'Normal (tok)', 'tok_depth': 'Depth (tok)', 'tok_semseg': 'Semseg (tok)', 'tok_clip': 'CLIP (tok)', 'tok_canny': 'Canny (tok)', 'tok_sam': 'SAM (tok)', 'sam_instance': 'SAM Instances (tok)', 'rgb@224': 'RGB@224', 'tok_rgb@224': 'RGB@224 (tok)', 'tok_normal@224': 'Normal@224 (tok)', 'tok_depth@224': 'Depth@224 (tok)', 'tok_semseg@224': 'Semseg@224 (tok)', 'tok_clip@224': 'CLIP@224 (tok)', 'rgb@448': 'RGB@448', 'tok_rgb@448': 'RGB@448 (tok)', 'tok_normal@448': 'Normal@448 (tok)', 'tok_depth@448': 'Depth@448 (tok)', 'tok_semseg@448': 'Semseg@448 (tok)', 'tok_clip@448': 'CLIP@448 (tok)', 'caption': 'Caption', 'det': 'Detection', 't5_caption': 'T5 XXL', 'metadata': 'Metadata', 'human_poses': 'Human poses', 'color_palette': 'Color palette', 'tok_dinov2': 'DINOv2 (tok)', 'tok_dinov2_global': 'DINOv2 global (tok)', 'tok_imagebind': 'ImageBind (tok)', 'tok_imagebind_global': 'ImageBind global (tok)', } def remove_ticks_and_labels(ax): """ Remove the axis ticks and labels Args: ax (matplotlib.axes.Axes): Axis to remove ticks and labels from """ ax.set_xticks([]) ax.set_yticks([]) ax.set_xticklabels([]) ax.set_yticklabels([]) def remove_spines(ax): """ Removes the spines from the given axis. Args: ax (matplotlib.axes.Axes): Axis to remove spines from """ ax.spines['top'].set_visible(False) ax.spines['right'].set_visible(False) ax.spines['bottom'].set_visible(False) ax.spines['left'].set_visible(False) def convert_string_to_bboxes(bboxes_str, bins=1000): """ Converts a string of bboxes to a list of bboxes. Args: bboxes_str (str): String of bboxes bins (int): Number of bins (default: 1000) """ bboxes_str = bboxes_str.split(" ") bboxes = [] for token in bboxes_str: if "=" in token: coord = token.split("=")[1] coord = float(coord) / (bins - 1) if token.startswith("v0="): bboxes.append([coord,]) else: bboxes[-1].append(coord) elif len(bboxes[-1]) == 4: bboxes[-1].append(token) else: bboxes[-1][4] = " ".join([bboxes[-1][4], token]) bboxes = [bbox for bbox in bboxes if len(bbox) == 5] return bboxes def visualize_palettes_multi(palettes): palettes = palettes.split() palettes = palettes[1:] all_colors = [] for ii in range(len(palettes)): all_colors.append(int(palettes[ii][3:])) w = h = 25 # construct palette image o = Image.new("RGB", size=(w * len(palettes)//3, h * len(palettes)//3)) arr = np.asarray(o).copy() for ii in range(len(palettes)//3): arr[:, ii * h : (ii + 1) * h, :] = all_colors[ii*3:(ii+1)*3] final_palette = arr / 255 return final_palette BOX_COLOR = (255, 0, 0) # Red TEXT_COLOR = (255, 255, 255) # White try: from fourm.utils.hmr2_utils.hmr2.models.smpl_wrapper import SMPL from fourm.utils.hmr2_utils.hmr2.utils.renderer import Renderer, cam_crop_to_full import pickle as pkl LIGHT_BLUE=(0.65098039, 0.74117647, 0.85882353) with open('./fourm/utils/hmr2_utils/model_cfg.pkl','rb') as f: pose_model_cfg = pkl.load(f) # Instantiate SMPL model smpl_cfg = {k.lower(): v for k,v in dict(pose_model_cfg.SMPL).items()} smpl_cfg['model_path'] = './fourm/utils/hmr2_utils/data/smpl' smpl_cfg['joint_regressor_extra'] = './fourm/utils/hmr2_utils/data/SMPL_to_J19.pkl' smpl_cfg['mean_params'] = './fourm/utils/hmr2_utils/data/smpl_mean_params.npz' smpl = SMPL(**smpl_cfg) # Setup the renderer renderer = Renderer(pose_model_cfg, faces=smpl.faces) except Exception as e: print(e) print('Human pose dependencies are not installed, hence poses will not be visualized. To visualize them (optional), you can do the following: \n' \ '1) Install via `pip install timm yacs smplx pyrender pyopengl==3.1.4` \n' \ ' You may need to follow the pyrender install instructions: https://pyrender.readthedocs.io/en/latest/install/index.html \n' \ '2) Download SMPL data from https://smpl.is.tue.mpg.de/. See https://github.com/shubham-goel/4D-Humans/ for an example. \n' \ '3) Copy the required SMPL files (smpl_mean_params.npz, SMPL_to_J19.pkl, smpl/SMPL_NEUTRAL.pkl) to fourm/utils/hmr2_utils/data .') def visualize_human_poses(pose, poses_tokenizer, mod_dict): full_gts = pose full_gts = full_gts.split() num_instances = len(full_gts) // 39 # total length of a pose instance seq is 39 all_verts = [] all_cam_t = [] for inst in range(num_instances): try: full_gt = full_gts[inst*39:(inst+1)*39] ##create the pose params dict all_params = {} all_params['bbox_xyxy'] = torch.Tensor((int(full_gt[1][3:])/999*224, int(full_gt[2][3:])/999*224, int(full_gt[3][3:])/999*224, int(full_gt[4][3:])/999*224)) all_params["box_center"] = torch.cat(( ((all_params["bbox_xyxy"][0] + all_params["bbox_xyxy"][2]) / 2.).unsqueeze(0).unsqueeze(1) , ( (all_params["bbox_xyxy"][1] + all_params["bbox_xyxy"][3]) / 2.).unsqueeze(0).unsqueeze(1) ), dim = 1) all_params["box_size"] = torch.max((all_params["box_center"][0,0] - all_params["bbox_xyxy"][0]) * 2 , (all_params["box_center"][0,1] - all_params["bbox_xyxy"][1]) * 2 ) all_params["img_size"] = torch.Tensor([224., 224.]) all_params["img_size"] = all_params["img_size"].unsqueeze(0) all_params["focal_length"] = torch.Tensor([5000., 5000.]) for ii in range(len(full_gt)): if full_gt[ii] == 'camera': all_params['pred_cam'] = torch.Tensor([ (int(full_gt[ii+1][3:])-49.95)/49.95, (int(full_gt[ii+2][3:])-49.95)/49.95, (int(full_gt[ii+3][3:])-49.95)/49.95 ]) break all_params['pred_cam'] = all_params['pred_cam'].unsqueeze(0) all_params['pred_smpl_params'] = {} for ii in range(len(full_gt)): if full_gt[ii] == 'shape': all_params['pred_smpl_params']['betas'] = torch.Tensor([ (int(full_gt[ii+1][3:])-499.5)/166.5, (int(full_gt[ii+2][3:])-499.5)/166.5, (int(full_gt[ii+3][3:])-499.5)/166.5, (int(full_gt[ii+4][3:])-499.5)/166.5, (int(full_gt[ii+5][3:])-499.5)/166.5, (int(full_gt[ii+6][3:])-499.5)/166.5, (int(full_gt[ii+7][3:])-499.5)/166.5, (int(full_gt[ii+8][3:])-499.5)/166.5, (int(full_gt[ii+9][3:])-499.5)/166.5, (int(full_gt[ii+10][3:])-499.5)/166.5 ]) break all_params['pred_smpl_params']['betas'] = all_params['pred_smpl_params']['betas'].unsqueeze(0) for ii in range(len(full_gt)): if full_gt[ii] == 'global': all_params['pred_smpl_params']['global_orient'] = torch.Tensor( [ [(int(full_gt[ii+1][3:])-499.5)/499.5, (int(full_gt[ii+2][3:])-499.5)/499.5, (int(full_gt[ii+3][3:])-499.5)/499.5 ] , [ (int(full_gt[ii+4][3:])-499.5)/499.5, (int(full_gt[ii+5][3:])-499.5)/499.5, (int(full_gt[ii+6][3:])-499.5)/499.5], [(int(full_gt[ii+7][3:])-499.5)/499.5, (int(full_gt[ii+8][3:])-499.5)/499.5, (int(full_gt[ii+9][3:])-499.5)/499.5 ] ] ) break all_params['pred_smpl_params']['global_orient'] = all_params['pred_smpl_params']['global_orient'].unsqueeze(0).unsqueeze(0) body_poses = torch.FloatTensor() for ii in range(len(full_gt)): if full_gt[ii] == 'pose': pose_start = ii break for ii in range(8): pose_curr = ii + pose_start + 1 if 'v1' in full_gt[pose_curr]: poses_curr = torch.Tensor([int(full_gt[pose_curr][3:])+512]) else: poses_curr = torch.Tensor([int(full_gt[pose_curr][3:])]) poses_curr = poses_curr body_poses = torch.cat((body_poses,poses_curr), dim=0) body_poses = body_poses.long() body_poses = body_poses.unsqueeze(0).unsqueeze(2).unsqueeze(2).to(device) body_poses = poses_tokenizer.decode_tokens(body_poses).squeeze(2).squeeze().reshape(1,23,3,3).cpu() all_params['pred_smpl_params']['body_pose'] = body_poses smpl_params = (all_params['pred_smpl_params']) smpl_output = smpl(**{k: v.float().cpu() for k,v in smpl_params.items()}, pose2rot=False) for n in range(smpl_output.vertices.size(0)): # Add all verts and cams to list verts = smpl_output.vertices[n].detach().cpu().numpy() img_size = all_params["img_size"].float() pred_cam = all_params['pred_cam'] box_center = all_params["box_center"].float() box_size = all_params["box_size"].float() scaled_focal_length = pose_model_cfg.EXTRA.FOCAL_LENGTH / pose_model_cfg.MODEL.IMAGE_SIZE * img_size.max() pred_cam_t_full = cam_crop_to_full(pred_cam, box_center, box_size, img_size, scaled_focal_length).detach().cpu().numpy() cam_t = pred_cam_t_full[n] all_verts.append(verts) all_cam_t.append(cam_t) except Exception as e: print('Error in decoding human poses: ', end='') print(e) continue try: input_img = denormalize(mod_dict['rgb@224']['tensor'].squeeze(), mean=(IMAGENET_DEFAULT_MEAN), std=IMAGENET_DEFAULT_STD).permute(1,2,0).cpu() except Exception as e: print(e) input_img = 1. if 'tok_rgb' in mod_dict: input_img = decode_tok_rgb(mod_dict, toks, key='tok_rgb') # Render front view input_img_overlay = 0.5* input_img[:,:,:3] if len(all_verts) > 0: misc_args = dict( mesh_base_color=LIGHT_BLUE, scene_bg_color=(1, 1, 1), focal_length=scaled_focal_length, ) cam_view = renderer.render_rgba_multiple(all_verts, cam_t=all_cam_t, render_res=img_size[n], **misc_args) mask = (cam_view[:,:,0]<1.).astype(int)[:,:,None] input_img_overlay = 0.5* input_img[:,:,:3] * (1-mask) + cam_view[:,:,:3] * mask return input_img_overlay def visualize_bboxes(img, bboxes_str, color=BOX_COLOR, thickness=2): """ Visualizes bounding boxes on the image. Args: img (np.array): Image to draw bounding boxes on. bboxes_str (str): String containing bounding boxes in the format: v0=1 v1=2 v2=3 v3=4 class_name ..., where v0 is xmin, v1 is ymin, v2 is xmax, v3 is ymax color (tuple): Color of the bounding box. thickness (int): Thickness of the bounding box. """ if img is None: img = 255 * np.ones((256,256,3), dtype=np.int32) img = img.copy() bboxes_str = bboxes_str.replace('[PAD]', '') if len(bboxes_str.replace('[EOS]', '')) == 0: return img try: bboxes = convert_string_to_bboxes(bboxes_str.replace(' [EOS]', '')) except: return img for bbox in bboxes: x_min, y_min, x_max, y_max, class_name = bbox img_h, img_w = img.shape[0], img.shape[1] x_min, x_max, y_min, y_max = int(x_min * img_w), int(x_max * img_w), int(y_min * img_h), int(y_max * img_h) cv2.rectangle(img, (x_min, y_min), (x_max, y_max), color=color, thickness=thickness) ((text_width, text_height), _) = cv2.getTextSize(class_name.rstrip(), cv2.FONT_HERSHEY_SIMPLEX, 0.35, 1) cv2.rectangle(img, (x_min, y_min - int(1.3 * text_height)), (x_min + text_width, y_min), BOX_COLOR, -1) cv2.putText( img, text=f"{class_name}", org=(x_min, y_min - int(0.3 * text_height)), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.35, color=TEXT_COLOR, lineType=cv2.LINE_AA, ) return img def plot_text_in_square(ax, text, padding=0.5, fontsize=14, wrap_width=50): """ Plots text in a square. Args: ax (matplotlib.axes.Axes): Matplotlib axis to plot on text (str): Text to plot padding (float): Padding around the text fontsize (int): Font size of the text wrap_width (int): Width of the text to wrap """ ax.set_xlim(0, 1) ax.set_ylim(0, 1) if isinstance(text, list): text = text[0] text = text.replace('[PAD]', '') # Wrap the text if necessary wrapped_text = textwrap.fill(text, int(wrap_width)) # Add the padding bbox_props = dict(boxstyle="square,pad=" + str(padding), facecolor="white", edgecolor="black") # Add the text to the plot ax.text(0.5, 0.5, wrapped_text, ha='center', va='center', fontsize=fontsize, bbox=bbox_props) remove_ticks_and_labels(ax) remove_spines(ax) def text_to_pil_image(text, padding=0.5, fontsize=14, wrap_width=40, image_size=(512, 512)): """ Converts text to a PIL image. Args: text (str): Text to convert to image padding (float): Padding around the text fontsize (int): Font size of the text wrap_width (int): Width of the text to wrap image_size (tuple): Size of the output image (width, height) Returns: PIL.Image.Image: Generated image with the text """ fig, ax = plt.subplots(figsize=(image_size[0] / 100, image_size[1] / 100), dpi=100) ax.set_xlim(0, 1) ax.set_ylim(0, 1) if isinstance(text, list): text = text[0] text = text.replace('[PAD]', '') # Wrap the text if necessary wrapped_text = textwrap.fill(text, wrap_width) # Add the padding bbox_props = dict(boxstyle="square,pad=" + str(padding), facecolor="white", edgecolor="black") # Add the text to the plot ax.text(0.5, 0.5, wrapped_text, ha='center', va='center', fontsize=fontsize, bbox=bbox_props) # Remove ticks, labels, and spines ax.set_xticks([]) ax.set_yticks([]) for spine in ax.spines.values(): spine.set_visible(False) # Convert the plot to a PIL image fig.canvas.draw() image = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8) image = image.reshape(fig.canvas.get_width_height()[::-1] + (3,)) plt.close(fig) return Image.fromarray(image) def plot_modality(dec_dict, key, ax, figscale=4.0): """ Plots a single modality. Function name has a typo because of legacy reasons. Args: dec_dict (dict): Dictionary of decoded modalities key (str): Key of the modality to plot ax (matplotlib.axes.Axes): Matplotlib axis to plot on figscale (float): Scaling factor for the figure (used to scale the caption box) """ modality = dec_dict[key] k = get_transform_key(key) if 'tok' in k or k == 'rgb' or k == 'human_poses' or k == 'color_palette': ax.imshow(modality.clip(0,1)) elif k == 'caption': plot_text_in_square(ax, modality, wrap_width=max(1,int(7*figscale))) # 7*figscale turns out to make caption box fit nicely elif k == 't5_caption': plot_text_in_square(ax, modality, wrap_width=max(1,int(7*figscale))) # 7*figscale turns out to make caption box fit nicely elif k == 'metadata': modality = ',\n'.join([f'{k}: {v:.2f}' if isinstance(v, float) else f'{k}: {v}' for k, v in modality.items()]) plot_text_in_square(ax, modality, wrap_width=max(1,int(7*figscale)), fontsize=11) elif k == 'det': bbox_img = visualize_bboxes(np.ones((224,224,3)), modality, thickness=2) ax.imshow(bbox_img.clip(0,1)) def plot_conds_and_targets(cond_domains, target_domains, dec_dicts, save_path=None, fs_titles=15, figscale=4.0, dpi=100): """ Plots the conditioning and target modalities for a batch of samples. Args: cond_domains (list of str): List of conditioning domains target_domains (list of str): List of target domains dec_dicts (list of dicts): List of dictionaries containing the decoded conditioning and target modalities save_path (str): Path to save the figure. If None, the figure is not saved but plotted instead. fs_titles (int): Font size of the titles figscale (float): Scaling factor for the figure size (minimum 4.0 for good results) dpi (float): Dots per inch for the saved figure """ n_cond = len(cond_domains) n_target = len(target_domains) n_samples = len(dec_dicts) ncols = n_samples + 1 if n_cond > 0 else n_samples nrows = max(n_cond, n_target) fig, ax = plt.subplots(nrows=nrows, ncols=ncols, figsize=(ncols*figscale, nrows*figscale), facecolor='white') if nrows == 1 and ncols == 1: ax = np.array([[ax]]) elif nrows == 1: ax = np.expand_dims(ax, axis=0) elif ncols == 1: ax = np.expand_dims(ax, axis=1) for cond_idx, cond_domain in enumerate(cond_domains): axi = ax[cond_idx, 0] plot_modality(dec_dicts[0], key=cond_domain, ax=axi) axi.set_title(f'Conditioning: {MOD_PRINT_NAMES[cond_domain]}', fontsize=fs_titles) # Remove spines that are not needed if n_cond > 0: for i in range(n_cond, nrows, 1): remove_spines(ax[i, 0]) offset = 0 if n_cond == 0 else 1 for sample_idx, dec_dict in enumerate(dec_dicts): for target_idx, target_domain in enumerate(target_domains): axi = ax[target_idx, sample_idx+offset] plot_modality(dec_dict, key=target_domain, ax=axi) axi.set_title(f'{sample_idx+1}.{target_idx+1}: {MOD_PRINT_NAMES[target_domain]}', fontsize=fs_titles) # Remove spines that are not needed for i in range(n_target, nrows, 1): remove_spines(ax[i, sample_idx+offset]) for ax in fig.axes: remove_ticks_and_labels(ax) plt.tight_layout() if save_path is not None: os.makedirs(os.path.dirname(save_path), exist_ok=True) plt.savefig(save_path, bbox_inches='tight', dpi=dpi) #, pil_kwargs={'quality': 30}) plt.close() else: plt.show() def save_conds_and_targets(cond_domains, target_domains, dec_dicts, save_dir, sample_idx, suffix=None, vis_det=False): """ Saves the conditioning and target modalities for a batch of samples. Args: cond_domains (list of str): List of conditioning domains target_domains (list of str): List of target domains dec_dicts (list of dicts): List of dictionaries containing the decoded conditioning and target modalities save_dir (str): Path to save the modalities sample_idx (int): Unique index of the dataset sample suffix (str): Suffix to append to the saved file names vis_det (bool): Whether to visualize detection """ for variant_idx, dec_dict in enumerate(dec_dicts): for domain in cond_domains + target_domains: if variant_idx != 0 and domain in cond_domains: continue variant_suffix = f'_{variant_idx}' if domain in target_domains else '' if suffix is not None: variant_suffix += f'_{suffix}' domain_save_dir = os.path.join(save_dir, 'conds' if domain in cond_domains else 'targets', domain) os.makedirs(domain_save_dir, exist_ok=True) if 'tok' in domain or domain in ['rgb', 'human_poses', 'color_palette']: img = Image.fromarray((255 * dec_dict[domain]).astype(np.uint8)) if domain in ['tok_clip', 'tok_dinov2', 'tok_imagebind']: img = img.resize((224,224), resample=Image.NEAREST) save_path = os.path.join(domain_save_dir, f'{sample_idx:06d}{variant_suffix}.png') img.save(save_path) elif domain in ['caption', 'det', 'metadata']: if vis_det: save_path = os.path.join(domain_save_dir, f'{sample_idx:06d}{variant_suffix}.png') bbox_img = visualize_bboxes(np.ones((512,512,3)), dec_dict[domain], thickness=2) bbox_img = Image.fromarray((255 * bbox_img.clip(0,1)).astype(np.uint8)) bbox_img.save(save_path) else: # Save caption as text file save_path = os.path.join(domain_save_dir, f'{sample_idx:06d}{variant_suffix}.txt') with open(save_path, 'w') as f: f.write(dec_dict[domain]) def plot_images_with_captions(images, captions, save_path=None, dpi=100, wrap_length=40, figscale=4.0): """ Plots images with their corresponding captions. Parameters: - images (torch.Tensor): A tensor of shape Bx3xHxW with images. - captions (list): A list of B captions. """ assert len(images) == len(captions), "Number of images must match number of captions!" B = len(images) sqrt_B = int(B**0.5) # Determine the number of rows and columns for subplots nrows = sqrt_B ncols = (B + nrows - 1) // nrows fig, axarr = plt.subplots(nrows=nrows, ncols=ncols, figsize=(figscale*ncols, figscale*nrows)) axarr = np.array([axarr]) if nrows == 1 and ncols == 1 else axarr.ravel() for i, ax in enumerate(axarr): if i < B: # Convert tensor image to numpy image_np = images[i].permute(1, 2, 0).cpu().float().numpy() ax.imshow(image_np) # Place caption below the image caption_wrapped = textwrap.fill(captions[i], width=wrap_length) ax.text(0.5, -0.1, caption_wrapped, ha='center', va='top', transform=ax.transAxes, wrap=True) ax.axis("off") else: ax.axis("off") # Hide any additional subplots plt.subplots_adjust(hspace=0.6) plt.tight_layout() if save_path is not None: os.makedirs(os.path.dirname(save_path), exist_ok=True) plt.savefig(save_path, bbox_inches='tight', dpi=dpi) plt.close() else: plt.show()