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