def test_get_windows(self): extent = Box(0, 0, 100, 100) windows = list(extent.get_windows(10, 10)) self.assertEqual(len(windows), 100) extent = Box(0, 0, 100, 100) windows = list(extent.get_windows(10, 5)) self.assertEqual(len(windows), 400) extent = Box(0, 0, 20, 20) windows = set( [window.tuple_format() for window in extent.get_windows(10, 10)]) expected_windows = [ Box.make_square(0, 0, 10), Box.make_square(10, 0, 10), Box.make_square(0, 10, 10), Box.make_square(10, 10, 10) ] expected_windows = set( [window.tuple_format() for window in expected_windows]) self.assertSetEqual(windows, expected_windows) extent = Box(10, 10, 20, 20) windows = set( [window.tuple_format() for window in extent.get_windows(6, 6)]) expected_windows = [ Box.make_square(10, 10, 6), Box.make_square(10, 16, 6), Box.make_square(16, 10, 6), Box.make_square(16, 16, 6) ] expected_windows = set( [window.tuple_format() for window in expected_windows]) self.assertSetEqual(windows, expected_windows)
def get_predict_windows(self, extent: Box) -> List[Box]: """Returns windows to compute predictions for. Args: extent: extent of RasterSource """ chip_sz = stride = self.config.predict_chip_sz return extent.get_windows(chip_sz, stride)