Dataset/sakuga_dataset.py

from datetime import timedelta
from pathlib import Path
from typing import List, Optional, Tuple, Union
from torch.utils.data import DataLoader, Dataset
from tqdm.auto import tqdm
from torchvision.transforms.functional import center_crop, resize
from torchvision.transforms import InterpolationMode
import torchvision.transforms as TT
import numpy as np
import accelerate
import torch
import pandas as pd
from pathlib import PosixPath
import os 
from datetime import datetime

try:
    import decord
except ImportError:
    raise ImportError(
        "The `decord` package is required for loading the video dataset. Install with `pip install decord`"
    )
decord.bridge.set_bridge("torch")


class Sakuga_Dataset(Dataset):
    def __init__(
        self,
        instance_data_root: Optional[str] = None,
        sketch_data_root: Optional[str] = None,
        dataset_name: Optional[str] = None,
        dataset_config_name: Optional[str] = None,
        caption_column: str = "text",
        video_column: str = "video",
        height: int = 480,
        width: int = 720,
        video_reshape_mode: str = "center",
        fps: int = 8,
        max_num_frames: int = 49,
        skip_frames_start: int = 0,
        skip_frames_end: int = 0,
        cache_dir: Optional[str] = None,
        id_token: Optional[str] = None,
        data_information:  Optional[str] = None,
        stage: Optional[str] = "1",
    ) -> None:
        super().__init__()

        self.instance_data_root = Path(instance_data_root) if instance_data_root is not None else None
        self.sketch_data_root = Path(sketch_data_root) if sketch_data_root is not None else None
        self.dataset_name = dataset_name
        self.dataset_config_name = dataset_config_name
        self.caption_column = caption_column
        self.video_column = video_column
        self.height = height
        self.width = width
        self.video_reshape_mode = video_reshape_mode
        self.fps = fps
        self.max_num_frames = max_num_frames
        self.skip_frames_start = skip_frames_start
        self.skip_frames_end = skip_frames_end
        self.cache_dir = cache_dir
        self.id_token = id_token or ""
        self.stage=stage
        
        '''
        if dataset_name is not None:
            self.instance_prompts, self.instance_video_paths = self._load_dataset_from_hub()
        else:
            self.instance_prompts, self.instance_video_paths = self._load_dataset_from_local_path()
        '''
        
        self.data_information=pd.read_parquet(data_information)
        self.num_instance_videos = self.data_information.shape[0]


        #self.detector = LineartDetector('cpu')
        #TODO: here just point the cuda maybe have some problem

        #we put the preprocess_data() in the get_item function
        #self.instance_videos = self._preprocess_data()
        #here, how to make it in the get_item?

    def __len__(self):
        return self.num_instance_videos
    
    def encode_video(self, video,vae,device):
        

        #vae,device
        video = video.to(device, dtype=vae.dtype).unsqueeze(0)
        video = video.permute(0, 2, 1, 3, 4)  # [B, C, F, H, W]
        image = video[:, :, :1].clone()

        latent_dist = vae.encode(video).latent_dist

        image_noise_sigma = torch.normal(mean=-3.0, std=0.5, size=(1,), device=image.device)
        image_noise_sigma = torch.exp(image_noise_sigma).to(dtype=image.dtype)
        noisy_image = torch.randn_like(image) * image_noise_sigma[:, None, None, None, None]
        image_latent_dist = vae.encode(noisy_image).latent_dist

        return latent_dist, image_latent_dist

    def read_video(self,video_path):
        filename=PosixPath(video_path)

        #this part have some wrong things
        try:
            video_reader = decord.VideoReader(uri=filename.as_posix())
            video_num_frames = len(video_reader)

            #需不需要这里强制一下从第10帧开始？
            start_frame = min(self.skip_frames_start, video_num_frames)
            end_frame = max(0, video_num_frames - self.skip_frames_end)
            if end_frame <= start_frame:
                frames = video_reader.get_batch([start_frame])
            elif end_frame - start_frame <= self.max_num_frames:
                frames = video_reader.get_batch(list(range(start_frame, end_frame)))
            else:
                #this has problem
                #indices = list(range(start_frame, end_frame, (end_frame - start_frame) // self.max_num_frames))
                indices=list(range(start_frame,self.max_num_frames))
                frames = video_reader.get_batch(indices)

            # Ensure that we don't go over the limit
            frames = frames[: self.max_num_frames]
            selected_num_frames = frames.shape[0]

            # Choose first (4k + 1) frames as this is how many is required by the VAE
            remainder = (3 + (selected_num_frames % 4)) % 4
            if remainder != 0:
                frames = frames[:-remainder]
            selected_num_frames = frames.shape[0]

            assert (selected_num_frames - 1) % 4 == 0

            # Training transforms
            
            frames = frames.permute(0, 3, 1, 2) # [F, C, H, W]
            #print("frame",frames.shape)
            frames = self._resize_for_rectangle_crop(frames)
            final_frames = frames.contiguous()
            if final_frames.dim()==3:
                final_frames=final_frames.unsqueeze(0)
            return final_frames
        except:
            return None
        
        
    
    def __getitem__(self, index):
        
        #output_video=self.encode_video(video,vae,device)
        #_encode_instance_video=self.encode_video(self.instance_prompts[index],device=)
        
        #处理selfinstance_videos
        
        folder_path=os.path.join(self.instance_data_root, str(self.data_information.iloc[index]['identifier_video']))
        
        #sketch_path=os.path.join(self.sketch_data_root, str(self.data_information.iloc[index]['identifier_video']))
        #注意这里的sketch存的是[0,255]的信息，不需要1-了，但是之后可能还是要变成2值的操作，得看一下如何2值化
        
        
        #这里寻找id号的过程有点问题，不太对
        #这个identifier是不是就是对应的片段名称？是的，identifier后面的对应的就是sence_number
        #indices = self.data_information.index[self.data_information['identifier_video'] == self.data_information.iloc[index]['identifier_video']].tolist()
        
        # video_name=self.data_information.iloc[index]['identifier_video']
        # sence_number=self.data_information.iloc[index]["identifier"].split(":")[1]
        # data_name=f"{video_name}-Scene-{int(sence_number):03d}.mp4"

        
        frames=self.data_information.iloc[index]["start_frame"]
        video_name=self.data_information.iloc[index]["identifier"].split(':')[0]
        #print(frames)
        
        data_path_1=f'{video_name}-Scene-{frames}.mp4'
        data_path_2=f'{video_name}-Scene-{frames+1}.mp4'
        data_path_3=f'{video_name}-Scene-{frames-1}.mp4'
        fd1=os.path.join(folder_path,data_path_1)
        #sketch_fd1=os.path.join(sketch_path,data_path_1)
        fd2=os.path.join(folder_path,data_path_2)
        #sketch_fd2=os.path.join(sketch_path,data_path_2)
        fd3=os.path.join(folder_path,data_path_3)
        #sketch_fd3=os.path.join(sketch_path,data_path_2)
        
        #print(fd1)
        
        if os.path.exists(fd1):
            file_path=fd1
        elif os.path.exists(fd2):
            file_path=fd2
        elif os.path.exists(fd3):
            file_path=fd3   
        

        
        prompt=self.data_information.iloc[index]["text_description"]
                
        final_frames=self.read_video(PosixPath(file_path))
        #final_sketch_frames=self.read_video(PosixPath(sketch_file_path))
        final_sketch_frames=None
        
        
        
        instance_prompt =  prompt + self.id_token
        
        return {
            "instance_prompt": instance_prompt,
            "instance_video": final_frames,
            "file_path":file_path,
            "sketch_video": final_sketch_frames,
            "instance_image": None,
            #"instance_sketch": final_sketch,
        }

    def _load_dataset_from_hub(self):
        try:
            from datasets import load_dataset
        except ImportError:
            raise ImportError(
                "You are trying to load your data using the datasets library. If you wish to train using custom "
                "captions please install the datasets library: `pip install datasets`. If you wish to load a "
                "local folder containing images only, specify --instance_data_root instead."
            )

        # Downloading and loading a dataset from the hub. See more about loading custom images at
        # https://huggingface.co/docs/datasets/v2.0.0/en/dataset_script
        dataset = load_dataset(
            self.dataset_name,
            self.dataset_config_name,
            cache_dir=self.cache_dir,
        )
        column_names = dataset["train"].column_names

        if self.video_column is None:
            video_column = column_names[0]
            #logger.info(f"`video_column` defaulting to {video_column}")
            print(f"`video_column` defaulting to {video_column}")
        else:
            video_column = self.video_column
            if video_column not in column_names:
                raise ValueError(
                    f"`--video_column` value '{video_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}"
                )

        if self.caption_column is None:
            caption_column = column_names[1]
            #logger.info(f"`caption_column` defaulting to {caption_column}")
            print(f"`caption_column` defaulting to {caption_column}")
        else:
            caption_column = self.caption_column
            if self.caption_column not in column_names:
                raise ValueError(
                    f"`--caption_column` value '{self.caption_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}"
                )

        instance_prompts = dataset["train"][caption_column]
        instance_videos = [Path(self.instance_data_root, filepath) for filepath in dataset["train"][video_column]]

        return instance_prompts, instance_videos

    def _load_dataset_from_local_path(self):
        if not self.instance_data_root.exists():
            raise ValueError("Instance videos root folder does not exist")

        prompt_path = self.instance_data_root.joinpath(self.caption_column)
        video_path = self.instance_data_root.joinpath(self.video_column)

        if not prompt_path.exists() or not prompt_path.is_file():
            raise ValueError(
                "Expected `--caption_column` to be path to a file in `--instance_data_root` containing line-separated text prompts."
            )
        if not video_path.exists() or not video_path.is_file():
            raise ValueError(
                "Expected `--video_column` to be path to a file in `--instance_data_root` containing line-separated paths to video data in the same directory."
            )

        with open(prompt_path, "r", encoding="utf-8") as file:
            instance_prompts = [line.strip() for line in file.readlines() if len(line.strip()) > 0]
        with open(video_path, "r", encoding="utf-8") as file:
            instance_videos = [
                self.instance_data_root.joinpath(line.strip()) for line in file.readlines() if len(line.strip()) > 0
            ]

        if any(not path.is_file() for path in instance_videos):
            raise ValueError(
                "Expected '--video_column' to be a path to a file in `--instance_data_root` containing line-separated paths to video data but found atleast one path that is not a valid file."
            )

        return instance_prompts, instance_videos

    def _resize_for_rectangle_crop(self, arr):
        image_size = self.height, self.width
        reshape_mode = self.video_reshape_mode
        if arr.shape[3] / arr.shape[2] > image_size[1] / image_size[0]:
            arr = resize(
                arr,
                size=[image_size[0], int(arr.shape[3] * image_size[0] / arr.shape[2])],
                interpolation=InterpolationMode.BICUBIC,
            )
        else:
            arr = resize(
                arr,
                size=[int(arr.shape[2] * image_size[1] / arr.shape[3]), image_size[1]],
                interpolation=InterpolationMode.BICUBIC,
            )

        h, w = arr.shape[2], arr.shape[3]
        arr = arr.squeeze(0)

        delta_h = h - image_size[0]
        delta_w = w - image_size[1]

        if reshape_mode == "random" or reshape_mode == "none":
            top = np.random.randint(0, delta_h + 1)
            left = np.random.randint(0, delta_w + 1)
        elif reshape_mode == "center":
            top, left = delta_h // 2, delta_w // 2
        else:
            raise NotImplementedError
        arr = TT.functional.crop(arr, top=top, left=left, height=image_size[0], width=image_size[1])
        return arr


    # here process the all data, we should make these processed in the get_item or other position

    def _preprocess_data(self):


        decord.bridge.set_bridge("torch")

        progress_dataset_bar = tqdm(
            range(0, len(self.instance_video_paths)),
            desc="Loading progress resize and crop videos",
        )

        videos = []

        for filename in self.instance_video_paths:
            video_reader = decord.VideoReader(uri=filename.as_posix())
            video_num_frames = len(video_reader)

            start_frame = min(self.skip_frames_start, video_num_frames)
            end_frame = max(0, video_num_frames - self.skip_frames_end)
            if end_frame <= start_frame:
                frames = video_reader.get_batch([start_frame])
            elif end_frame - start_frame <= self.max_num_frames:
                frames = video_reader.get_batch(list(range(start_frame, end_frame)))
            else:
                indices = list(range(start_frame, end_frame, (end_frame - start_frame) // self.max_num_frames))
                frames = video_reader.get_batch(indices)

            # Ensure that we don't go over the limit
            frames = frames[: self.max_num_frames]
            selected_num_frames = frames.shape[0]

            # Choose first (4k + 1) frames as this is how many is required by the VAE
            remainder = (3 + (selected_num_frames % 4)) % 4
            if remainder != 0:
                frames = frames[:-remainder]
            selected_num_frames = frames.shape[0]

            assert (selected_num_frames - 1) % 4 == 0

            # Training transforms
            
            frames = frames.permute(0, 3, 1, 2) # [F, C, H, W]
            progress_dataset_bar.set_description(
                f"Loading progress Resizing video from {frames.shape[2]}x{frames.shape[3]} to {self.height}x{self.width}"
            )
            frames = self._resize_for_rectangle_crop(frames)  #here the tensor should be processed to right size
            
            frames = (frames - 127.5) / 127.5
            videos.append(frames.contiguous())  # [F, C, H, W]
            progress_dataset_bar.update(1)

        progress_dataset_bar.close()

        return videos
        



if __name__=="__main__":
    train_dataset = Sakuga_Dataset(
        instance_data_root='',
        height= 480,
        width=  720,
        video_reshape_mode="center",
        fps=8,
        max_num_frames=49,
        skip_frames_start=0,
        skip_frames_end=0,
        cache_dir="~/.cache",
        id_token="",
        data_information="../../../Datasets/SakugaDataset/parquet/fliter_59_aesthetic_precise.parquet"
    )
    data=train_dataset.__getitem__(0)
    print(data["instance_video"].shape)