util/dataloader.py

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