dataset/dataset_cls.py

import os
import json
import jsonlines
import torch
import math
import random
import cv2

from tqdm import tqdm
from collections import OrderedDict

from PIL import Image
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True

import numpy as np
import subprocess
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torchvision.transforms.functional import InterpolationMode
from torchvision.transforms import functional as F


class ImageTextDataset(Dataset):
    """
        Usage:
            The dataset class for image-text pairs, used for image generation training
            It supports multi-aspect ratio training
        params:
            anno_file: The annotation file list
            add_normalize: whether to normalize the input image pixel to [-1, 1], default: True
            ratios: The aspect ratios during training, format: width / height
            sizes: The resoultion of training images, format: (width, height)
    """
    def __init__(
        self, anno_file, add_normalize=True,
        ratios=[1/1, 3/5, 5/3], 
        sizes=[(1024, 1024), (768, 1280), (1280, 768)],
        crop_mode='random', p_random_ratio=0.0,
    ):  
        # Ratios and Sizes : (w h)
        super().__init__()
        
        self.image_annos = []
        if not isinstance(anno_file, list):
            anno_file = [anno_file]

        for anno_file_ in anno_file:
            print(f"Load image annotation files from {anno_file_}")
            with jsonlines.open(anno_file_, 'r') as reader:
                for item in reader:
                    self.image_annos.append(item)

        print(f"Totally Remained {len(self.image_annos)} images")

        transform_list = [
            transforms.ToTensor(),
        ]    
    
        if add_normalize:
            transform_list.append(transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)))
        
        self.transform = transforms.Compose(transform_list)

        print(f"Transform List is {transform_list}")

        assert crop_mode in ['center', 'random']
        self.crop_mode = crop_mode
        self.ratios = ratios
        self.sizes = sizes
        self.p_random_ratio = p_random_ratio

    def get_closest_size(self, x):
        if self.p_random_ratio > 0 and np.random.rand() < self.p_random_ratio:
            best_size_idx = np.random.randint(len(self.ratios))
        else:
            w, h = x.width, x.height
            best_size_idx = np.argmin([abs(w/h-r) for r in self.ratios])
        return self.sizes[best_size_idx]

    def get_resize_size(self, orig_size, tgt_size):
        if (tgt_size[1]/tgt_size[0] - 1) * (orig_size[1]/orig_size[0] - 1) >= 0:
            alt_min = int(math.ceil(max(tgt_size)*min(orig_size)/max(orig_size)))
            resize_size = max(alt_min, min(tgt_size))
        else:
            alt_max = int(math.ceil(min(tgt_size)*max(orig_size)/min(orig_size)))
            resize_size = max(alt_max, max(tgt_size))
        return resize_size

    def __len__(self):
        return len(self.image_annos)

    def __getitem__(self, index):
        image_anno = self.image_annos[index]

        try:
            img = Image.open(image_anno['image']).convert("RGB")
            text = image_anno['text']

            assert isinstance(text, str), "Text should be str"

            size = self.get_closest_size(img)
            resize_size = self.get_resize_size((img.width, img.height), size)

            img = transforms.functional.resize(img, resize_size, interpolation=transforms.InterpolationMode.BICUBIC, antialias=True)
        
            if self.crop_mode == 'center':
                img = transforms.functional.center_crop(img, (size[1], size[0]))
            elif self.crop_mode == 'random':
                img = transforms.RandomCrop((size[1], size[0]))(img)
            else:
                img = transforms.functional.center_crop(img, (size[1], size[0]))

            image_tensor = self.transform(img)
            
            return {
                "video": image_tensor,    # using keyname `video`, to be compatible with video
                "text" : text,
                "identifier": 'image',
            }
            
        except Exception as e:
            print(f'Load Image Error with {e}')
            return self.__getitem__(random.randint(0, self.__len__() - 1))


