65 lines
2.0 KiB
Python
65 lines
2.0 KiB
Python
import os
|
|
from glob import glob
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torchvision
|
|
from PIL import Image
|
|
from torch.utils.data import Dataset
|
|
|
|
from opensora.utils.dataset_utils import DecordInit, is_image_file
|
|
|
|
|
|
class ExtractVideo2Feature(Dataset):
|
|
def __init__(self, args, transform):
|
|
self.data_path = args.data_path
|
|
self.transform = transform
|
|
self.v_decoder = DecordInit()
|
|
self.samples = list(glob(f'{self.data_path}'))
|
|
|
|
def __len__(self):
|
|
return len(self.samples)
|
|
|
|
def __getitem__(self, idx):
|
|
video_path = self.samples[idx]
|
|
video = self.decord_read(video_path)
|
|
video = self.transform(video) # T C H W -> T C H W
|
|
return video, video_path
|
|
|
|
def tv_read(self, path):
|
|
vframes, aframes, info = torchvision.io.read_video(filename=path, pts_unit='sec', output_format='TCHW')
|
|
total_frames = len(vframes)
|
|
frame_indice = list(range(total_frames))
|
|
video = vframes[frame_indice]
|
|
return video
|
|
|
|
def decord_read(self, path):
|
|
decord_vr = self.v_decoder(path)
|
|
total_frames = len(decord_vr)
|
|
frame_indice = list(range(total_frames))
|
|
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
|
|
|
|
|
|
|
|
class ExtractImage2Feature(Dataset):
|
|
def __init__(self, args, transform):
|
|
self.data_path = args.data_path
|
|
self.transform = transform
|
|
self.data_all = list(glob(f'{self.data_path}'))
|
|
|
|
def __len__(self):
|
|
return len(self.data_all)
|
|
|
|
def __getitem__(self, index):
|
|
path = self.data_all[index]
|
|
video_frame = torch.as_tensor(np.array(Image.open(path), dtype=np.uint8, copy=True)).unsqueeze(0)
|
|
video_frame = video_frame.permute(0, 3, 1, 2)
|
|
video_frame = self.transform(video_frame) # T C H W
|
|
# video_frame = video_frame.transpose(0, 1) # T C H W -> C T H W
|
|
|
|
return video_frame, path
|
|
|