File size: 4,624 Bytes
8d015d4 |
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 |
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)
|