class LengthGroupedVideoTextDataset(Dataset):
    """
        Usage:
            The dataset class for video-text pairs, used for video generation training
            It groups the video with the same frames together
            Now only supporting fixed resolution during training
        params:
            anno_file: The annotation file list
            max_frames: The maximum temporal lengths (This is the vae latent temporal length) 16 => (16 - 1) * 8 + 1 = 121 frames
            load_vae_latent: Loading the pre-extracted vae latents during training, we recommend to extract the latents in advance
                to reduce the time cost per batch
            load_text_fea: Loading the pre-extracted text features during training, we recommend to extract the prompt textual features
                in advance, since the T5 encoder will cost many GPU memories
    """
    
    def __init__(self, anno_file, max_frames=16, resolution='384p', load_vae_latent=True, load_text_fea=True):
        super().__init__()

        self.video_annos = []
        self.max_frames = max_frames
        self.load_vae_latent = load_vae_latent
        self.load_text_fea = load_text_fea
        self.resolution = resolution

        assert load_vae_latent, "Now only support loading vae latents, we will support to directly load video frames in the future"

        if not isinstance(anno_file, list):
            anno_file = [anno_file]

        for anno_file_ in anno_file:
            with jsonlines.open(anno_file_, 'r') as reader:
                for item in tqdm(reader):
                    self.video_annos.append(item)
        
        print(f"Totally Remained {len(self.video_annos)} videos") 

    def __len__(self):
        return len(self.video_annos)

    def __getitem__(self, index):
        try:
            video_anno = self.video_annos[index]
            text = video_anno['text']
            latent_path = video_anno['latent']
            latent = torch.load(latent_path, map_location='cpu')  # loading the pre-extracted video latents

            # TODO: remove the hard code latent shape checking
            if self.resolution == '384p':
                assert latent.shape[-1] == 640 // 8
                assert latent.shape[-2] == 384 // 8
            else:
                assert self.resolution == '768p'
                assert latent.shape[-1] == 1280 // 8
                assert latent.shape[-2] == 768 // 8

            cur_temp = latent.shape[2]
            cur_temp = min(cur_temp, self.max_frames)

            video_latent = latent[:,:,:cur_temp].float()
            assert video_latent.shape[1] == 16

            if self.load_text_fea:
                text_fea_path = video_anno['text_fea']
                text_fea = torch.load(text_fea_path, map_location='cpu')
                return {
                    'video': video_latent,
                    'prompt_embed': text_fea['prompt_embed'],
                    'prompt_attention_mask': text_fea['prompt_attention_mask'],
                    'pooled_prompt_embed': text_fea['pooled_prompt_embed'],
                    "identifier": 'video',
                }

            else:
                return {
                    'video': video_latent,
                    'text': text,
                    "identifier": 'video',
                }

        except Exception as e:
            print(f'Load Video Error with {e}')
            return self.__getitem__(random.randint(0, self.__len__() - 1))


class VideoFrameProcessor:
    # load a video and transform
    def __init__(self, resolution=256, num_frames=24, add_normalize=True, sample_fps=24):
    
        image_size = resolution

        transform_list = [
            transforms.Resize(image_size, interpolation=InterpolationMode.BICUBIC, antialias=True),
            transforms.CenterCrop(image_size),
        ]
        
        if add_normalize:
            transform_list.append(transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)))

        print(f"Transform List is {transform_list}")
        self.num_frames = num_frames
        self.transform = transforms.Compose(transform_list)
        self.sample_fps = sample_fps

    def __call__(self, video_path):
        try:
            video_capture = cv2.VideoCapture(video_path)
            fps = video_capture.get(cv2.CAP_PROP_FPS)
            frames = []

            while True:
                flag, frame = video_capture.read()
                if not flag:
                    break

                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                frame = torch.from_numpy(frame)
                frame = frame.permute(2, 0, 1)
                frames.append(frame)

            video_capture.release()
            sample_fps = self.sample_fps
            interval = max(int(fps / sample_fps), 1)
            frames = frames[::interval]

            if len(frames) < self.num_frames:
                num_frame_to_pack = self.num_frames - len(frames)
                recurrent_num = num_frame_to_pack // len(frames)
                frames = frames + recurrent_num * frames + frames[:(num_frame_to_pack % len(frames))]
                assert len(frames) >= self.num_frames, f'{len(frames)}'

            start_indexs = list(range(0, max(0, len(frames) - self.num_frames + 1)))            
            start_index = random.choice(start_indexs)

            filtered_frames = frames[start_index : start_index+self.num_frames]
            assert len(filtered_frames) == self.num_frames, f"The sampled frames should equals to {self.num_frames}"

            filtered_frames = torch.stack(filtered_frames).float() / 255
            filtered_frames = self.transform(filtered_frames)
            filtered_frames = filtered_frames.permute(1, 0, 2, 3)

            return filtered_frames, None
            
        except Exception as e:
            print(f"Load video: {video_path} Error, Exception {e}")
            return None, None


