def setUp(self): self.windows = [Box.make_square(0, 0, 10), Box.make_square(0, 10, 10)] self.label_arr0 = np.random.choice([0, 1], (10, 10)) self.label_arr1 = np.random.choice([0, 1], (10, 10)) self.labels = SemanticSegmentationLabels() self.labels.set_label_arr(self.windows[0], self.label_arr0) self.labels.set_label_arr(self.windows[1], self.label_arr1)
def get_labels(self, window: Optional[Box] = None, chip_size=1000) -> SemanticSegmentationLabels: """Get labels for a window. To avoid running out of memory, if window is None and defaults to using the full extent, a set of sub-windows of size chip_size are used which cover the full extent with no overlap. Args: window: Either None or a window given as a Box object. Uses full extent of scene if window is not provided. chip_size: size of sub-windows to use if full extent is used (in units of pixels) Returns: SemanticSegmentationLabels """ labels = SemanticSegmentationLabels() windows = (self.raster_source.get_extent().get_windows( chip_size, chip_size) if window is None else [window]) for w in windows: raw_labels = self.raster_source.get_raw_chip(w) label_arr = (np.squeeze(raw_labels) if self.class_transformer is None else self.class_transformer.rgb_to_class(raw_labels)) labels.set_label_arr(w, label_arr) return labels
def predict(self, chips, windows): if self.learner is None: self.load_model() batch_out = self.learner.numpy_predict(chips, raw_out=False) labels = SemanticSegmentationLabels() for out, window in zip(batch_out, windows): labels.set_label_arr(window, out) return labels
def get_labels(self, chip_size=1000): """Get all labels. Returns: SemanticSegmentationLabels with windows of size chip_size covering the scene with no overlap. """ if self.source is None: raise Exception('Raster source at {} does not exist'.format( self.uri)) labels = SemanticSegmentationLabels() extent = self.source.get_extent() windows = extent.get_windows(chip_size, chip_size) for w in windows: raw_labels = self.source.get_raw_chip(w) label_arr = (np.squeeze(raw_labels) if self.class_trans is None else self.class_trans.rgb_to_class(raw_labels)) labels.set_label_arr(w, label_arr) return labels
class TestSemanticSegmentationLabels(unittest.TestCase): def setUp(self): self.windows = [Box.make_square(0, 0, 10), Box.make_square(0, 10, 10)] self.label_arr0 = np.random.choice([0, 1], (10, 10)) self.label_arr1 = np.random.choice([0, 1], (10, 10)) self.labels = SemanticSegmentationLabels() self.labels.set_label_arr(self.windows[0], self.label_arr0) self.labels.set_label_arr(self.windows[1], self.label_arr1) def test_get(self): np.testing.assert_array_equal( self.labels.get_label_arr(self.windows[0]), self.label_arr0) def test_get_with_aoi(self): null_class_id = 2 aoi_polygons = [Box.make_square(5, 15, 2).to_shapely()] exp_label_arr = np.full(self.label_arr1.shape, null_class_id) exp_label_arr[5:7, 5:7] = self.label_arr1[5:7, 5:7] labels = self.labels.filter_by_aoi(aoi_polygons, null_class_id) label_arr = labels.get_label_arr(self.windows[1]) np.testing.assert_array_equal(label_arr, exp_label_arr) self.assertEqual(1, len(labels.window_to_label_arr))
def empty_labels(self): """Returns an empty SemanticSegmentationLabels object.""" return SemanticSegmentationLabels()