Spaces:
Running
on
T4
Running
on
T4
| import random | |
| import pickle | |
| import logging | |
| import torch | |
| import cv2 | |
| import os | |
| from torch.utils.data.dataset import Dataset | |
| import numpy as np | |
| from skimage.feature import canny | |
| from .util.STTN_mask import create_random_shape_with_random_motion | |
| from cvbase import read_flow, flow2rgb | |
| from .util.flow_utils import region_fill as rf | |
| import imageio | |
| logger = logging.getLogger('base') | |
| class VideoBasedDataset(Dataset): | |
| def __init__(self, opt, dataInfo): | |
| self.opt = opt | |
| self.mode = opt['mode'] | |
| self.sampleMethod = opt['sample'] | |
| self.dataInfo = dataInfo | |
| self.flow_height, self.flow_width = dataInfo['flow']['flow_height'], dataInfo['flow']['flow_width'] | |
| self.data_path = dataInfo['flow_path'] | |
| self.frame_path = dataInfo['frame_path'] | |
| self.train_list = os.listdir(self.data_path) | |
| self.name2length = self.dataInfo['name2len'] | |
| self.require_edge = opt['use_edges'] | |
| self.sigma = dataInfo['edge']['sigma'] | |
| self.low_threshold = dataInfo['edge']['low_threshold'] | |
| self.high_threshold = dataInfo['edge']['high_threshold'] | |
| with open(self.name2length, 'rb') as f: | |
| self.name2len = pickle.load(f) | |
| self.norm = opt['norm'] | |
| self.sequenceLen = self.opt['num_flows'] | |
| self.flow_interval = self.opt['flow_interval'] | |
| self.halfLen = self.sequenceLen // 2 | |
| def __len__(self): | |
| return len(self.train_list) | |
| def __getitem__(self, idx): | |
| try: | |
| item = self.load_item(idx) | |
| except: | |
| print('Loading error: ' + self.train_list[idx]) | |
| item = self.load_item(0) | |
| return item | |
| def frameSample(self, flowLen): | |
| if self.sampleMethod == 'random': | |
| indices = [i for i in range(flowLen)] | |
| sampledIndices = random.sample(indices, self.sequenceLen) | |
| else: | |
| sampledIndices = [] | |
| pivot = random.randint(0, flowLen - 1) | |
| for i in range(-self.halfLen, self.halfLen + 1): | |
| index = pivot + i * self.flow_interval | |
| if index < 0: | |
| index = 0 | |
| if index >= flowLen: | |
| index = flowLen - 1 | |
| sampledIndices.append(index) | |
| return sampledIndices | |
| def load_item(self, idx): | |
| info = {} | |
| video = self.train_list[idx] | |
| info['name'] = video | |
| if np.random.uniform(0, 1) > 0.5: | |
| direction = 'forward_flo' | |
| else: | |
| direction = 'backward_flo' | |
| flow_dir = os.path.join(self.data_path, video, direction) | |
| frame_dir = os.path.join(self.frame_path, video) | |
| flowLen = self.name2len[video] - 1 | |
| assert flowLen > self.sequenceLen, 'Flow length {} is not enough'.format(flowLen) | |
| sampledIndices = self.frameSample(flowLen) | |
| candidateMasks = create_random_shape_with_random_motion(self.sequenceLen, 0.9, 1.1, 1, | |
| 10) | |
| flows, diffused_flows, masks = [], [], [] | |
| current_frames, shift_frames = None, None | |
| mask_counter = 0 | |
| for i in sampledIndices: | |
| flow = read_flow(os.path.join(flow_dir, '{:05d}.flo'.format(i))) | |
| mask = self.read_mask(candidateMasks[mask_counter], self.flow_height, self.flow_width) | |
| mask_counter += 1 | |
| flow = self.flow_tf(flow, self.flow_height, self.flow_width) | |
| diffused_flow = self.diffusion_fill(flow, mask) | |
| flows.append(flow) | |
| masks.append(mask) | |
| diffused_flows.append(diffused_flow) | |
| targetIndex = sampledIndices[self.sequenceLen // 2] | |
| current_frames, shift_frames = self.read_frames(frame_dir, targetIndex, direction, self.flow_width, | |
| self.flow_height) | |
| flow_gray, edge = self.load_edge(flows[self.halfLen]) | |
| inputs = {'flows': flows, 'diffused_flows': diffused_flows, 'current_frame': current_frames, | |
| 'shift_frame': shift_frames, 'edges': edge, 'masks': masks, 'flow_gray': flow_gray} | |
| return self.to_tensor(inputs) | |
| def read_frames(self, frame_dir, index, direction, width, height): | |
| if direction == 'forward_flo': | |
| current_frame = os.path.join(frame_dir, '{:05d}.jpg'.format(index)) | |
| shift_frame = os.path.join(frame_dir, '{:05d}.jpg'.format(index + 1)) | |
| else: | |
| current_frame = os.path.join(frame_dir, '{:05d}.jpg'.format(index + 1)) | |
| shift_frame = os.path.join(frame_dir, '{:05d}.jpg'.format(index)) | |
| current_frame = imageio.imread(current_frame) | |
| shift_frame = imageio.imread(shift_frame) | |
| current_frame = cv2.resize(current_frame, (width, height), cv2.INTER_LINEAR) | |
| shift_frame = cv2.resize(shift_frame, (width, height), cv2.INTER_LINEAR) | |
| current_frame = current_frame / 255. | |
| shift_frame = shift_frame / 255. | |
| return current_frame, shift_frame | |
| def diffusion_fill(self, flow, mask): | |
| flow_filled = np.zeros(flow.shape) | |
| flow_filled[:, :, 0] = rf.regionfill(flow[:, :, 0] * (1 - mask), mask) | |
| flow_filled[:, :, 1] = rf.regionfill(flow[:, :, 1] * (1 - mask), mask) | |
| return flow_filled | |
| def flow_tf(self, flow, height, width): | |
| flow_shape = flow.shape | |
| flow_resized = cv2.resize(flow, (width, height), cv2.INTER_LINEAR) | |
| flow_resized[:, :, 0] *= (float(width) / float(flow_shape[1])) | |
| flow_resized[:, :, 1] *= (float(height) / float(flow_shape[0])) | |
| return flow_resized | |
| def read_mask(self, mask, height, width): | |
| mask = np.array(mask) | |
| mask = mask / 255. | |
| raw_mask = (mask > 0.5).astype(np.uint8) | |
| raw_mask = cv2.resize(raw_mask, dsize=(width, height), interpolation=cv2.INTER_NEAREST) | |
| return raw_mask | |
| def load_edge(self, flow): | |
| gray_flow = (flow[:, :, 0] ** 2 + flow[:, :, 1] ** 2) ** 0.5 | |
| factor = gray_flow.max() | |
| gray_flow = gray_flow / factor | |
| flow_rgb = flow2rgb(flow) | |
| flow_gray = cv2.cvtColor(flow_rgb, cv2.COLOR_RGB2GRAY) | |
| return gray_flow, canny(flow_gray, sigma=self.sigma, mask=None, low_threshold=self.low_threshold, | |
| high_threshold=self.high_threshold).astype(np.float) | |
| def to_tensor(self, data_list): | |
| """ | |
| Args: | |
| data_list: a numpy.array list | |
| Returns: a torch.tensor list with the None entries removed | |
| """ | |
| keys = list(data_list.keys()) | |
| for key in keys: | |
| if data_list[key] is None or data_list[key] == []: | |
| data_list.pop(key) | |
| else: | |
| item = data_list[key] | |
| if not isinstance(item, list): | |
| if len(item.shape) == 2: | |
| item = item[:, :, np.newaxis] | |
| item = torch.from_numpy(np.transpose(item, (2, 0, 1))).float() | |
| else: | |
| item = np.stack(item, axis=0) | |
| if len(item.shape) == 3: | |
| item = item[:, :, :, np.newaxis] | |
| item = torch.from_numpy(np.transpose(item, (3, 0, 1, 2))).float() | |
| data_list[key] = item | |
| return data_list | |