def test_shift_y_from_shift_scale_rotate(target):
img = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]], dtype=np.uint8)
expected = np.array([[0, 0, 0, 0], [0, 0, 0, 0], [1, 2, 3, 4], [5, 6, 7, 8]], dtype=np.uint8)
img, expected = convert_2d_to_target_format([img, expected], target=target)
shifted_along_y_img = F.shift_scale_rotate(
img, angle=0, scale=1, dx=0, dy=0.5, interpolation=cv2.INTER_NEAREST, border_mode=cv2.BORDER_CONSTANT
)
assert np.array_equal(shifted_along_y_img, expected)
def apply(self,
img,
angle=0,
scale=0,
dx=0,
dy=0,
interpolation=cv2.INTER_LINEAR,
**params):
return augf.shift_scale_rotate(img, angle, scale, dx, dy,
interpolation, self.border_mode,
self.value)
def aug3D(self, vols):
num_rot_xy = self.rng.randint(4)
# num_rot_yz = self.rng.randint (4)
num_rot_yz = self.rng.choice([1, 3], 1)[0]
useFlipX = self.rng.randint(2)
useFlipY = self.rng.randint(2)
useFlipZ = self.rng.randint(2)
ret = []
for vol in vols:
vol = np.rot90(vol, num_rot_xy, axes=(1, 2))
vol = np.rot90(vol, num_rot_yz, axes=(0, 1))
if useFlipZ:
vol = np.flip(vol, axis=0)
if useFlipX:
vol = np.flip(vol, axis=1)
if useFlipY:
vol = np.flip(vol, axis=2)
ret.append(vol)
# Full AUG
if not self.alpha_only:
if self.raw.shape[0] < 50:
angle = self.rng.randint(10)
scale = self.rng.uniform(0.9, 1.1)
dx = self.rng.randint(-2, 2)
dy = self.rng.randint(-2, 2)
else:
angle = self.rng.randint(30)
scale = self.rng.uniform(0.8, 1.2)
dx = self.rng.randint(-4, 4)
dy = self.rng.randint(-4, 4)
for vol in ret:
for i, img in enumerate(vol):
vol[i] = F.shift_scale_rotate(
img,
angle,
scale,
dx,
dy,
interpolation=cv2.INTER_LINEAR,
border_mode=cv2.BORDER_REFLECT_101)
return ret
def test_shift_scale_separate_shift_x_shift_y(image, mask):
aug = A.ShiftScaleRotate(shift_limit=(0.3, 0.3),
shift_limit_y=(0.4, 0.4),
scale_limit=0,
rotate_limit=0,
p=1)
data = aug(image=image, mask=mask)
expected_image = F.shift_scale_rotate(image,
angle=0,
scale=1,
dx=0.3,
dy=0.4,
interpolation=cv2.INTER_LINEAR,
border_mode=cv2.BORDER_REFLECT_101)
expected_mask = F.shift_scale_rotate(mask,
angle=0,
scale=1,
dx=0.3,
dy=0.4,
interpolation=cv2.INTER_NEAREST,
border_mode=cv2.BORDER_REFLECT_101)
assert np.array_equal(data["image"], expected_image)
assert np.array_equal(data["mask"], expected_mask)
def ghosting(img, dis, alpha=0.5, angle=0, scale=0):
dx = random.uniform(-dis, dis)
dy = dis - abs(dx)
dy = random.choice([-dy, dy])
tmp1 = F.shift_scale_rotate(img, angle, scale + 1, dx, dy, value=0)
if random.random() < 0.5:
tmp1 = F.gaussian_blur(tmp1, 3)
tmp2 = img
if random.random() < 0.5:
tmp2 = F.gaussian_blur(tmp2, 3)
beta = 1 - alpha
gamma = 0
return cv2.addWeighted(tmp1, alpha, tmp2, beta, gamma)
def test_shift_y_float_from_shift_scale_rotate(target):
img = np.array(
[[0.01, 0.02, 0.03, 0.04], [0.05, 0.06, 0.07, 0.08], [0.09, 0.10, 0.11, 0.12], [0.13, 0.14, 0.15, 0.16]],
dtype=np.float32,
)
expected = np.array(
[[0.00, 0.00, 0.00, 0.00], [0.00, 0.00, 0.00, 0.00], [0.01, 0.02, 0.03, 0.04], [0.05, 0.06, 0.07, 0.08]],
dtype=np.float32,
)
img, expected = convert_2d_to_target_format([img, expected], target=target)
shifted_along_y_img = F.shift_scale_rotate(
img, angle=0, scale=1, dx=0, dy=0.5, interpolation=cv2.INTER_NEAREST, border_mode=cv2.BORDER_CONSTANT
)
assert_array_almost_equal_nulp(shifted_along_y_img, expected)
def test_shift_y_from_shift_scale_rotate(target):
img = np.array([
[1, 2, 3, 4],
[5, 6, 7, 8],
[9, 10, 11, 12],
[13, 14, 15, 16]], dtype=np.uint8)
expected = np.array([
[0, 0, 0, 0],
[0, 0, 0, 0],
[1, 2, 3, 4],
[5, 6, 7, 8]], dtype=np.uint8)
if target == 'image':
img = convert_2d_to_3d(img)
expected = convert_2d_to_3d(expected)
shifted_along_y_img = F.shift_scale_rotate(img, angle=0, scale=1, dx=0, dy=0.5, interpolation=cv2.INTER_NEAREST,
border_mode=cv2.BORDER_CONSTANT)
assert np.array_equal(shifted_along_y_img, expected)
def test_shift_x_from_shift_scale_rotate(target):
img = np.array(
[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]],
dtype=np.uint8)
expected = np.array(
[[0, 0, 1, 2], [0, 0, 5, 6], [0, 0, 9, 10], [0, 0, 13, 14]],
dtype=np.uint8)
if target == 'image':
img = convert_2d_to_3d(img)
expected = convert_2d_to_3d(expected)
shifted_along_x_img = F.shift_scale_rotate(img,
angle=0,
scale=1,
dx=0.5,
dy=0,
interpolation=cv2.INTER_NEAREST,
border_mode=cv2.BORDER_CONSTANT)
assert np.array_equal(shifted_along_x_img, expected)
def test_compare_rotate_and_shift_scale_rotate(image):
rotated_img_1 = F.rotate(image, angle=60)
rotated_img_2 = F.shift_scale_rotate(image, angle=60, scale=1, dx=0, dy=0)
assert np.array_equal(rotated_img_1, rotated_img_2)
def albumentations(self, img):
return albumentations.shift_scale_rotate(img, angle=-45, scale=2, dx=0.2, dy=0.2)
def apply_to_mask(self, img, angle=0, scale=0, dx=0, dy=0, **params):
return augf.shift_scale_rotate(img, angle, scale, dx, dy,
cv2.INTER_NEAREST, self.border_mode,
self.mask_value)
def apply_to_mask(self, img, dx=0, **params):
return F.shift_scale_rotate(img, 0, 1, dx, 0, cv2.INTER_NEAREST, self.border_mode, self.mask_value)
def __getitem__(self, index):
if index < len(self.img_paths):
img_path, mask_path, label = self.img_paths[
index], self.mask_paths[index], self.labels[index]
num_objs = len(label)
img = cv2.imread(img_path)
height, width = img.shape[:2]
img = cv2.resize(img, (self.input_w, self.input_h))
mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
if mask is not None:
mask = cv2.resize(mask, (self.output_w, self.output_h))
mask = 1 - mask.astype('float32') / 255
else:
mask = np.ones((self.output_h, self.output_w), dtype='float32')
if self.test:
img = img.astype('float32') / 255
img = (img - self.mean) / self.std
img = img.transpose(2, 0, 1)
mask = mask[None, ...]
if self.lhalf:
img = img[:, self.input_h // 2:]
mask = mask[:, self.output_h // 2:]
return {
'img_path': img_path,
'input': img,
'mask': mask,
}
kpts = []
poses = []
for k in range(num_objs):
ann = label[k]
kpts.append([ann['x'], ann['y'], ann['z']])
poses.append([ann['yaw'], ann['pitch'], ann['roll']])
kpts = np.array(kpts)
poses = np.array(poses)
if np.random.random() < self.hflip:
img = img[:, ::-1].copy()
mask = mask[:, ::-1].copy()
kpts[:, 0] *= -1
poses[:, [0, 2]] *= -1
if np.random.random() < self.scale:
scale = np.random.uniform(-self.scale_limit,
self.scale_limit) + 1.0
img = F.shift_scale_rotate(img,
angle=0,
scale=scale,
dx=0,
dy=0)
mask = F.shift_scale_rotate(mask,
angle=0,
scale=scale,
dx=0,
dy=0)
kpts[:, 2] /= scale
kpts = np.array(
convert_3d_to_2d(kpts[:, 0], kpts[:, 1], kpts[:, 2])).T
kpts[:, 0] *= self.input_w / width
kpts[:, 1] *= self.input_h / height
if self.transform is not None:
data = self.transform(image=img, mask=mask, keypoints=kpts)
img = data['image']
mask = data['mask']
kpts = data['keypoints']
for k, ((x, y), (yaw, pitch, roll)) in enumerate(zip(kpts, poses)):
label[k]['x'] = x
label[k]['y'] = y
label[k]['yaw'] = yaw
label[k]['pitch'] = pitch
label[k]['roll'] = roll
img = img.astype('float32') / 255
img = (img - self.mean) / self.std
img = img.transpose(2, 0, 1)
mask = mask[None, ...]
hm = np.zeros((1, self.output_h, self.output_w), dtype=np.float32)
reg_mask = np.zeros((1, self.output_h, self.output_w),
dtype=np.float32)
reg = np.zeros((2, self.output_h, self.output_w), dtype=np.float32)
wh = np.zeros((2, self.output_h, self.output_w), dtype=np.float32)
depth = np.zeros((1, self.output_h, self.output_w),
dtype=np.float32)
eular = np.zeros((3, self.output_h, self.output_w),
dtype=np.float32)
trig = np.zeros((6, self.output_h, self.output_w),
dtype=np.float32)
quat = np.zeros((4, self.output_h, self.output_w),
dtype=np.float32)
gt = np.zeros((self.max_objs, 7), dtype=np.float32)
for k in range(num_objs):
ann = label[k]
x, y = ann['x'], ann['y']
x *= self.output_w / self.input_w
y *= self.output_h / self.input_h
if x < 0 or y < 0 or x > self.output_w or y > self.output_h:
continue
bbox = get_bbox(
ann['yaw'], ann['pitch'], ann['roll'],
*convert_2d_to_3d(ann['x'] * width / self.input_w,
ann['y'] * height / self.input_h,
ann['z']), ann['z'], width, height,
self.output_w, self.output_h)
h, w = bbox[3] - bbox[1], bbox[2] - bbox[0]
radius = gaussian_radius((math.ceil(h), math.ceil(w)))
radius = max(0, int(radius))
ct = np.array([x, y], dtype=np.float32)
ct_int = ct.astype(np.int32)
draw_umich_gaussian(hm[0], ct_int, radius)
reg_mask[0, ct_int[1], ct_int[0]] = 1
reg[:, ct_int[1], ct_int[0]] = ct - ct_int
wh[0, ct_int[1], ct_int[0]] = w
wh[1, ct_int[1], ct_int[0]] = h
depth[0, ct_int[1], ct_int[0]] = ann['z']
yaw = ann['yaw']
pitch = ann['pitch']
roll = ann['roll']
eular[0, ct_int[1], ct_int[0]] = yaw
eular[1, ct_int[1], ct_int[0]] = pitch
eular[2, ct_int[1], ct_int[0]] = rotate(roll, np.pi)
trig[0, ct_int[1], ct_int[0]] = math.cos(yaw)
trig[1, ct_int[1], ct_int[0]] = math.sin(yaw)
trig[2, ct_int[1], ct_int[0]] = math.cos(pitch)
trig[3, ct_int[1], ct_int[0]] = math.sin(pitch)
trig[4, ct_int[1], ct_int[0]] = math.cos(rotate(roll, np.pi))
trig[5, ct_int[1], ct_int[0]] = math.sin(rotate(roll, np.pi))
qx, qy, qz, qw = (R.from_euler('xyz',
[yaw, pitch, roll])).as_quat()
norm = (qx**2 + qy**2 + qz**2 + qw**2)**(1 / 2)
quat[0, ct_int[1], ct_int[0]] = qx / norm
quat[1, ct_int[1], ct_int[0]] = qy / norm
quat[2, ct_int[1], ct_int[0]] = qz / norm
quat[3, ct_int[1], ct_int[0]] = qw / norm
gt[k, 0] = ann['pitch']
gt[k, 1] = ann['yaw']
gt[k, 2] = ann['roll']
gt[k, 3:5] = convert_2d_to_3d(ann['x'] * width / self.input_w,
ann['y'] * height / self.input_h,
ann['z'])
gt[k, 5] = ann['z']
gt[k, 6] = 1
if self.lhalf:
img = img[:, self.input_h // 2:]
mask = mask[:, self.output_h // 2:]
hm = hm[:, self.output_h // 2:]
reg_mask = reg_mask[:, self.output_h // 2:]
reg = reg[:, self.output_h // 2:]
wh = wh[:, self.output_h // 2:]
depth = depth[:, self.output_h // 2:]
eular = eular[:, self.output_h // 2:]
trig = trig[:, self.output_h // 2:]
quat = quat[:, self.output_h // 2:]
else:
index -= len(self.img_paths)
img_path, mask_path = self.test_img_paths[
index], self.test_mask_paths[index]
output = self.test_outputs[os.path.splitext(
os.path.basename(img_path))[0]]
img = cv2.imread(img_path)
height, width = img.shape[:2]
img = cv2.resize(img, (self.input_w, self.input_h))
mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
if mask is not None:
mask = cv2.resize(mask, (self.output_w, self.output_h))
mask = 1 - mask.astype('float32') / 255
else:
mask = np.ones((self.output_h, self.output_w), dtype='float32')
if self.test:
img = img.astype('float32') / 255
img = (img - self.mean) / self.std
img = img.transpose(2, 0, 1)
mask = mask[None, ...]
if self.lhalf:
img = img[:, self.input_h // 2:]
mask = mask[:, self.output_h // 2:]
return {
'img_path': img_path,
'input': img,
'mask': mask,
}
img = img.astype('float32') / 255
img = (img - self.mean) / self.std
img = img.transpose(2, 0, 1)
mask = mask[None, ...]
hm = np.zeros((1, self.output_h // 2, self.output_w),
dtype=np.float32)
reg_mask = np.zeros((1, self.output_h // 2, self.output_w),
dtype=np.float32)
reg = np.zeros((2, self.output_h // 2, self.output_w),
dtype=np.float32)
wh = np.zeros((2, self.output_h // 2, self.output_w),
dtype=np.float32)
depth = np.zeros((1, self.output_h // 2, self.output_w),
dtype=np.float32)
eular = np.zeros((3, self.output_h // 2, self.output_w),
dtype=np.float32)
trig = np.zeros((6, self.output_h // 2, self.output_w),
dtype=np.float32)
quat = np.zeros((4, self.output_h // 2, self.output_w),
dtype=np.float32)
gt = np.zeros((self.max_objs, 7), dtype=np.float32)
hm = torch.sigmoid(output['hm']).numpy()[0]
reg_mask = hm
reg = output['reg'].numpy()[0]
wh = output['wh'].numpy()[0]
depth = output['depth'].numpy()[0]
eular = eular if output['eular'] is None else output[
'eular'].numpy()[0]
trig = trig if output['trig'] is None else output['trig'].numpy(
)[0]
quat = quat if output['quat'] is None else output['quat'].numpy(
)[0]
if self.lhalf:
img = img[:, self.input_h // 2:]
mask = mask[:, self.output_h // 2:]
ret = {
'img_path': img_path,
'input': img,
'mask': mask,
# 'label': label,
'hm': hm,
'reg_mask': reg_mask,
'reg': reg,
'wh': wh,
'depth': depth,
'eular': eular,
'trig': trig,
'quat': quat,
'gt': gt,
}
# plt.imshow(ret['hm'][0])
# plt.show()
# img = visualize(((img.transpose(1, 2, 0) * self.std + self.mean) * 255).astype('uint8'),
# gt[gt[:, -1] > 0], scale_w=self.input_w / width, scale_h=self.input_h / height)
# plt.imshow(img)
# plt.show()
return ret
def func(img):
return albF.shift_scale_rotate(img, degree, 1.0 + scale, shift[0], shift[1], cv2.INTER_LINEAR, cv2.BORDER_CONSTANT)
def scale(img, level):
"Scale image with level. Zoom in/out at random"
v = float_parameter(level, 1.0)
if random.random() < 0.5:
v = -v * 0.4
return Image.fromarray(shift_scale_rotate(np.array(img), 0, v + 1.0, 0, 0))
def test_compare_rotate_and_shift_scale_rotate(image):
rotated_img_1 = F.rotate(image, angle=60)
rotated_img_2 = F.shift_scale_rotate(image, angle=60, scale=1, dx=0, dy=0)
assert np.array_equal(rotated_img_1, rotated_img_2)
def albumentations(self, img):
return albumentations.shift_scale_rotate(img,
angle=-45,
scale=2,
dx=0.2,
dy=0.2)
def apply(self, img, dx=0, interpolation=cv2.INTER_LINEAR, **params):
return F.shift_scale_rotate(img, 0, 1, dx, 0, interpolation, self.border_mode, self.value)
