128 lines
5.0 KiB
Python
128 lines
5.0 KiB
Python
import os
|
|
import torch
|
|
import random
|
|
import torch.utils.data as data
|
|
|
|
import numpy as np
|
|
|
|
from PIL import Image
|
|
|
|
from opensora.utils.dataset_utils import is_image_file
|
|
|
|
|
|
class Sky(data.Dataset):
|
|
def __init__(self, args, transform, temporal_sample=None, train=True):
|
|
|
|
self.args = args
|
|
self.data_path = args.data_path
|
|
self.transform = transform
|
|
self.temporal_sample = temporal_sample
|
|
self.num_frames = self.args.num_frames
|
|
self.sample_rate = self.args.sample_rate
|
|
self.data_all = self.load_video_frames(self.data_path)
|
|
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 __getitem__(self, index):
|
|
|
|
vframes = self.data_all[index]
|
|
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, num=self.num_frames, dtype=int) # start, stop, num=50
|
|
|
|
select_video_frames = vframes[frame_indice[0]: frame_indice[-1]+1: self.sample_rate]
|
|
|
|
video_frames = []
|
|
for path in select_video_frames:
|
|
video_frame = torch.as_tensor(np.array(Image.open(path), dtype=np.uint8, copy=True)).unsqueeze(0)
|
|
video_frames.append(video_frame)
|
|
video_clip = torch.cat(video_frames, dim=0).permute(0, 3, 1, 2)
|
|
video_clip = self.transform(video_clip)
|
|
video_clip = video_clip.transpose(0, 1) # T C H W -> C T H W
|
|
|
|
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_clip[:, select_image_idx] # c, num_img, h, w
|
|
video_clip = torch.cat([video_clip, images], dim=1) # c, num_frame+num_img, h, w
|
|
elif self.use_image_num != 0 and not self.use_img_from_vid:
|
|
images, captions = self.img_cap_list[index]
|
|
raise NotImplementedError
|
|
else:
|
|
pass
|
|
|
|
return video_clip, 1
|
|
|
|
def __len__(self):
|
|
return self.video_num
|
|
|
|
def load_video_frames(self, dataroot):
|
|
data_all = []
|
|
frame_list = os.walk(dataroot)
|
|
for _, meta in enumerate(frame_list):
|
|
root = meta[0]
|
|
try:
|
|
frames = [i for i in meta[2] if is_image_file(i)]
|
|
frames = sorted(frames, key=lambda item: int(item.split('.')[0].split('_')[-1]))
|
|
except:
|
|
pass
|
|
# print(meta[0]) # root
|
|
# print(meta[2]) # files
|
|
frames = [os.path.join(root, item) for item in frames if is_image_file(item)]
|
|
if len(frames) > max(0, self.num_frames * self.sample_rate): # need all > (16 * frame-interval) videos
|
|
# if len(frames) >= max(0, self.target_video_len): # need all > 16 frames videos
|
|
data_all.append(frames)
|
|
self.video_num = len(data_all)
|
|
return data_all
|
|
|
|
def get_img_cap_list(self):
|
|
raise NotImplementedError
|
|
|
|
if __name__ == '__main__':
|
|
|
|
import argparse
|
|
import torchvision
|
|
import video_transforms
|
|
import torch.utils.data as data
|
|
|
|
from torchvision import transforms
|
|
from torchvision.utils import save_image
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--num_frames", type=int, default=16)
|
|
parser.add_argument("--frame_interval", type=int, default=4)
|
|
parser.add_argument("--data-path", type=str, default="/path/to/datasets/sky_timelapse/sky_train/")
|
|
config = parser.parse_args()
|
|
|
|
|
|
target_video_len = config.num_frames
|
|
|
|
temporal_sample = video_transforms.TemporalRandomCrop(target_video_len * config.frame_interval)
|
|
trans = transforms.Compose([
|
|
video_transforms.ToTensorVideo(),
|
|
# video_transforms.CenterCropVideo(256),
|
|
video_transforms.CenterCropResizeVideo(256),
|
|
# video_transforms.RandomHorizontalFlipVideo(),
|
|
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
|
|
])
|
|
|
|
taichi_dataset = Sky(config, transform=trans, temporal_sample=temporal_sample)
|
|
print(len(taichi_dataset))
|
|
taichi_dataloader = data.DataLoader(dataset=taichi_dataset, batch_size=1, shuffle=False, num_workers=1)
|
|
|
|
for i, video_data in enumerate(taichi_dataloader):
|
|
print(video_data['video'].shape)
|
|
|
|
# print(video_data.dtype)
|
|
# for i in range(target_video_len):
|
|
# save_image(video_data[0][i], os.path.join('./test_data', '%04d.png' % i), normalize=True, value_range=(-1, 1))
|
|
|
|
# video_ = ((video_data[0] * 0.5 + 0.5) * 255).add_(0.5).clamp_(0, 255).to(dtype=torch.uint8).cpu().permute(0, 2, 3, 1)
|
|
# torchvision.io.write_video('./test_data' + 'test.mp4', video_, fps=8)
|
|
# exit() |