def setUp(self):
self.windows = [Box.make_square(0, 0, 10), Box.make_square(0, 10, 10)]
self.label_arr0 = np.random.choice([1, 2], (10, 10))
self.label_arr1 = np.random.choice([1, 2], (10, 10))
def label_fn(window):
if window == self.windows[0]:
return self.label_arr0.copy()
elif window == self.windows[1]:
return self.label_arr1.copy()
else:
raise ValueError('Unknown window: {}'.format(window))
self.label_fn = label_fn
self.labels = SemanticSegmentationLabels(self.windows, label_fn)
def predict(self, chips, windows, tmp_dir):
"""Return a prediction for a single chip.
Args:
chips: (numpy.ndarray) of shape (1, height, width, nb_channels)
containing a single imagery chip
windows: List containing a single (Box) window which is aligned
with the chip
Return:
(SemanticSegmentationLabels) containing predictions
"""
self.load_model(tmp_dir)
chips = torch.Tensor(chips).permute((0, 3, 1, 2)) / 255.
chips = chips.to(self.device)
model = self.model.eval()
with torch.no_grad():
out = model(chips)['out'].cpu()
def label_fn(_window):
if _window == windows[0]:
return out[0].argmax(0).squeeze().numpy()
else:
raise ValueError('Trying to get labels for unknown window.')
return SemanticSegmentationLabels(windows, label_fn)
def predict(self, chips: np.ndarray, windows: List[Box],
tmp_dir: str) -> SemanticSegmentationLabels:
"""Predict using an already-trained DeepLab model.
Args:
chips: An np.ndarray containing the image data in a batch of size 1.
tmp_dir: (str) temporary directory to use
Returns:
SemanticSegmentationLabels object with predictions for a single chip
"""
self.load_model(tmp_dir)
label_arr = self.sess.run(
OUTPUT_TENSOR_NAME, feed_dict={INPUT_TENSOR_NAME: [chips[0]]})[0]
# Return "trivial" instance of SemanticSegmentationLabels that holds a single
# window and has ability to get labels for that one window.
def label_fn(_window):
if _window == windows[0]:
return label_arr
else:
raise ValueError('Trying to get labels for unknown window.')
labels = SemanticSegmentationLabels(windows, label_fn)
return labels
def predict(self, chips, windows, tmp_dir):
"""Return a prediction for a single chip.
Args:
chips: (numpy.ndarray) of shape (1, height, width, nb_channels)
containing a single imagery chip
windows: List containing a single (Box) window which is aligned
with the chip
Return:
(SemanticSegmentationLabels) containing predictions
"""
self.load_model(tmp_dir)
chip = torch.Tensor(chips[0]).permute((2, 0, 1)) / 255.
im = Image(chip)
self.inf_learner.data.single_ds.tfmargs['size'] = self.task_config.predict_chip_size
self.inf_learner.data.single_ds.tfmargs_y['size'] = self.task_config.predict_chip_size
if self.train_opts.tta:
label_arr = tta_predict(self.inf_learner, chip)
else:
label_arr = self.inf_learner.predict(im)[1].squeeze().numpy()
# TODO better explanation
# Return "trivial" instance of SemanticSegmentationLabels that holds a single
# window and has ability to get labels for that one window.
def label_fn(_window):
if _window == windows[0]:
return label_arr
else:
raise ValueError('Trying to get labels for unknown window.')
return SemanticSegmentationLabels(windows, label_fn)
def get_labels(self,
window: Union[Box, None] = 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
"""
def label_fn(_window):
raw_labels = self.source.get_raw_chip(_window)
if self.class_transformer is not None:
labels = self.class_transformer.rgb_to_class(raw_labels)
else:
labels = np.squeeze(raw_labels)
return labels
windows = [window]
if window is None:
window = self.source.get_extent()
windows = window.get_windows(chip_size, chip_size)
return SemanticSegmentationLabels(windows, label_fn)
def predict(self, chips, windows, tmp_dir):
"""Return predictions for a chip using model.
Args:
chips: [[height, width, channels], ...] numpy array of chips
windows: List of boxes that are the windows aligned with the chips.
Return:
Labels object containing predictions
"""
self.load_model(tmp_dir)
# TODO get it to work on a whole batch at a time
chip = torch.Tensor(chips[0]).permute((2, 0, 1)) / 255.
im = Image(chip)
label_arr = self.inf_learner.predict(im)[1].squeeze().numpy()
# Return "trivial" instance of SemanticSegmentationLabels that holds a single
# window and has ability to get labels for that one window.
def label_fn(_window):
if _window == windows[0]:
return label_arr
else:
raise ValueError('Trying to get labels for unknown window.')
return SemanticSegmentationLabels(windows, label_fn)
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([1, 2], (10, 10))
self.label_arr1 = np.random.choice([1, 2], (10, 10))
def label_fn(window):
if window == self.windows[0]:
return self.label_arr0.copy()
elif window == self.windows[1]:
return self.label_arr1.copy()
else:
raise ValueError('Unknown window: {}'.format(window))
self.label_fn = label_fn
self.labels = SemanticSegmentationLabels(self.windows, label_fn)
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):
aoi_polygons = [Box.make_square(5, 15, 2).to_shapely()]
exp_label_arr = self.label_arr1.copy()
mask = np.zeros(exp_label_arr.shape)
mask[5:7, 5:7] = 1
exp_label_arr = exp_label_arr * mask
labels = self.labels.filter_by_aoi(aoi_polygons)
label_arr = labels.get_label_arr(self.windows[1])
np.testing.assert_array_equal(label_arr, exp_label_arr)
# Set clip_extent
clip_extent = Box(0, 0, 10, 18)
label_arr = labels.get_label_arr(self.windows[1],
clip_extent=clip_extent)
np.testing.assert_array_equal(label_arr, exp_label_arr[:, 0:8])
def predict(self, chips, windows, tmp_dir):
"""Return predictions for a chip using model.
Args:
chips: [[height, width, channels], ...] numpy array of chips
windows: List of boxes that are the windows aligned with the chips.
Return:
Labels object containing predictions
"""
self.load_model(tmp_dir)
chip = torch.Tensor(chips[0]).permute((2, 0, 1)) / 255.
im = Image(chip)
self.inf_learner.data.single_ds.tfmargs[
'size'] = self.task_config.predict_chip_size
self.inf_learner.data.single_ds.tfmargs_y[
'size'] = self.task_config.predict_chip_size
if self.train_opts.tta:
probs = []
for k in range(8):
trans_im = dihedral(Image(chip), k)
o = self.inf_learner.predict(trans_im)[2]
# https://forums.fast.ai/t/how-best-to-have-get-preds-or-tta-apply-specified-transforms/40731/9
o = Image(o)
if k == 5:
o = dihedral(o, 6)
elif k == 6:
o = dihedral(o, 5)
else:
o = dihedral(o, k)
probs.append(o.data)
label_arr = torch.stack(probs).mean(0).argmax(0).numpy()
else:
label_arr = self.inf_learner.predict(im)[1].squeeze().numpy()
# Return "trivial" instance of SemanticSegmentationLabels that holds a single
# window and has ability to get labels for that one window.
def label_fn(_window):
if _window == windows[0]:
return label_arr
else:
raise ValueError('Trying to get labels for unknown window.')
return SemanticSegmentationLabels(windows, label_fn)
def predict_scene(self, scene, tmp_dir):
"""Predict on a single scene, and return the labels."""
log.info('Making predictions for scene')
raster_source = scene.raster_source
windows = self.get_predict_windows(raster_source.get_extent())
def label_fn(window):
chip = raster_source.get_chip(window)
labels = self.backend.predict([chip], [window], tmp_dir)
label_arr = labels.get_label_arr(window)
# Set NODATA pixels in imagery to predicted value of 0 (ie. ignore)
label_arr[np.sum(chip, axis=2) == 0] = 0
print('.', end='', flush=True)
return label_arr
return SemanticSegmentationLabels(windows, label_fn)
def get_labels(self, chip_size=1000):
"""Get all labels.
Returns:
SemanticSegmentationLabels with windows of size chip_size covering the
scene with no overlap.
"""
def label_fn(window):
raw_labels = self.source.get_raw_chip(window)
if self.class_trans:
labels = self.class_trans.rgb_to_class(raw_labels)
else:
labels = np.squeeze(raw_labels)
return labels
if self.source is None:
raise Exception('Raster source at {} does not exist'.format(
self.uri))
extent = self.source.get_extent()
windows = extent.get_windows(chip_size, chip_size)
return SemanticSegmentationLabels(windows, label_fn)
def predict(self, chips, windows, tmp_dir):
"""Return predictions for a chip using model.
Args:
chips: [[height, width, channels], ...] numpy array of chips
windows: List of boxes that are the windows aligned with the chips.
Return:
Labels object containing predictions
"""
# chips and windows both are arrays, but they always only contain one element (we think)
# Labels should be an array of your classifications (integers), in the same shape as the
# chip. RV assumes that the class IDs start at 1, 0 is an "ignore" class, so it shouldn't
# be included during training. They need to be integers. So I need to do the snapping here
# (and it should happen in the training too, in the same way).
self.load_model(tmp_dir)
func_chips = [np.transpose(chip, (2, 0, 1)) for chip in chips]
# The expected output shape is only one band with the same length and width dimensions
label_arr = apply_to_raster(
self.raster_func,
func_chips[0],
(chips[0].shape[0], chips[0].shape[1])
)
print(label_arr.shape, label_arr.dtype, str(label_arr))
# Return "trivial" instance of SemanticSegmentationLabels that holds a single
# window and has ability to get labels for that one window.
# This is designed to access the prediction for a particular window lazily.
# If the data isn't huge this is just a pass through for the results, which is what is
# happening here.
def label_fn(_window):
if _window == windows[0]:
return label_arr # Do the prediction in here.
else:
raise ValueError('Trying to get labels for unknown window.')
return SemanticSegmentationLabels(windows, label_fn)
def empty_labels(self):
"""Returns an empty SemanticSegmentationLabels object."""
return SemanticSegmentationLabels()