File size: 4,893 Bytes
a814eb3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
import os
import random
from torch.utils.data import Dataset
from PIL import Image
from .augmentation import MotionAugmentation
class VideoMatteDataset(Dataset):
def __init__(self,
videomatte_dir,
background_image_dir,
background_video_dir,
size,
seq_length,
seq_sampler,
transform=None):
self.background_image_dir = background_image_dir
self.background_image_files = os.listdir(background_image_dir)
self.background_video_dir = background_video_dir
self.background_video_clips = sorted(os.listdir(background_video_dir))
self.background_video_frames = [sorted(os.listdir(os.path.join(background_video_dir, clip)))
for clip in self.background_video_clips]
self.videomatte_dir = videomatte_dir
self.videomatte_clips = sorted(os.listdir(os.path.join(videomatte_dir, 'fgr')))
self.videomatte_frames = [sorted(os.listdir(os.path.join(videomatte_dir, 'fgr', clip)))
for clip in self.videomatte_clips]
self.videomatte_idx = [(clip_idx, frame_idx)
for clip_idx in range(len(self.videomatte_clips))
for frame_idx in range(0, len(self.videomatte_frames[clip_idx]), seq_length)]
self.size = size
self.seq_length = seq_length
self.seq_sampler = seq_sampler
self.transform = transform
def __len__(self):
return len(self.videomatte_idx)
def __getitem__(self, idx):
if random.random() < 0.5:
bgrs = self._get_random_image_background()
else:
bgrs = self._get_random_video_background()
fgrs, phas = self._get_videomatte(idx)
if self.transform is not None:
return self.transform(fgrs, phas, bgrs)
return fgrs, phas, bgrs
def _get_random_image_background(self):
with Image.open(os.path.join(self.background_image_dir, random.choice(self.background_image_files))) as bgr:
bgr = self._downsample_if_needed(bgr.convert('RGB'))
bgrs = [bgr] * self.seq_length
return bgrs
def _get_random_video_background(self):
clip_idx = random.choice(range(len(self.background_video_clips)))
frame_count = len(self.background_video_frames[clip_idx])
frame_idx = random.choice(range(max(1, frame_count - self.seq_length)))
clip = self.background_video_clips[clip_idx]
bgrs = []
for i in self.seq_sampler(self.seq_length):
frame_idx_t = frame_idx + i
frame = self.background_video_frames[clip_idx][frame_idx_t % frame_count]
with Image.open(os.path.join(self.background_video_dir, clip, frame)) as bgr:
bgr = self._downsample_if_needed(bgr.convert('RGB'))
bgrs.append(bgr)
return bgrs
def _get_videomatte(self, idx):
clip_idx, frame_idx = self.videomatte_idx[idx]
clip = self.videomatte_clips[clip_idx]
frame_count = len(self.videomatte_frames[clip_idx])
fgrs, phas = [], []
for i in self.seq_sampler(self.seq_length):
frame = self.videomatte_frames[clip_idx][(frame_idx + i) % frame_count]
with Image.open(os.path.join(self.videomatte_dir, 'fgr', clip, frame)) as fgr, \
Image.open(os.path.join(self.videomatte_dir, 'pha', clip, frame)) as pha:
fgr = self._downsample_if_needed(fgr.convert('RGB'))
pha = self._downsample_if_needed(pha.convert('L'))
fgrs.append(fgr)
phas.append(pha)
return fgrs, phas
def _downsample_if_needed(self, img):
w, h = img.size
if min(w, h) > self.size:
scale = self.size / min(w, h)
w = int(scale * w)
h = int(scale * h)
img = img.resize((w, h))
return img
class VideoMatteTrainAugmentation(MotionAugmentation):
def __init__(self, size):
super().__init__(
size=size,
prob_fgr_affine=0.3,
prob_bgr_affine=0.3,
prob_noise=0.1,
prob_color_jitter=0.3,
prob_grayscale=0.02,
prob_sharpness=0.1,
prob_blur=0.02,
prob_hflip=0.5,
prob_pause=0.03,
)
class VideoMatteValidAugmentation(MotionAugmentation):
def __init__(self, size):
super().__init__(
size=size,
prob_fgr_affine=0,
prob_bgr_affine=0,
prob_noise=0,
prob_color_jitter=0,
prob_grayscale=0,
prob_sharpness=0,
prob_blur=0,
prob_hflip=0,
prob_pause=0,
)
|