Spaces:
Runtime error
Runtime error
File size: 8,472 Bytes
04fbff5 |
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 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 |
'''
This code is partially borrowed from IFRNet (https://github.com/ltkong218/IFRNet).
In the consideration of the difficulty in flow supervision generation, we abort
flow loss in the 8x case.
'''
import os
import cv2
import torch
import random
import numpy as np
from torch.utils.data import Dataset
from utils.utils import read, img2tensor
def random_resize_woflow(img0, imgt, img1, p=0.1):
if random.uniform(0, 1) < p:
img0 = cv2.resize(img0, dsize=None, fx=2.0, fy=2.0, interpolation=cv2.INTER_LINEAR)
imgt = cv2.resize(imgt, dsize=None, fx=2.0, fy=2.0, interpolation=cv2.INTER_LINEAR)
img1 = cv2.resize(img1, dsize=None, fx=2.0, fy=2.0, interpolation=cv2.INTER_LINEAR)
return img0, imgt, img1
def random_crop_woflow(img0, imgt, img1, crop_size=(224, 224)):
h, w = crop_size[0], crop_size[1]
ih, iw, _ = img0.shape
x = np.random.randint(0, ih-h+1)
y = np.random.randint(0, iw-w+1)
img0 = img0[x: x + h, y : y + w, :]
imgt = imgt[x: x + h, y : y + w, :]
img1 = img1[x: x + h, y : y + w, :]
return img0, imgt, img1
def center_crop_woflow(img0, imgt, img1, crop_size=(512, 512)):
h, w = crop_size[0], crop_size[1]
ih, iw, _ = img0.shape
img0 = img0[ih // 2 - h // 2: ih // 2 + h // 2, iw // 2 - w // 2: iw // 2 + w // 2, :]
imgt = imgt[ih // 2 - h // 2: ih // 2 + h // 2, iw // 2 - w // 2: iw // 2 + w // 2, :]
img1 = img1[ih // 2 - h // 2: ih // 2 + h // 2, iw // 2 - w // 2: iw // 2 + w // 2, :]
return img0, imgt, img1
def random_reverse_channel_woflow(img0, imgt, img1, p=0.5):
if random.uniform(0, 1) < p:
img0 = img0[:, :, ::-1]
imgt = imgt[:, :, ::-1]
img1 = img1[:, :, ::-1]
return img0, imgt, img1
def random_vertical_flip_woflow(img0, imgt, img1, p=0.3):
if random.uniform(0, 1) < p:
img0 = img0[::-1]
imgt = imgt[::-1]
img1 = img1[::-1]
return img0, imgt, img1
def random_horizontal_flip_woflow(img0, imgt, img1, p=0.5):
if random.uniform(0, 1) < p:
img0 = img0[:, ::-1]
imgt = imgt[:, ::-1]
img1 = img1[:, ::-1]
return img0, imgt, img1
def random_rotate_woflow(img0, imgt, img1, p=0.05):
if random.uniform(0, 1) < p:
img0 = img0.transpose((1, 0, 2))
imgt = imgt.transpose((1, 0, 2))
img1 = img1.transpose((1, 0, 2))
return img0, imgt, img1
def random_reverse_time_woflow(img0, imgt, img1, embt, p=0.5):
if random.uniform(0, 1) < p:
tmp = img1
img1 = img0
img0 = tmp
embt = 1 - embt
return img0, imgt, img1, embt
class GoPro_Train_Dataset(Dataset):
def __init__(self, dataset_dir='data/GOPRO', interFrames=7, augment=True):
self.dataset_dir = dataset_dir + '/train'
self.interFrames = interFrames
self.augment = augment
self.setLength = interFrames + 2
video_list = [
'GOPR0372_07_00', 'GOPR0374_11_01', 'GOPR0378_13_00', 'GOPR0384_11_01',
'GOPR0384_11_04', 'GOPR0477_11_00', 'GOPR0868_11_02', 'GOPR0884_11_00',
'GOPR0372_07_01', 'GOPR0374_11_02', 'GOPR0379_11_00', 'GOPR0384_11_02',
'GOPR0385_11_00', 'GOPR0857_11_00', 'GOPR0871_11_01', 'GOPR0374_11_00',
'GOPR0374_11_03', 'GOPR0380_11_00', 'GOPR0384_11_03', 'GOPR0386_11_00',
'GOPR0868_11_01', 'GOPR0881_11_00']
self.frames_list = []
self.file_list = []
for video in video_list:
frames = sorted(os.listdir(os.path.join(self.dataset_dir, video)))
n_sets = (len(frames) - self.setLength) // (interFrames+1) + 1
videoInputs = [frames[(interFrames + 1) * i: (interFrames + 1) * i + self.setLength
] for i in range(n_sets)]
videoInputs = [[os.path.join(video, f) for f in group] for group in videoInputs]
self.file_list.extend(videoInputs)
def __len__(self):
return len(self.file_list) * self.interFrames
def __getitem__(self, idx):
clip_idx = idx // self.interFrames
embt_idx = idx % self.interFrames
imgpaths = [os.path.join(self.dataset_dir, fp) for fp in self.file_list[clip_idx]]
pick_idxs = list(range(0, self.setLength, self.interFrames + 1))
imgt_beg = self.setLength // 2 - self.interFrames // 2
imgt_end = self.setLength // 2 + self.interFrames // 2 + self.interFrames % 2
imgt_idx = list(range(imgt_beg, imgt_end))
input_paths = [imgpaths[idx] for idx in pick_idxs]
imgt_paths = [imgpaths[idx] for idx in imgt_idx]
embt = torch.from_numpy(np.array((embt_idx + 1) / (self.interFrames+1)
).reshape(1, 1, 1).astype(np.float32))
img0 = np.array(read(input_paths[0]))
imgt = np.array(read(imgt_paths[embt_idx]))
img1 = np.array(read(input_paths[1]))
if self.augment == True:
img0, imgt, img1 = random_resize_woflow(img0, imgt, img1, p=0.1)
img0, imgt, img1 = random_crop_woflow(img0, imgt, img1, crop_size=(224, 224))
img0, imgt, img1 = random_reverse_channel_woflow(img0, imgt, img1, p=0.5)
img0, imgt, img1 = random_vertical_flip_woflow(img0, imgt, img1, p=0.3)
img0, imgt, img1 = random_horizontal_flip_woflow(img0, imgt, img1, p=0.5)
img0, imgt, img1 = random_rotate_woflow(img0, imgt, img1, p=0.05)
img0, imgt, img1, embt = random_reverse_time_woflow(img0, imgt, img1,
embt=embt, p=0.5)
else:
img0, imgt, img1 = center_crop_woflow(img0, imgt, img1, crop_size=(512, 512))
img0 = img2tensor(img0.copy()).squeeze(0)
imgt = img2tensor(imgt.copy()).squeeze(0)
img1 = img2tensor(img1.copy()).squeeze(0)
return {'img0': img0.float(),
'imgt': imgt.float(),
'img1': img1.float(),
'embt': embt}
class GoPro_Test_Dataset(Dataset):
def __init__(self, dataset_dir='data/GOPRO', interFrames=7):
self.dataset_dir = dataset_dir + '/test'
self.interFrames = interFrames
self.setLength = interFrames + 2
video_list = [
'GOPR0384_11_00', 'GOPR0385_11_01', 'GOPR0410_11_00',
'GOPR0862_11_00', 'GOPR0869_11_00', 'GOPR0881_11_01',
'GOPR0384_11_05', 'GOPR0396_11_00', 'GOPR0854_11_00',
'GOPR0868_11_00', 'GOPR0871_11_00']
self.frames_list = []
self.file_list = []
for video in video_list:
frames = sorted(os.listdir(os.path.join(self.dataset_dir, video)))
n_sets = (len(frames) - self.setLength)//(interFrames+1) + 1
videoInputs = [frames[(interFrames + 1) * i:(interFrames + 1) * i + self.setLength
] for i in range(n_sets)]
videoInputs = [[os.path.join(video, f) for f in group] for group in videoInputs]
self.file_list.extend(videoInputs)
def __len__(self):
return len(self.file_list) * self.interFrames
def __getitem__(self, idx):
clip_idx = idx // self.interFrames
embt_idx = idx % self.interFrames
imgpaths = [os.path.join(self.dataset_dir, fp) for fp in self.file_list[clip_idx]]
pick_idxs = list(range(0, self.setLength, self.interFrames + 1))
imgt_beg = self.setLength // 2 - self.interFrames // 2
imgt_end = self.setLength // 2 + self.interFrames // 2 + self.interFrames % 2
imgt_idx = list(range(imgt_beg, imgt_end))
input_paths = [imgpaths[idx] for idx in pick_idxs]
imgt_paths = [imgpaths[idx] for idx in imgt_idx]
img0 = np.array(read(input_paths[0]))
imgt = np.array(read(imgt_paths[embt_idx]))
img1 = np.array(read(input_paths[1]))
img0, imgt, img1 = center_crop_woflow(img0, imgt, img1, crop_size=(512, 512))
img0 = img2tensor(img0).squeeze(0)
imgt = img2tensor(imgt).squeeze(0)
img1 = img2tensor(img1).squeeze(0)
embt = torch.from_numpy(np.array((embt_idx + 1) / (self.interFrames + 1)
).reshape(1, 1, 1).astype(np.float32))
return {'img0': img0.float(),
'imgt': imgt.float(),
'img1': img1.float(),
'embt': embt} |