|
import json |
|
import random |
|
|
|
import torch |
|
import torchvision.transforms as transforms |
|
from decord import VideoReader |
|
from PIL import Image |
|
from torch.utils.data import Dataset |
|
from transformers import CLIPImageProcessor |
|
|
|
|
|
class HumanDanceDataset(Dataset): |
|
def __init__( |
|
self, |
|
img_size, |
|
img_scale=(1.0, 1.0), |
|
img_ratio=(0.9, 1.0), |
|
drop_ratio=0.1, |
|
data_meta_paths=["./data/fahsion_meta.json"], |
|
sample_margin=30, |
|
): |
|
super().__init__() |
|
|
|
self.img_size = img_size |
|
self.img_scale = img_scale |
|
self.img_ratio = img_ratio |
|
self.sample_margin = sample_margin |
|
|
|
|
|
|
|
|
|
|
|
|
|
vid_meta = [] |
|
for data_meta_path in data_meta_paths: |
|
vid_meta.extend(json.load(open(data_meta_path, "r"))) |
|
self.vid_meta = vid_meta |
|
|
|
self.clip_image_processor = CLIPImageProcessor() |
|
|
|
self.transform = transforms.Compose( |
|
[ |
|
|
|
|
|
|
|
|
|
|
|
|
|
transforms.Resize( |
|
self.img_size, |
|
), |
|
transforms.ToTensor(), |
|
transforms.Normalize([0.5], [0.5]), |
|
] |
|
) |
|
|
|
self.cond_transform = transforms.Compose( |
|
[ |
|
|
|
|
|
|
|
|
|
|
|
|
|
transforms.Resize( |
|
self.img_size, |
|
), |
|
transforms.ToTensor(), |
|
] |
|
) |
|
|
|
self.drop_ratio = drop_ratio |
|
|
|
def augmentation(self, image, transform, state=None): |
|
if state is not None: |
|
torch.set_rng_state(state) |
|
return transform(image) |
|
|
|
def __getitem__(self, index): |
|
video_meta = self.vid_meta[index] |
|
video_path = video_meta["video_path"] |
|
kps_path = video_meta["kps_path"] |
|
|
|
video_reader = VideoReader(video_path) |
|
kps_reader = VideoReader(kps_path) |
|
|
|
assert len(video_reader) == len( |
|
kps_reader |
|
), f"{len(video_reader) = } != {len(kps_reader) = } in {video_path}" |
|
|
|
video_length = len(video_reader) |
|
|
|
margin = min(self.sample_margin, video_length) |
|
|
|
ref_img_idx = random.randint(0, video_length - 1) |
|
if ref_img_idx + margin < video_length: |
|
tgt_img_idx = random.randint(ref_img_idx + margin, video_length - 1) |
|
elif ref_img_idx - margin > 0: |
|
tgt_img_idx = random.randint(0, ref_img_idx - margin) |
|
else: |
|
tgt_img_idx = random.randint(0, video_length - 1) |
|
|
|
ref_img = video_reader[ref_img_idx] |
|
ref_img_pil = Image.fromarray(ref_img.asnumpy()) |
|
tgt_img = video_reader[tgt_img_idx] |
|
tgt_img_pil = Image.fromarray(tgt_img.asnumpy()) |
|
|
|
tgt_pose = kps_reader[tgt_img_idx] |
|
tgt_pose_pil = Image.fromarray(tgt_pose.asnumpy()) |
|
|
|
state = torch.get_rng_state() |
|
tgt_img = self.augmentation(tgt_img_pil, self.transform, state) |
|
tgt_pose_img = self.augmentation(tgt_pose_pil, self.cond_transform, state) |
|
ref_img_vae = self.augmentation(ref_img_pil, self.transform, state) |
|
clip_image = self.clip_image_processor( |
|
images=ref_img_pil, return_tensors="pt" |
|
).pixel_values[0] |
|
|
|
sample = dict( |
|
video_dir=video_path, |
|
img=tgt_img, |
|
tgt_pose=tgt_pose_img, |
|
ref_img=ref_img_vae, |
|
clip_images=clip_image, |
|
) |
|
|
|
return sample |
|
|
|
def __len__(self): |
|
return len(self.vid_meta) |
|
|
