| import cv2 | |
| import numpy as np | |
| import random | |
| perm = [(0, 1, 2), (0, 2, 1), (1, 0, 2), (1, 2, 0), (2, 0, 1), (2, 1, 0)] | |
| rotate = [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_180, cv2.ROTATE_90_COUNTERCLOCKWISE] | |
| def random_crop(img0, imgt, img1, crop_size, flowt0=None, flowt1=None, distance=None): | |
| im_h, im_w = img0.shape[:2] | |
| crop_h, crop_w = crop_size, crop_size | |
| i = random.randint(0, im_h - crop_h) | |
| j = random.randint(0, im_w - crop_w) | |
| img0 = img0[i:i + crop_h, j:j + crop_w] | |
| imgt = imgt[i:i + crop_h, j:j + crop_w] | |
| img1 = img1[i:i + crop_h, j:j + crop_w] | |
| if flowt0 is not None and flowt1 is not None: | |
| flowt0 = flowt0[i:i + crop_h, j:j + crop_w] | |
| flowt1 = flowt1[i:i + crop_h, j:j + crop_w] | |
| if distance is not None: | |
| distance = distance[i:i + crop_h, j:j + crop_w] | |
| return img0, imgt, img1, flowt0, flowt1, distance | |
| def random_hor_flip(img0, imgt, img1, flowt0=None, flowt1=None, distance=None): | |
| img0, imgt, img1 = img0[::-1, :, :], imgt[::-1, :, :], img1[::-1, :, :] | |
| if flowt0 is not None and flowt1 is not None: | |
| flowt0, flowt1 = flowt0[::-1, :, :] * np.array([1, -1]).reshape(1, 1, 2), \ | |
| flowt1[::-1, :, :] * np.array([1, -1]).reshape(1, 1, 2) | |
| if distance is not None: | |
| distance = distance[::-1, :, :] | |
| return img0, imgt, img1, flowt0, flowt1, distance | |
| def random_ver_flip(img0, imgt, img1, flowt0=None, flowt1=None, distance=None): | |
| img0, imgt, img1 = img0[:, ::-1, :], imgt[:, ::-1, :], img1[:, ::-1, :] | |
| if flowt0 is not None and flowt1 is not None: | |
| flowt0, flowt1 = flowt0[:, ::-1, :] * np.array([-1, 1]).reshape(1, 1, 2), \ | |
| flowt1[:, ::-1, :] * np.array([-1, 1]).reshape(1, 1, 2) | |
| if distance is not None: | |
| distance = distance[:, ::-1, :] | |
| return img0, imgt, img1, flowt0, flowt1, distance | |
| def random_color_permutation(img0, imgt, img1): | |
| perm_idx = random.randint(0, 5) | |
| img0, imgt, img1 = img0[:, :, perm[perm_idx]], imgt[:, :, perm[perm_idx]], img1[:, :, perm[perm_idx]] | |
| return img0, imgt, img1 | |
| def random_temporal_flip(img0, imgt, img1, time_step, flowt0=None, flowt1=None): | |
| tmp = img1 | |
| img1 = img0 | |
| img0 = tmp | |
| time_step = 1 - time_step | |
| if flowt0 is not None and flowt1 is not None: | |
| tmp = flowt0 | |
| flowt0 = flowt1 | |
| flowt1 = tmp | |
| return img0, imgt, img1, time_step, flowt0, flowt1 | |
| def random_rotation(img0, imgt, img1, degree, flowt0=None, flowt1=None, distance=None): | |
| if degree != 3: | |
| img0 = cv2.rotate(img0, rotate[degree]) | |
| imgt = cv2.rotate(imgt, rotate[degree]) | |
| img1 = cv2.rotate(img1, rotate[degree]) | |
| if flowt0 is not None and flowt1 is not None: | |
| flowt0 = cv2.rotate(flowt0, rotate[degree]) | |
| flowt1 = cv2.rotate(flowt1, rotate[degree]) | |
| if degree == 0: | |
| flowt0 = np.concatenate((-flowt0[:, :, 1:2], flowt0[:, :, 0:1]), 2) | |
| flowt1 = np.concatenate((-flowt1[:, :, 1:2], flowt1[:, :, 0:1]), 2) | |
| elif degree == 1: | |
| flowt0 = -flowt0 | |
| flowt1 = -flowt1 | |
| elif degree == 2: | |
| flowt0 = np.concatenate((flowt0[:, :, 1:2], -flowt0[:, :, 0:1]), 2) | |
| flowt1 = np.concatenate((flowt1[:, :, 1:2], -flowt1[:, :, 0:1]), 2) | |
| if distance is not None: | |
| H,W,_ = distance.shape | |
| distance = cv2.rotate(distance, rotate[degree]).reshape(H,W,1) | |
| return img0, imgt, img1, flowt0, flowt1, distance | |
| def random_resize(img0, imgt, img1, flowt0=None, flowt1=None, distance=None): | |
| ''' | |
| 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) | |
| if flowt0 is not None and flowt1 is not None: | |
| flowt0 = cv2.resize(flowt0, dsize=None, fx=2.0, fy=2.0, interpolation=cv2.INTER_LINEAR) * 2.0 | |
| flowt1 = cv2.resize(flowt1, dsize=None, fx=2.0, fy=2.0, interpolation=cv2.INTER_LINEAR) * 2.0 | |
| ''' | |
| return img0, imgt, img1, flowt0, flowt1, distance | |
| def read_flow(name): | |
| with open(name, "rb") as f: | |
| header = f.read(4) | |
| if header.decode("utf-8") != 'PIEH': | |
| raise Exception('Flow file header does not contain PIEH') | |
| width = np.fromfile(f, np.int32, 1).squeeze() | |
| height = np.fromfile(f, np.int32, 1).squeeze() | |
| flow = np.fromfile(f, np.float32, width * height * 2).reshape((height, width, 2)) | |
| return flow.astype(np.float32) | |