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()