def __call__(self, handler): for data in self.iter(handler): for name in self.shape_dict.keys(): assert isinstance(data[name], np.ndarray) t, n, h, w, c = data[name].shape crop_h, crop_w = self.shape_dict[name] max_r = h - crop_h max_c = w - crop_w row_indices = self.rnd.random_integers(0, max_r, n) col_indices = self.rnd.random_integers(0, max_c, n) cropped = np.zeros((t, n, crop_h, crop_w, c)) _crop_images(data[name], crop_h, crop_w, row_indices, col_indices, cropped) data[name] = cropped yield data
def test_crop_images_operation(): a = np.random.randn(3, 2, 5, 5, 4) out = np.zeros((3, 2, 3, 3, 4)) _crop_images(a, 3, 3, np.array([0, 1]), np.array([0, 2]), out) assert np.allclose(out[:, 0, ...], a[:, 0, 0:3, 0:3, :]) assert np.allclose(out[:, 1, ...], a[:, 1, 1:4, 2:5, :])