import os import gc import cv2 import json import math import decord import random import numpy as np from PIL import Image from tqdm import tqdm from decord import VideoReader from contextlib import contextmanager from func_timeout import FunctionTimedOut from typing import Optional, Sized, Iterator import torch from torch.utils.data import Dataset, Sampler import torch.nn.functional as F from torchvision.transforms import ToPILImage from torchvision import transforms from accelerate.logging import get_logger logger = get_logger(__name__) import threading log_lock = threading.Lock() def log_error_to_file(error_message, video_path): with log_lock: with open("error_log.txt", "a") as f: f.write(f"Error: {error_message}\n") f.write(f"Video Path: {video_path}\n") f.write("-" * 50 + "\n") def draw_kps(image_pil, kps, color_list=[(255,0,0), (0,255,0), (0,0,255), (255,255,0), (255,0,255)]): stickwidth = 4 limbSeq = np.array([[0, 2], [1, 2], [3, 2], [4, 2]]) kps = np.array(kps) w, h = image_pil.size out_img = np.zeros([h, w, 3]) for i in range(len(limbSeq)): index = limbSeq[i] color = color_list[index[0]] x = kps[index][:, 0] y = kps[index][:, 1] length = ((x[0] - x[1]) ** 2 + (y[0] - y[1]) ** 2) ** 0.5 angle = math.degrees(math.atan2(y[0] - y[1], x[0] - x[1])) polygon = cv2.ellipse2Poly((int(np.mean(x)), int(np.mean(y))), (int(length / 2), stickwidth), int(angle), 0, 360, 1) out_img = cv2.fillConvexPoly(out_img.copy(), polygon, color) out_img = (out_img * 0.6).astype(np.uint8) for idx_kp, kp in enumerate(kps): color = color_list[idx_kp] x, y = kp out_img = cv2.circle(out_img.copy(), (int(x), int(y)), 10, color, -1) out_img_pil = Image.fromarray(out_img.astype(np.uint8)) return out_img_pil @contextmanager def VideoReader_contextmanager(*args, **kwargs): vr = VideoReader(*args, **kwargs) try: yield vr finally: del vr gc.collect() def get_valid_segments(valid_frame, tolerance=5): valid_positions = sorted(set(valid_frame['face']).union(set(valid_frame['head']))) valid_segments = [] current_segment = [valid_positions[0]] for i in range(1, len(valid_positions)): if valid_positions[i] - valid_positions[i - 1] <= tolerance: current_segment.append(valid_positions[i]) else: valid_segments.append(current_segment) current_segment = [valid_positions[i]] if current_segment: valid_segments.append(current_segment) return valid_segments def get_frame_indices_adjusted_for_face(valid_frames, n_frames): valid_length = len(valid_frames) if valid_length >= n_frames: return valid_frames[:n_frames] additional_frames_needed = n_frames - valid_length repeat_indices = [] for i in range(additional_frames_needed): index_to_repeat = i % valid_length repeat_indices.append(valid_frames[index_to_repeat]) all_indices = valid_frames + repeat_indices all_indices.sort() return all_indices def generate_frame_indices_for_face(n_frames, sample_stride, valid_frame, tolerance=7, skip_frames_start_percent=0.0, skip_frames_end_percent=1.0, skip_frames_start=0, skip_frames_end=0): valid_segments = get_valid_segments(valid_frame, tolerance) selected_segment = max(valid_segments, key=len) valid_length = len(selected_segment) if skip_frames_start_percent != 0.0 or skip_frames_end_percent != 1.0: # print("use skip frame percent") valid_start = int(valid_length * skip_frames_start_percent) valid_end = int(valid_length * skip_frames_end_percent) elif skip_frames_start != 0 or skip_frames_end != 0: # print("use skip frame") valid_start = skip_frames_start valid_end = valid_length - skip_frames_end else: # print("no use skip frame") valid_start = 0 valid_end = valid_length if valid_length <= n_frames: return get_frame_indices_adjusted_for_face(selected_segment, n_frames), valid_length else: adjusted_length = valid_end - valid_start if adjusted_length <= 0: print(f"video_length: {valid_length}, adjusted_length: {adjusted_length}, valid_start:{valid_start}, skip_frames_end: {valid_end}") raise ValueError("Skipping too many frames results in no frames left to sample.") clip_length = min(adjusted_length, (n_frames - 1) * sample_stride + 1) start_idx_position = random.randint(valid_start, valid_end - clip_length) start_frame = selected_segment[start_idx_position] selected_frames = [] for i in range(n_frames): next_frame = start_frame + i * sample_stride if next_frame in selected_segment: selected_frames.append(next_frame) else: break if len(selected_frames) < n_frames: return get_frame_indices_adjusted_for_face(selected_frames, n_frames), len(selected_frames) return selected_frames, len(selected_frames) def frame_has_required_confidence(bbox_data, frame, ID, conf_threshold=0.88): frame_str = str(frame) if frame_str not in bbox_data: return False frame_data = bbox_data[frame_str] face_conf = any( item['confidence'] > conf_threshold and item['new_track_id'] == ID for item in frame_data.get('face', []) ) head_conf = any( item['confidence'] > conf_threshold and item['new_track_id'] == ID for item in frame_data.get('head', []) ) return face_conf and head_conf def select_mask_frames_from_index(batch_frame, original_batch_frame, valid_id, corresponding_data, control_sam2_frame, valid_frame, bbox_data, base_dir, min_distance=3, min_frames=1, max_frames=5, mask_type='face', control_mask_type='head', dense_masks=False, ensure_control_frame=True): """ Selects frames with corresponding mask images while ensuring a minimum distance constraint between frames, and that the frames exist in both batch_frame and valid_frame. Parameters: base_path (str): Base directory where the JSON files and mask results are located. min_distance (int): Minimum distance between selected frames. min_frames (int): Minimum number of frames to select. max_frames (int): Maximum number of frames to select. mask_type (str): Type of mask to select frames for ('face' or 'head'). control_mask_type (str): Type of mask used for control frame selection ('face' or 'head'). Returns: dict: A dictionary where keys are IDs and values are lists of selected mask PNG paths. """ # Helper function to randomly select frames with at least X frames apart def select_frames_with_distance_constraint(frames, num_frames, min_distance, control_frame, bbox_data, ID, ensure_control_frame=True, fallback=True): """ Selects frames with a minimum distance constraint. If not enough frames can be selected, a fallback plan is applied. Parameters: frames (list): List of frame indices to select from. num_frames (int): Number of frames to select. min_distance (int): Minimum distance between selected frames. control_frame (int): The control frame that must always be included. fallback (bool): Whether to apply a fallback strategy if not enough frames meet the distance constraint. Returns: list: List of selected frames. """ conf_thresholds = [0.95, 0.94, 0.93, 0.92, 0.91, 0.90] if ensure_control_frame: selected_frames = [control_frame] # Ensure control frame is always included else: valid_initial_frames = [] for conf_threshold in conf_thresholds: valid_initial_frames = [ f for f in frames if frame_has_required_confidence(bbox_data, f, ID, conf_threshold=conf_threshold) ] if valid_initial_frames: break if valid_initial_frames: selected_frames = [random.choice(valid_initial_frames)] else: # If no frame meets the initial confidence, fall back to a random frame (or handle as per your preference) selected_frames = [random.choice(frames)] available_frames = [f for f in frames if f != selected_frames[0]] # Exclude control frame for random selection random.shuffle(available_frames) # Shuffle to introduce randomness while available_frames and len(selected_frames) < num_frames: last_selected_frame = selected_frames[-1] valid_choices = [] for conf_threshold in conf_thresholds: valid_choices = [ f for f in available_frames if abs(f - last_selected_frame) >= min_distance and frame_has_required_confidence(bbox_data, f, ID, conf_threshold=conf_threshold) ] if valid_choices: break if valid_choices: frame = random.choice(valid_choices) available_frames.remove(frame) selected_frames.append(frame) else: if fallback: # Fallback strategy: uniformly distribute remaining frames if distance constraint cannot be met remaining_needed = num_frames - len(selected_frames) remaining_frames = available_frames[:remaining_needed] # Distribute the remaining frames evenly if possible if remaining_frames: step = max(1, len(remaining_frames) // remaining_needed) evenly_selected = remaining_frames[::step][:remaining_needed] selected_frames.extend(evenly_selected) break else: break # No valid choices remain and no fallback strategy is allowed if len(selected_frames) < num_frames: return None return selected_frames # Convert batch_frame list to a set to remove duplicates batch_frame_set = set(batch_frame) # Dictionary to store selected mask PNGs selected_masks_dict = {} selected_bboxs_dict = {} dense_masks_dict = {} selected_frames_dict = {} # ID try: mask_valid_frames = valid_frame[mask_type] # Select frames based on the specified mask type control_valid_frames = valid_frame[control_mask_type] # Control frames for control_mask_type except KeyError: if mask_type not in valid_frame.keys(): print(f"no valid {mask_type}") if control_mask_type not in valid_frame.keys(): print(f"no valid {control_mask_type}") # Get the control frame for the control mask type control_frame = control_sam2_frame[valid_id][control_mask_type] # Filter frames to only those which are in both valid_frame and batch_frame_set valid_frames = [] # valid_frames = [frame for frame in mask_valid_frames if frame in control_valid_frames and frame in batch_frame_set] for frame in mask_valid_frames: if frame in control_valid_frames and frame in batch_frame_set: # Check if bbox_data has 'head' or 'face' for the frame if str(frame) in bbox_data: frame_data = bbox_data[str(frame)] if 'head' in frame_data or 'face' in frame_data: valid_frames.append(frame) # Ensure the control frame is included in the valid frames if ensure_control_frame and (control_frame not in valid_frames): valid_frames.append(control_frame) # Select a random number of frames between min_frames and max_frames num_frames_to_select = random.randint(min_frames, max_frames) selected_frames = select_frames_with_distance_constraint(valid_frames, num_frames_to_select, min_distance, control_frame, bbox_data, valid_id, ensure_control_frame) # Store the selected frames as mask PNGs and bbox selected_masks_dict[valid_id] = [] selected_bboxs_dict[valid_id] = [] # Initialize the dense_masks_dict entry for the current ID dense_masks_dict[valid_id] = [] # Store the selected frames in the dictionary selected_frames_dict[valid_id] = selected_frames if dense_masks: for frame in original_batch_frame: mask_data_path = f"{base_dir}/{valid_id}/annotated_frame_{int(frame):05d}.png" mask_array = np.array(Image.open(mask_data_path)) binary_mask = np.where(mask_array > 0, 1, 0).astype(np.uint8) dense_masks_dict[valid_id].append(binary_mask) for frame in selected_frames: mask_data_path = f"{base_dir}/{valid_id}/annotated_frame_{frame:05d}.png" mask_array = np.array(Image.open(mask_data_path)) binary_mask = np.where(mask_array > 0, 1, 0).astype(np.uint8) selected_masks_dict[valid_id].append(binary_mask) try: for item in bbox_data[f"{frame}"]["head"]: if item['new_track_id'] == int(valid_id): temp_bbox = item['box'] break except (KeyError, StopIteration): try: for item in bbox_data[f"{frame}"]["face"]: if item['new_track_id'] == int(valid_id): temp_bbox = item['box'] break except (KeyError, StopIteration): temp_bbox = None selected_bboxs_dict[valid_id].append(temp_bbox) return selected_frames_dict, selected_masks_dict, selected_bboxs_dict, dense_masks_dict def pad_tensor(tensor, target_size, dim=0): padding_size = target_size - tensor.size(dim) if padding_size > 0: pad_shape = list(tensor.shape) pad_shape[dim] = padding_size padding_tensor = torch.zeros(pad_shape, dtype=tensor.dtype, device=tensor.device) return torch.cat([tensor, padding_tensor], dim=dim) else: return tensor[:target_size] def crop_images(selected_frame_index, selected_bboxs_dict, video_reader, return_ori=False): """ Crop images based on given bounding boxes and frame indices from a video. Args: selected_frame_index (list): List of frame indices to be cropped. selected_bboxs_dict (list of dict): List of dictionaries, each containing 'x1', 'y1', 'x2', 'y2' bounding box coordinates. video_reader (VideoReader or list of numpy arrays): Video frames accessible by index, where each frame is a numpy array (H, W, C). Returns: list: A list of cropped images in PIL Image format. """ expanded_cropped_images = [] original_cropped_images = [] for frame_idx, bbox in zip(selected_frame_index, selected_bboxs_dict): # Get the specific frame from the video reader using the frame index frame = video_reader[frame_idx] # torch.tensor # (H, W, C) # Extract bounding box coordinates and convert them to integers x1, y1, x2, y2 = int(bbox['x1']), int(bbox['y1']), int(bbox['x2']), int(bbox['y2']) # Crop to minimize the bounding box to a square width = x2 - x1 # Calculate the width of the bounding box height = y2 - y1 # Calculate the height of the bounding box side_length = max(width, height) # Determine the side length of the square (max of width or height) # Calculate the center of the bounding box center_x = (x1 + x2) // 2 center_y = (y1 + y2) // 2 # Calculate new coordinates for the square region centered around the original bounding box new_x1 = max(0, center_x - side_length // 2) # Ensure x1 is within image bounds new_y1 = max(0, center_y - side_length // 2) # Ensure y1 is within image bounds new_x2 = min(frame.shape[1], new_x1 + side_length) # Ensure x2 does not exceed image width new_y2 = min(frame.shape[0], new_y1 + side_length) # Ensure y2 does not exceed image height # Adjust coordinates if the cropped area is smaller than the desired side length # Ensure final width and height are equal, keeping it a square actual_width = new_x2 - new_x1 actual_height = new_y2 - new_y1 if actual_width < side_length: # Adjust x1 or x2 to ensure the correct side length, while staying in bounds if new_x1 == 0: new_x2 = min(frame.shape[1], new_x1 + side_length) else: new_x1 = max(0, new_x2 - side_length) if actual_height < side_length: # Adjust y1 or y2 to ensure the correct side length, while staying in bounds if new_y1 == 0: new_y2 = min(frame.shape[0], new_y1 + side_length) else: new_y1 = max(0, new_y2 - side_length) # Expand the square by 20% expansion_ratio = 0.2 # Define the expansion ratio expansion_amount = int(side_length * expansion_ratio) # Calculate the number of pixels to expand by # Calculate expanded coordinates, ensuring they stay within image bounds expanded_x1 = max(0, new_x1 - expansion_amount) # Expand left, ensuring x1 is within bounds expanded_y1 = max(0, new_y1 - expansion_amount) # Expand up, ensuring y1 is within bounds expanded_x2 = min(frame.shape[1], new_x2 + expansion_amount) # Expand right, ensuring x2 does not exceed bounds expanded_y2 = min(frame.shape[0], new_y2 + expansion_amount) # Expand down, ensuring y2 does not exceed bounds # Ensure the expanded area is still a square expanded_width = expanded_x2 - expanded_x1 expanded_height = expanded_y2 - expanded_y1 final_side_length = min(expanded_width, expanded_height) # Adjust to ensure square shape if necessary if expanded_width != expanded_height: if expanded_width > expanded_height: expanded_x2 = expanded_x1 + final_side_length else: expanded_y2 = expanded_y1 + final_side_length expanded_cropped_rgb_tensor = frame[expanded_y1:expanded_y2, expanded_x1:expanded_x2, :] expanded_cropped_rgb = Image.fromarray(np.array(expanded_cropped_rgb_tensor)).convert('RGB') expanded_cropped_images.append(expanded_cropped_rgb) if return_ori: original_cropped_rgb_tensor = frame[new_y1:new_y2, new_x1:new_x2, :] original_cropped_rgb = Image.fromarray(np.array(original_cropped_rgb_tensor)).convert('RGB') original_cropped_images.append(original_cropped_rgb) return expanded_cropped_images, original_cropped_images return expanded_cropped_images, None def process_cropped_images(expand_images_pil, original_images_pil, target_size=(480, 480)): """ Process a list of cropped images in PIL format. Parameters: expand_images_pil (list of PIL.Image): List of cropped images in PIL format. target_size (tuple of int): The target size for resizing images, default is (480, 480). Returns: torch.Tensor: A tensor containing the processed images. """ expand_face_imgs = [] original_face_imgs = [] if len(original_images_pil) != 0: for expand_img, original_img in zip(expand_images_pil, original_images_pil): expand_resized_img = expand_img.resize(target_size, Image.LANCZOS) expand_src_img = np.array(expand_resized_img) expand_src_img = np.transpose(expand_src_img, (2, 0, 1)) expand_src_img = torch.from_numpy(expand_src_img).unsqueeze(0).float() expand_face_imgs.append(expand_src_img) original_resized_img = original_img.resize(target_size, Image.LANCZOS) original_src_img = np.array(original_resized_img) original_src_img = np.transpose(original_src_img, (2, 0, 1)) original_src_img = torch.from_numpy(original_src_img).unsqueeze(0).float() original_face_imgs.append(original_src_img) expand_face_imgs = torch.cat(expand_face_imgs, dim=0) original_face_imgs = torch.cat(original_face_imgs, dim=0) else: for expand_img in expand_images_pil: expand_resized_img = expand_img.resize(target_size, Image.LANCZOS) expand_src_img = np.array(expand_resized_img) expand_src_img = np.transpose(expand_src_img, (2, 0, 1)) expand_src_img = torch.from_numpy(expand_src_img).unsqueeze(0).float() expand_face_imgs.append(expand_src_img) expand_face_imgs = torch.cat(expand_face_imgs, dim=0) original_face_imgs = None return expand_face_imgs, original_face_imgs class RandomSampler(Sampler[int]): r"""Samples elements randomly. If without replacement, then sample from a shuffled dataset. If with replacement, then user can specify :attr:`num_samples` to draw. Args: data_source (Dataset): dataset to sample from replacement (bool): samples are drawn on-demand with replacement if ``True``, default=``False`` num_samples (int): number of samples to draw, default=`len(dataset)`. generator (Generator): Generator used in sampling. """ data_source: Sized replacement: bool def __init__(self, data_source: Sized, replacement: bool = False, num_samples: Optional[int] = None, generator=None) -> None: self.data_source = data_source self.replacement = replacement self._num_samples = num_samples self.generator = generator self._pos_start = 0 if not isinstance(self.replacement, bool): raise TypeError(f"replacement should be a boolean value, but got replacement={self.replacement}") if not isinstance(self.num_samples, int) or self.num_samples <= 0: raise ValueError(f"num_samples should be a positive integer value, but got num_samples={self.num_samples}") @property def num_samples(self) -> int: # dataset size might change at runtime if self._num_samples is None: return len(self.data_source) return self._num_samples def __iter__(self) -> Iterator[int]: n = len(self.data_source) if self.generator is None: seed = int(torch.empty((), dtype=torch.int64).random_().item()) generator = torch.Generator() generator.manual_seed(seed) else: generator = self.generator if self.replacement: for _ in range(self.num_samples // 32): yield from torch.randint(high=n, size=(32,), dtype=torch.int64, generator=generator).tolist() yield from torch.randint(high=n, size=(self.num_samples % 32,), dtype=torch.int64, generator=generator).tolist() else: for _ in range(self.num_samples // n): xx = torch.randperm(n, generator=generator).tolist() if self._pos_start >= n: self._pos_start = 0 print("xx top 10", xx[:10], self._pos_start) for idx in range(self._pos_start, n): yield xx[idx] self._pos_start = (self._pos_start + 1) % n self._pos_start = 0 yield from torch.randperm(n, generator=generator).tolist()[:self.num_samples % n] def __len__(self) -> int: return self.num_samples class SequentialSampler(Sampler[int]): r"""Samples elements sequentially, always in the same order. Args: data_source (Dataset): dataset to sample from """ data_source: Sized def __init__(self, data_source: Sized) -> None: self.data_source = data_source self._pos_start = 0 def __iter__(self) -> Iterator[int]: n = len(self.data_source) for idx in range(self._pos_start, n): yield idx self._pos_start = (self._pos_start + 1) % n self._pos_start = 0 def __len__(self) -> int: return len(self.data_source) class ConsisID_Dataset(Dataset): def __init__( self, instance_data_root: Optional[str] = None, id_token: Optional[str] = None, height=480, width=640, max_num_frames=49, sample_stride=3, skip_frames_start_percent=0.0, skip_frames_end_percent=1.0, skip_frames_start=0, skip_frames_end=0, text_drop_ratio=-1, is_train_face=False, is_single_face=False, miss_tolerance=6, min_distance=3, min_frames=1, max_frames=5, is_cross_face=False, is_reserve_face=False, ): self.id_token = id_token or "" # ConsisID self.skip_frames_start_percent = skip_frames_start_percent self.skip_frames_end_percent = skip_frames_end_percent self.skip_frames_start = skip_frames_start self.skip_frames_end = skip_frames_end self.is_train_face = is_train_face self.is_single_face = is_single_face if is_train_face: self.miss_tolerance = miss_tolerance self.min_distance = min_distance self.min_frames = min_frames self.max_frames = max_frames self.is_cross_face = is_cross_face self.is_reserve_face = is_reserve_face # Loading annotations from files print(f"loading annotations from {instance_data_root} ...") with open(instance_data_root, 'r') as f: folder_anno = [i.strip().split(',') for i in f.readlines() if len(i.strip()) > 0] self.instance_prompts = [] self.instance_video_paths = [] self.instance_annotation_base_paths = [] for sub_root, anno, anno_base in tqdm(folder_anno): print(anno) self.instance_annotation_base_paths.append(anno_base) with open(anno, 'r') as f: sub_list = json.load(f) for i in tqdm(sub_list): path = os.path.join(sub_root, os.path.basename(i['path'])) cap = i.get('cap', None) fps = i.get('fps', 0) duration = i.get('duration', 0) if fps * duration < 49.0: continue self.instance_prompts.append(cap) self.instance_video_paths.append(path) self.num_instance_videos = len(self.instance_video_paths) self.text_drop_ratio = text_drop_ratio # Video params self.sample_stride = sample_stride self.max_num_frames = max_num_frames self.height = height self.width = width def _get_frame_indices_adjusted(self, video_length, n_frames): indices = list(range(video_length)) additional_frames_needed = n_frames - video_length repeat_indices = [] for i in range(additional_frames_needed): index_to_repeat = i % video_length repeat_indices.append(indices[index_to_repeat]) all_indices = indices + repeat_indices all_indices.sort() return all_indices def _generate_frame_indices(self, video_length, n_frames, sample_stride, skip_frames_start_percent=0.0, skip_frames_end_percent=1.0, skip_frames_start=0, skip_frames_end=0): if skip_frames_start_percent != 0.0 or skip_frames_end_percent != 1.0: print("use skip frame percent") valid_start = int(video_length * skip_frames_start_percent) valid_end = int(video_length * skip_frames_end_percent) elif skip_frames_start != 0 or skip_frames_end != 0: print("use skip frame") valid_start = skip_frames_start valid_end = video_length - skip_frames_end else: print("no use skip frame") valid_start = 0 valid_end = video_length adjusted_length = valid_end - valid_start if adjusted_length <= 0: print(f"video_length: {video_length}, adjusted_length: {adjusted_length}, valid_start:{valid_start}, skip_frames_end: {valid_end}") raise ValueError("Skipping too many frames results in no frames left to sample.") if video_length <= n_frames: return self._get_frame_indices_adjusted(video_length, n_frames) else: # clip_length = min(video_length, (n_frames - 1) * sample_stride + 1) # start_idx = random.randint(0, video_length - clip_length) # frame_indices = np.linspace(start_idx, start_idx + clip_length - 1, n_frames, dtype=int).tolist() clip_length = min(adjusted_length, (n_frames - 1) * sample_stride + 1) start_idx = random.randint(valid_start, valid_end - clip_length) frame_indices = np.linspace(start_idx, start_idx + clip_length - 1, n_frames, dtype=int).tolist() return frame_indices def _short_resize_and_crop(self, frames, target_width, target_height): """ Resize frames and crop to the specified size. Args: frames (torch.Tensor): Input frames of shape [T, H, W, C]. target_width (int): Desired width. target_height (int): Desired height. Returns: torch.Tensor: Cropped frames of shape [T, target_height, target_width, C]. """ T, H, W, C = frames.shape aspect_ratio = W / H # Determine new dimensions ensuring they are at least target size if aspect_ratio > target_width / target_height: new_width = target_width new_height = int(target_width / aspect_ratio) if new_height < target_height: new_height = target_height new_width = int(target_height * aspect_ratio) else: new_height = target_height new_width = int(target_height * aspect_ratio) if new_width < target_width: new_width = target_width new_height = int(target_width / aspect_ratio) resize_transform = transforms.Resize((new_height, new_width)) crop_transform = transforms.CenterCrop((target_height, target_width)) frames_tensor = frames.permute(0, 3, 1, 2) # (T, H, W, C) -> (T, C, H, W) resized_frames = resize_transform(frames_tensor) cropped_frames = crop_transform(resized_frames) sample = cropped_frames.permute(0, 2, 3, 1) return sample def _resize_with_aspect_ratio(self, frames, target_width, target_height): """ Resize frames while maintaining the aspect ratio by padding or cropping. Args: frames (torch.Tensor): Input frames of shape [T, H, W, C]. target_width (int): Desired width. target_height (int): Desired height. Returns: torch.Tensor: Resized and padded frames of shape [T, target_height, target_width, C]. """ T, frame_height, frame_width, C = frames.shape aspect_ratio = frame_width / frame_height # 1.77, 1280 720 -> 720 406 target_aspect_ratio = target_width / target_height # 1.50, 720 480 -> # If the frame is wider than the target, resize based on width if aspect_ratio > target_aspect_ratio: new_width = target_width new_height = int(target_width / aspect_ratio) else: new_height = target_height new_width = int(target_height * aspect_ratio) # Resize using batch processing frames = frames.permute(0, 3, 1, 2) # [T, C, H, W] frames = F.interpolate(frames, size=(new_height, new_width), mode='bilinear', align_corners=False) # Calculate padding pad_top = (target_height - new_height) // 2 pad_bottom = target_height - new_height - pad_top pad_left = (target_width - new_width) // 2 pad_right = target_width - new_width - pad_left # Apply padding frames = F.pad(frames, (pad_left, pad_right, pad_top, pad_bottom), mode='constant', value=0) frames = frames.permute(0, 2, 3, 1) # [T, H, W, C] return frames def _save_frame(self, frame, name="1.png"): # [H, W, C] -> [C, H, W] img = frame img = img.permute(2, 0, 1) to_pil = ToPILImage() img = to_pil(img) img.save(name) def _save_video(self, torch_frames, name="output.mp4"): from moviepy.editor import ImageSequenceClip frames_np = torch_frames.cpu().numpy() if frames_np.dtype != 'uint8': frames_np = frames_np.astype('uint8') frames_list = [frame for frame in frames_np] desired_fps = 24 clip = ImageSequenceClip(frames_list, fps=desired_fps) clip.write_videofile(name, codec="libx264") def get_batch(self, idx): decord.bridge.set_bridge("torch") video_dir = self.instance_video_paths[idx] text = self.instance_prompts[idx] train_transforms = transforms.Compose( [ transforms.Lambda(lambda x: x / 255.0 * 2.0 - 1.0), ] ) with VideoReader_contextmanager(video_dir, num_threads=2) as video_reader: video_num_frames = len(video_reader) if self.is_train_face: reserve_face_imgs = None file_base_name = os.path.basename(video_dir.replace(".mp4", "")) anno_base_path = self.instance_annotation_base_paths[idx] valid_frame_path = os.path.join(anno_base_path, "track_masks_data", file_base_name, "valid_frame.json") control_sam2_frame_path = os.path.join(anno_base_path, "track_masks_data", file_base_name, "control_sam2_frame.json") corresponding_data_path = os.path.join(anno_base_path, "track_masks_data", file_base_name, "corresponding_data.json") masks_data_path = os.path.join(anno_base_path, "track_masks_data", file_base_name, "tracking_mask_results") bboxs_data_path = os.path.join(anno_base_path, "refine_bbox_jsons", f"{file_base_name}.json") with open(corresponding_data_path, 'r') as f: corresponding_data = json.load(f) with open(control_sam2_frame_path, 'r') as f: control_sam2_frame = json.load(f) with open(valid_frame_path, 'r') as f: valid_frame = json.load(f) with open(bboxs_data_path, 'r') as f: bbox_data = json.load(f) if self.is_single_face: if len(corresponding_data) != 1: raise ValueError(f"Using single face, but {idx} is multi person.") # get random valid id valid_ids = [] backup_ids = [] for id_key, data in corresponding_data.items(): if 'face' in data and 'head' in data: valid_ids.append(id_key) valid_id = random.choice(valid_ids) if valid_ids else (random.choice(backup_ids) if backup_ids else None) if valid_id is None: raise ValueError("No valid ID found: both valid_ids and backup_ids are empty.") # get video total_index = list(range(video_num_frames)) batch_index, _ = generate_frame_indices_for_face(self.max_num_frames, self.sample_stride, valid_frame[valid_id], self.miss_tolerance, self.skip_frames_start_percent, self.skip_frames_end_percent, self.skip_frames_start, self.skip_frames_end) if self.is_cross_face: remaining_batch_index_index = [i for i in total_index if i not in batch_index] try: selected_frame_index, selected_masks_dict, selected_bboxs_dict, dense_masks_dict = select_mask_frames_from_index( remaining_batch_index_index, batch_index, valid_id, corresponding_data, control_sam2_frame, valid_frame[valid_id], bbox_data, masks_data_path, min_distance=self.min_distance, min_frames=self.min_frames, max_frames=self.max_frames, dense_masks=True, ensure_control_frame=False, ) except: selected_frame_index, selected_masks_dict, selected_bboxs_dict, dense_masks_dict = select_mask_frames_from_index( batch_index, batch_index, valid_id, corresponding_data, control_sam2_frame, valid_frame[valid_id], bbox_data, masks_data_path, min_distance=self.min_distance, min_frames=self.min_frames, max_frames=self.max_frames, dense_masks=True, ensure_control_frame=False, ) else: selected_frame_index, selected_masks_dict, selected_bboxs_dict, dense_masks_dict = select_mask_frames_from_index( batch_index, batch_index, valid_id, corresponding_data, control_sam2_frame, valid_frame[valid_id], bbox_data, masks_data_path, min_distance=self.min_distance, min_frames=self.min_frames, max_frames=self.max_frames, dense_masks=True, ensure_control_frame=True, ) if self.is_reserve_face: reserve_frame_index, _, reserve_bboxs_dict, _ = select_mask_frames_from_index( batch_index, batch_index, valid_id, corresponding_data, control_sam2_frame, valid_frame[valid_id], bbox_data, masks_data_path, min_distance=3, min_frames=4, max_frames=4, dense_masks=False, ensure_control_frame=False, ) # get mask and aligned_face_img selected_frame_index = selected_frame_index[valid_id] valid_frame = valid_frame[valid_id] selected_masks_dict = selected_masks_dict[valid_id] selected_bboxs_dict = selected_bboxs_dict[valid_id] dense_masks_dict = dense_masks_dict[valid_id] if self.is_reserve_face: reserve_frame_index = reserve_frame_index[valid_id] reserve_bboxs_dict = reserve_bboxs_dict[valid_id] selected_masks_tensor = torch.stack([torch.tensor(mask) for mask in selected_masks_dict]) temp_dense_masks_tensor = torch.stack([torch.tensor(mask) for mask in dense_masks_dict]) dense_masks_tensor = self._short_resize_and_crop(temp_dense_masks_tensor.unsqueeze(-1), self.width, self.height).squeeze(-1) # [T, H, W] -> [T, H, W, 1] -> [T, H, W] expand_images_pil, original_images_pil = crop_images(selected_frame_index, selected_bboxs_dict, video_reader, return_ori=True) expand_face_imgs, original_face_imgs = process_cropped_images(expand_images_pil, original_images_pil, target_size=(480, 480)) if self.is_reserve_face: reserve_images_pil, _ = crop_images(reserve_frame_index, reserve_bboxs_dict, video_reader, return_ori=False) reserve_face_imgs, _ = process_cropped_images(reserve_images_pil, [], target_size=(480, 480)) if len(expand_face_imgs) == 0 or len(original_face_imgs) == 0: raise ValueError(f"No face detected in input image pool") # post process id related data expand_face_imgs = pad_tensor(expand_face_imgs, self.max_frames, dim=0) original_face_imgs = pad_tensor(original_face_imgs, self.max_frames, dim=0) selected_frame_index = torch.tensor(selected_frame_index) # torch.Size(([15, 13]) [N1] selected_frame_index = pad_tensor(selected_frame_index, self.max_frames, dim=0) else: batch_index = self._generate_frame_indices(video_num_frames, self.max_num_frames, self.sample_stride, self.skip_frames_start_percent, self.skip_frames_end_percent, self.skip_frames_start, self.skip_frames_end) try: frames = video_reader.get_batch(batch_index) # torch [T, H, W, C] frames = self._short_resize_and_crop(frames, self.width, self.height) # [T, H, W, C] except FunctionTimedOut: raise ValueError(f"Read {idx} timeout.") except Exception as e: raise ValueError(f"Failed to extract frames from video. Error is {e}.") # Apply training transforms in batch frames = frames.float() frames = train_transforms(frames) pixel_values = frames.permute(0, 3, 1, 2).contiguous() # [T, C, H, W] del video_reader # Random use no text generation if random.random() < self.text_drop_ratio: text = '' if self.is_train_face: return pixel_values, text, 'video', video_dir, expand_face_imgs, dense_masks_tensor, selected_frame_index, reserve_face_imgs, original_face_imgs else: return pixel_values, text, 'video', video_dir def __len__(self): return self.num_instance_videos def __getitem__(self, idx): sample = {} if self.is_train_face: pixel_values, cap, data_type, video_dir, expand_face_imgs, dense_masks_tensor, selected_frame_index, reserve_face_imgs, original_face_imgs = self.get_batch(idx) sample["instance_prompt"] = self.id_token + cap sample["instance_video"] = pixel_values sample["video_path"] = video_dir if self.is_train_face: sample["expand_face_imgs"] = expand_face_imgs sample["dense_masks_tensor"] = dense_masks_tensor sample["selected_frame_index"] = selected_frame_index if reserve_face_imgs is not None: sample["reserve_face_imgs"] = reserve_face_imgs if original_face_imgs is not None: sample["original_face_imgs"] = original_face_imgs else: pixel_values, cap, data_type, video_dir = self.get_batch(idx) sample["instance_prompt"] = self.id_token + cap sample["instance_video"] = pixel_values sample["video_path"] = video_dir return sample # while True: # sample = {} # try: # if self.is_train_face: # pixel_values, cap, data_type, video_dir, expand_face_imgs, dense_masks_tensor, selected_frame_index, reserve_face_imgs, original_face_imgs = self.get_batch(idx) # sample["instance_prompt"] = self.id_token + cap # sample["instance_video"] = pixel_values # sample["video_path"] = video_dir # if self.is_train_face: # sample["expand_face_imgs"] = expand_face_imgs # sample["dense_masks_tensor"] = dense_masks_tensor # sample["selected_frame_index"] = selected_frame_index # if reserve_face_imgs is not None: # sample["reserve_face_imgs"] = reserve_face_imgs # if original_face_imgs is not None: # sample["original_face_imgs"] = original_face_imgs # else: # pixel_values, cap, data_type, video_dir, = self.get_batch(idx) # sample["instance_prompt"] = self.id_token + cap # sample["instance_video"] = pixel_values # sample["video_path"] = video_dir # break # except Exception as e: # error_message = str(e) # video_path = self.instance_video_paths[idx % len(self.instance_video_paths)] # print(error_message, video_path) # log_error_to_file(error_message, video_path) # idx = random.randint(0, self.num_instance_videos - 1) # return sample