def test_fix_broadcast():
shape1 = [4, 6, 8, 10]
arr1 = np.zeros(shape1)
shape2 = [4, 6]
arr2 = np.zeros(shape2)
new_arr1, new_arr2 = pt_util.fix_broadcast(arr1, arr2)
assert new_arr1.shape == (4, 6, 8, 10)
assert new_arr2.shape == (4, 6, 1, 1)
new_arr2, new_arr1 = pt_util.fix_broadcast(arr2, arr1)
assert new_arr1.shape == (4, 6, 8, 10)
assert new_arr2.shape == (4, 6, 1, 1)
shape1 = [4, 6, 8, 10]
arr1 = np.zeros(shape1)
shape2 = [8, 10]
arr2 = np.zeros(shape2)
new_arr1, new_arr2 = pt_util.fix_broadcast(arr1, arr2)
assert new_arr1.shape == (4, 6, 8, 10)
assert new_arr2.shape == (1, 1, 8, 10)
shape1 = [4, 6, 8, 10]
arr1 = np.zeros(shape1)
shape2 = [4, 10]
arr2 = np.zeros(shape2)
new_arr1, new_arr2 = pt_util.fix_broadcast(arr1, arr2)
assert new_arr1.shape == (4, 6, 8, 10)
shape1 = [10, 6, 8, 10]
arr1 = np.zeros(shape1)
shape2 = [10]
arr2 = np.zeros(shape2)
new_arr1, new_arr2 = pt_util.fix_broadcast(arr1, arr2)
assert new_arr1.shape == (10, 6, 8, 10)
assert new_arr2.shape == (1, 1, 1, 10)
def draw_attention(image, attention, height_channel=0, width_channel=1):
dtype = image.dtype
image = pt_util.to_numpy(image)
attention = pt_util.to_numpy(attention)
attention = np.clip(attention, 0, 1)
im_width = image.shape[width_channel]
im_height = image.shape[height_channel]
attention = misc_util.resize(
attention,
(im_width, im_height),
interpolation=cv2.INTER_LINEAR,
height_channel=height_channel,
width_channel=width_channel,
)
image, attention = pt_util.fix_broadcast(image, attention)
image = (image * (1 - attention) + 255 * attention).astype(dtype)
return image
def make_mean_and_std(self, inputs):
_, mean = pt_util.fix_broadcast(inputs, self.mean)
_, std = pt_util.fix_broadcast(inputs, self.std)
mean = mean.to(dtype=inputs.dtype, device=inputs.device)
std = std.to(dtype=inputs.dtype, device=inputs.device)
return mean, std