class VideoDataset(Dataset):
    def __init__(self, anno_file, resolution=256, max_frames=6, add_normalize=True):
        super().__init__()

        self.video_annos = []
        self.max_frames = max_frames

        if not isinstance(anno_file, list):
            anno_file = [anno_file]

        print(f"The training video clip frame number is {max_frames} ")

        for anno_file_ in anno_file:
            print(f"Load annotation file from {anno_file_}")

            with jsonlines.open(anno_file_, 'r') as reader:
                for item in tqdm(reader):
                    self.video_annos.append(item)
        
        print(f"Totally Remained {len(self.video_annos)} videos")
        
        self.video_processor = VideoFrameProcessor(resolution, max_frames, add_normalize)        

    def __len__(self):
        return len(self.video_annos)

    def __getitem__(self, index):
        video_anno = self.video_annos[index]
        video_path = video_anno['video']

        try:
            video_tensors, video_frames = self.video_processor(video_path)

            assert video_tensors.shape[1] == self.max_frames
            
            return {
                "video": video_tensors,
                "identifier": 'video',
            }

        except Exception as e:
            print('Loading Video Error with {e}')
            return self.__getitem__(random.randint(0, self.__len__() - 1))


class ImageDataset(Dataset):
    def __init__(self, anno_file, resolution=256, max_frames=8, add_normalize=True):
        super().__init__()

        self.image_annos = []
        self.max_frames = max_frames   
        image_paths = []

        if not isinstance(anno_file, list):
            anno_file = [anno_file]

        for anno_file_ in anno_file:
            print(f"Load annotation file from {anno_file_}")
            with jsonlines.open(anno_file_, 'r') as reader:
                for item in tqdm(reader):
                    image_paths.append(item['image'])
        
        print(f"Totally Remained {len(image_paths)} images")

        # pack multiple frames
        for idx in range(0, len(image_paths), self.max_frames):
            image_path_shard = image_paths[idx : idx + self.max_frames]
            if len(image_path_shard) < self.max_frames:
                image_path_shard = image_path_shard + image_paths[:self.max_frames - len(image_path_shard)]
            assert len(image_path_shard) == self.max_frames
            self.image_annos.append(image_path_shard)

        image_size = resolution
        transform_list = [
            transforms.Resize(image_size, interpolation=InterpolationMode.BICUBIC, antialias=True),
            transforms.CenterCrop(image_size),
            transforms.ToTensor(),
        ]    
        if add_normalize:
            transform_list.append(transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)))

        print(f"Transform List is {transform_list}")
        self.transform = transforms.Compose(transform_list)

    def __len__(self):
        return len(self.image_annos)

    def __getitem__(self, index):
        image_paths = self.image_annos[index]

        try:
            packed_pil_frames = [Image.open(image_path).convert("RGB") for image_path in image_paths]
            filtered_frames = [self.transform(frame) for frame in packed_pil_frames]
            filtered_frames = torch.stack(filtered_frames)  # [t, c, h, w]
            filtered_frames = filtered_frames.permute(1, 0, 2, 3) # [c, t, h, w]
            
            return {
                "video": filtered_frames,
                "identifier": 'image',
            }

        except Exception as e:
            print(f'Load Images Error with {e}')
            return self.__getitem__(random.randint(0, self.__len__() - 1))