110 lines
4.5 KiB
Python
110 lines
4.5 KiB
Python
import json
|
|
import os, io, csv, math, random
|
|
import numpy as np
|
|
import torchvision
|
|
from einops import rearrange
|
|
from decord import VideoReader
|
|
|
|
import torch
|
|
import torchvision.transforms as transforms
|
|
from torch.utils.data.dataset import Dataset
|
|
from tqdm import tqdm
|
|
|
|
from opensora.utils.dataset_utils import DecordInit
|
|
from opensora.utils.utils import text_preprocessing
|
|
|
|
|
|
|
|
class T2V_dataset(Dataset):
|
|
def __init__(self, args, transform, temporal_sample, tokenizer):
|
|
|
|
# with open(args.data_path, 'r') as csvfile:
|
|
# self.samples = list(csv.DictReader(csvfile))
|
|
self.video_folder = args.video_folder
|
|
self.num_frames = args.num_frames
|
|
self.transform = transform
|
|
self.temporal_sample = temporal_sample
|
|
self.tokenizer = tokenizer
|
|
self.model_max_length = args.model_max_length
|
|
self.v_decoder = DecordInit()
|
|
|
|
with open('sharegpt4v_path_cap_.json', 'r') as f:
|
|
self.samples = json.load(f)
|
|
self.use_image_num = args.use_image_num
|
|
self.use_img_from_vid = args.use_img_from_vid
|
|
if self.use_image_num != 0 and not self.use_img_from_vid:
|
|
self.img_cap_list = self.get_img_cap_list()
|
|
|
|
def __len__(self):
|
|
return len(self.samples)
|
|
|
|
def __getitem__(self, idx):
|
|
try:
|
|
# video = torch.randn(3, 16, 128, 128)
|
|
# input_ids = torch.ones(1, 120).to(torch.long).squeeze(0)
|
|
# cond_mask = torch.cat([torch.ones(1, 60).to(torch.long), torch.ones(1, 60).to(torch.long)], dim=1).squeeze(0)
|
|
# return video, input_ids, cond_mask
|
|
video = self.decord_read(self.samples[idx]['path'])
|
|
video = self.transform(video) # T C H W -> T C H W
|
|
video = video.transpose(0, 1) # T C H W -> C T H W
|
|
text = self.samples[idx]['cap'][0]
|
|
|
|
text = text_preprocessing(text)
|
|
text_tokens_and_mask = self.tokenizer(
|
|
text,
|
|
max_length=self.model_max_length,
|
|
padding='max_length',
|
|
truncation=True,
|
|
return_attention_mask=True,
|
|
add_special_tokens=True,
|
|
return_tensors='pt'
|
|
)
|
|
input_ids = text_tokens_and_mask['input_ids'].squeeze(0)
|
|
cond_mask = text_tokens_and_mask['attention_mask'].squeeze(0)
|
|
|
|
if self.use_image_num != 0 and self.use_img_from_vid:
|
|
select_image_idx = np.linspace(0, self.num_frames-1, self.use_image_num, dtype=int)
|
|
assert self.num_frames >= self.use_image_num
|
|
images = video[:, select_image_idx] # c, num_img, h, w
|
|
video = torch.cat([video, images], dim=1) # c, num_frame+num_img, h, w
|
|
input_ids = torch.stack([input_ids] * (1+self.use_image_num)) # 1+self.use_image_num, l
|
|
cond_mask = torch.stack([cond_mask] * (1+self.use_image_num)) # 1+self.use_image_num, l
|
|
elif self.use_image_num != 0 and not self.use_img_from_vid:
|
|
images, captions = self.img_cap_list[idx]
|
|
raise NotImplementedError
|
|
else:
|
|
pass
|
|
|
|
return video, input_ids, cond_mask
|
|
except Exception as e:
|
|
print(f'Error with {e}, {self.samples[idx]}')
|
|
return self.__getitem__(random.randint(0, self.__len__() - 1))
|
|
|
|
def tv_read(self, path):
|
|
vframes, aframes, info = torchvision.io.read_video(filename=path, pts_unit='sec', output_format='TCHW')
|
|
total_frames = len(vframes)
|
|
|
|
# Sampling video frames
|
|
start_frame_ind, end_frame_ind = self.temporal_sample(total_frames)
|
|
# assert end_frame_ind - start_frame_ind >= self.num_frames
|
|
frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.num_frames, dtype=int)
|
|
|
|
video = vframes[frame_indice] # (T, C, H, W)
|
|
|
|
return video
|
|
|
|
def decord_read(self, path):
|
|
decord_vr = self.v_decoder(path)
|
|
total_frames = len(decord_vr)
|
|
# Sampling video frames
|
|
start_frame_ind, end_frame_ind = self.temporal_sample(total_frames)
|
|
# assert end_frame_ind - start_frame_ind >= self.num_frames
|
|
frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.num_frames, dtype=int)
|
|
|
|
video_data = decord_vr.get_batch(frame_indice).asnumpy()
|
|
video_data = torch.from_numpy(video_data)
|
|
video_data = video_data.permute(0, 3, 1, 2) # (T, H, W, C) -> (T C H W)
|
|
return video_data
|
|
|
|
def get_img_cap_list(self):
|
|
raise NotImplementedError |