diffsynth/datasets/mvdataset.py

import imageio, os, torch, warnings, torchvision, argparse, json
from peft import LoraConfig, inject_adapter_in_model
from PIL import Image
import pandas as pd
from tqdm import tqdm
from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs
import random
from decord import VideoReader
from decord import cpu, gpu
import imageio.v3 as iio

from torchvision import transforms
import torchvision
import random
import decord
from torchvision import transforms
import re
decord.bridge.set_bridge('torch')
import random
import numpy as np
from PIL import Image, ImageOps

class MulltiShot_MultiView_Dataset(torch.utils.data.Dataset):
    def __init__(self, dataset_base_path='/root/paddlejob/workspace/qizipeng/baidu/personal-code/Multi-view/multi_view/datasets/merged_mark_paishe_ds_meg_merge_dwposefilter_paishe.json',
                        ref_image_path='/root/paddlejob/workspace/qizipeng/code/longvideogen/output.json',
                        time_division_factor=4, 
                        time_division_remainder=1,
                        max_pixels=1920*1080,
                        height_division_factor=16, width_division_factor=16,
                        transform=None,
                        length=None,
                        resolution=None,
                        prev_length=5,
                        ref_num = 3,
                        training = True):
        self.data_path = dataset_base_path
        self.data = []
        self.length = length
        self.resolution = resolution  
        self.height, self.width = resolution
        self.num_frames = length
        self.time_division_factor = time_division_factor
        self.time_division_remainder = time_division_remainder
        self.max_pixels = max_pixels
        self.height_division_factor = height_division_factor
        self.width_division_factor = width_division_factor
        self.prev_length = prev_length
        self.training = training
        self.ref_num = ref_num

        with open(self.data_path, 'r') as f:
            meta_datas = json.load(f)
    
        for video_path in tqdm(meta_datas.keys()):
            context = meta_datas[video_path]
            candidate_labels = list(context.keys())
            candidate_labels.remove('text')

            disk_path = meta_datas[video_path]["disk_path"]
            if not disk_path.lower().endswith(".mp4"):
                continue


            # reader = imageio.get_reader(meta_datas[video_path]["disk_path"])
            # total_original_frames = reader.count_frames()
            # total_frame = total_original_frames # context["end_index"] - context["start_index"] - 1
            total_frame = None
            ref_id  = self.get_ref_id(face_crop_angle = context['facedetect_v1'], facedetect_v1_frame_index = context['facedetect_v1_frame_index'], total_frame = total_frame)
            if ref_id == []:
                continue
            ref_id_all = []
            for ids in ref_id:
                ref_id_grop = []
                for id in ids:
                    coordinate = context['facedetect_v1'][id][0]['detect']
                    if context['facedetect_v1'][id][0]['detect']["prob"] < 0.99:
                        continue
                    top, height, width, left = coordinate['top'], coordinate['height'], coordinate['width'], coordinate['left']
                    if not(min(height, width) > 80 ):
                        continue
                    # enlarge bbox 1.5x
                    width = int(width * 1)
                    height = int(height * 1)
                    frame_index = context['facedetect_v1_frame_index'][id]
                    ref_id_grop.append([top, height, width, left, int(frame_index)])
                if ref_id_grop != []:
                    if len(ref_id_grop) >= 3: #self.ref_num: ### 为了和ref_num = 3 保持数据一致
                        ref_id_all.append(ref_id_grop)
            if ref_id_all == []:
                continue
            meta_prompt = {}
            meta_prompt["global_caption"] = None
            meta_prompt["per_shot_prompt"] = []
            meta_prompt["single_prompt"] = context['text']
            self.data.append({'video_path': disk_path, 'meta_prompt': meta_prompt, "ref_id_all": ref_id_all})
            # self.data.append({'video_path':video_path, 'meta_prompt': meta_prompt, "ref_id_all": ref_id_all})

        random.seed(42)  # 让每次划分一致（可选）
        total = len(self.data)
        test_count = max(1, int(total * 0.05))  # 至少一个

        # 随机选择 test 的 index
        test_indices = set(random.sample(range(total), test_count))

        self.data_test = [self.data[i] for i in range(total) if i in test_indices]
        self.data_train = [self.data[i] for i in range(total) if i not in test_indices]
        print(f"🔥 数据集划分完成：Train={len(self.data_train)}, Test={len(self.data_test)}")

        if self.height is not None and self.width is not None:
            print("Height and width are fixed. Setting `dynamic_resolution` to False.")
            self.dynamic_resolution = False
        elif self.height is None and self.width is None:
            print("Height and width are none. Setting `dynamic_resolution` to True.")
            self.dynamic_resolution = True

    def get_ref_id(self, face_crop_angle, facedetect_v1_frame_index = None, total_frame = None, angle_threshold=50):
        """
            返回满足角度差异要求的三元组 [i, j, k]
            要求：
                - face_crop_angle[i] / [j] / [k] 都必须非空
                - i,j 两者任意 yaw/pitch/roll 差值 > angle_threshold
                - k != i != j，且 k 也必须非空
            """
        ref_id = []
        max_try = 5
        need_max = 3
        try_num = 0
            
        # 过滤空元素，保留有效索引
        valid_indices = [idx for idx, item in enumerate(face_crop_angle) if item]
        N = len(valid_indices)

        if N < 3:
            return ref_id   # 不足 3 张有效图，无法组成三元组

        # 两两组合检查角度差
        for a in range(N - 1):
            i = valid_indices[a]
            # if facedetect_v1_frame_index[i] > total_frame:
            #     continue
            angle_i = face_crop_angle[i][0]["angle"]

            for b in range(a + 1, N):
                j = valid_indices[b]
                # if facedetect_v1_frame_index[j] > total_frame:
                #     continue
                angle_j = face_crop_angle[j][0]["angle"]

                # 判断是否满足阈值
                if (
                    abs(angle_i["yaw"]   - angle_j["yaw"])   > angle_threshold or
                    abs(angle_i["pitch"] - angle_j["pitch"]) > angle_threshold or
                    abs(angle_i["roll"]  - angle_j["roll"])  > angle_threshold
                ):
                    # 找第三个 k
                    for c in range(N):
                        k = valid_indices[c]
                        # if facedetect_v1_frame_index[k] > total_frame:
                        #     continue
                        if k != i and k != j:
                            ref_id.append([i, j, k])
                            break

                    try_num += 1
                    if try_num >= max_try or len(ref_id) >= need_max:
                        return ref_id

        return ref_id
    def crop_and_resize(self, image, target_height, target_width):
        width, height = image.size
        scale = max(target_width / width, target_height / height)
        image = torchvision.transforms.functional.resize(
            image,
            (round(height*scale), round(width*scale)),
            interpolation=torchvision.transforms.InterpolationMode.BILINEAR
        )
        image = torchvision.transforms.functional.center_crop(image, (target_height, target_width))
        return image
    
    def get_height_width(self, image):
        if self.dynamic_resolution:
            width, height = image.size
            if width * height > self.max_pixels:
                scale = (width * height / self.max_pixels) ** 0.5
                height, width = int(height / scale), int(width / scale)
            height = height // self.height_division_factor * self.height_division_factor
            width = width // self.width_division_factor * self.width_division_factor
        else:
            height, width = self.height, self.width
        return height, width
    
    # def 
    #     img_ratio = img.width / img.height
    #     target_ratio = w / h
    #     if img_ratio > target_ratio:  # Image is wider than target
    #         new_width = w
    #         new_height = int(new_width / img_ratio)
    #     else:  # Image is taller than target
    #         new_height = h
    #         new_width = int(new_height * img_ratio)
        
    #     # img = img.resize((new_width, new_height), Image.ANTIALIAS)
    #     img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)

    #     # Create a new image with the target size and place the resized image in the center
    #     delta_w = w - img.size[0]
    #     delta_h = h - img.size[1]
    #     padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
    #     new_img = ImageOps.expand(img, padding, fill=(255, 255, 255))

    def resize_ref(self, img, target_h, target_w):
        h = target_h
        w = target_w
        img = img.convert("RGB")
        # Calculate the required size to keep aspect ratio and fill the rest with padding.
        img_ratio = img.width / img.height
        target_ratio = w / h
        
        if img_ratio > target_ratio:  # Image is wider than target
            new_width = w
            new_height = int(new_width / img_ratio)
        else:  # Image is taller than target
            new_height = h
            new_width = int(new_height * img_ratio)
        
        # img = img.resize((new_width, new_height), Image.ANTIALIAS)
        img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)

        # Create a new image with the target size and place the resized image in the center
        delta_w = w - img.size[0]
        delta_h = h - img.size[1]
        padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
        new_img = ImageOps.expand(img, padding, fill=(255, 255, 255))

        return new_img


    def load_video_crop_ref_image(self, video_path=None, ref_id_all=[[]]):
        ### fps 转化
        reader = imageio.get_reader(video_path)
        meta = reader.get_meta_data()
        original_fps = meta.get("fps", 24)
        target_fps = 16
        duration_seconds = 5
        target_frames = target_fps * duration_seconds + 1   # = 80 frames

        # ---- 获取原视频帧数 ----
        try:
            total_original_frames = reader.count_frames()
        except:
            total_original_frames = int(meta.get("duration", 5) * original_fps)



        # ---- 需要多少原始帧（5秒）----
        need_orig_frames = int(original_fps * duration_seconds)

        # ---- Case 1: 原视频 >= 5秒 → 随机选择 5 秒起点 ----
        if total_original_frames > need_orig_frames:
            max_start = total_original_frames - need_orig_frames
            start_frame = random.randint(0, max_start)
            segment_start = start_frame
            segment_end = start_frame + need_orig_frames
        else:
            # ---- Case 2: 原视频 < 5秒 → 用全部帧 ----
            segment_start = 0
            segment_end = total_original_frames

        # ---- 均匀采样 80 帧 ----
        sample_ids = np.linspace(segment_start, segment_end - 1, num=target_frames, dtype=int)

        frames = []
        for frame_id in sample_ids:
            frame = reader.get_data(int(frame_id))
            frame = Image.fromarray(frame)
            frame = self.crop_and_resize(frame, *self.get_height_width(frame))
            frames.append(frame)

        # ===========================
        #  选择参考图部分（你要求的）
        # ===========================

        # 1）从 ref_images_all（三维 list）里随机选一组
        #    ref_images_all = [ [img1, img2, img3], [imgA, imgB, imgC], ... ]
        ref_group = random.choice(ref_id_all)

        # 2）检查资源是否足够
        if len(ref_group) < self.ref_num:
            raise ValueError(f"需要 {self.ref_num} 张参考图，但该组只有 {len(ref_group)} 张。")

        # 3）从该组中随机选 self.ref_num 张
        selected_refs = random.sample(ref_group, self.ref_num)
        random.shuffle(selected_refs)
        
        ref_images = []
        for sf in selected_refs:
            top, height, width, left, frame_index = sf
            # import pdb; pdb.set_trace()
            if frame_index > total_original_frames:
                print(f"{video_path}, frame_index({frame_index}) out of range")
            frame = reader.get_data(int(frame_index))
            frame = Image.fromarray(frame)
            xmin, ymin, xmax, ymax = left, top, left + width, top + height
            cropped_image = frame.crop((xmin, ymin, xmax, ymax)).convert("RGB")
            cropped_image = self.resize_ref(cropped_image, self.height, self.width)
            # Calculate the required size to keep aspect ratio and fill the rest with padding.
            ref_images.append(cropped_image)
        reader.close()

        return frames, ref_images

    def __getitem__(self, index):
        max_retry = 10   # 最多重试 10 次，避免死循环
        retry = 0

        while retry < max_retry:
            # ----- 选择 train / test 数据 -----
            if self.training:
                meta_data = self.data_train[index % len(self.data_train)]
            else:
                meta_data = self.data_test[index % len(self.data_test)]

            video_path = meta_data['video_path']
            meta_prompt = meta_data['meta_prompt']
            ref_id_all = meta_data['ref_id_all']

            # ----- 尝试读取 video + ref -----
            try:
                input_video, ref_images = self.load_video_crop_ref_image(
                    video_path=video_path,
                    ref_id_all=ref_id_all
                )
            except Exception as e:
                print("❌ Exception in load_video_crop_ref_image")
                print(f"   video_path: {video_path}")
                print(f"   error type: {type(e).__name__}")
                print(f"   error msg : {e}")

                # 打印 traceback，定位问题更容易
                import traceback
                traceback.print_exc()
                input_video = None
                ref_images = None
            # ----- 如果成功，并且 video 不为空，返回结果 -----
            if input_video is not None and len(input_video) > 0:
                return {
                    "global_caption": None,
                    "shot_num": 1,
                    "pre_shot_caption": [],
                    "single_caption": meta_prompt["single_prompt"],
                    "video": input_video,
                    "ref_num": self.ref_num,
                    "ref_images": ref_images,
                    "video_path": video_path
                }

            # ----- 如果失败，换 index，并继续尝试 -----
            retry += 1
            index = random.randint(0, len(self.data_train) - 1 if self.training else len(self.data_test) - 1)

        # 若 10 次都失败，返回最后一次的错误内容
        raise RuntimeError(f"❌ [Dataset] Failed to load video/ref after {max_retry} retries.")

    def __len__(self):
        if self.training:
            return len(self.data_train)
        else:
            return len(self.data_test)

if __name__ == '__main__':
    from torch.utils.data import DataLoader
    dataset = MulltiShot_MultiView_Dataset(length=49, resolution=(384, 640), training=True)
    print(len(dataset))
    metadata = dataset[0]
    # results = dataset[0]
    # loader = DataLoader(
    #     dataset,
    #     batch_size=1,          # 视频一般 batch=1
    #     shuffle=False,         # 你想打乱就 True
    #     num_workers=10,         # ⭐ 重点：开启 8 个子进程并行加载
    #     pin_memory=True,
    #     prefetch_factor=2,     # 每个 worker 预读取 2 个样本
    #     collate_fn=lambda x: x[0],   # ⭐ 不做任何 collate
    # )

    # for batch in tqdm(loader):
    #     pass
    for i in tqdm(range(len(dataset))):
        file = dataset[i]

    assert 0