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