def test_channel_order_error(self):
img_path = join(self.tmp_dir, 'img.tif')
chip = np.ones((2, 2, 3)).astype(np.uint8)
chip[:, :, :] *= np.array([0, 1, 2]).astype(np.uint8)
save_img(chip, img_path)
channel_order = [3, 1, 0]
with self.assertRaises(ChannelOrderError):
config = RasterioSourceConfig(
uris=[img_path], channel_order=channel_order)
config.build(tmp_dir=self.tmp_dir)
def test_uses_channel_order(self):
img_path = join(self.tmp_dir, 'img.tif')
chip = np.ones((2, 2, 4)).astype(np.uint8)
chip[:, :, :] *= np.array([0, 1, 2, 3]).astype(np.uint8)
save_img(chip, img_path)
channel_order = [0, 1, 2]
config = RasterioSourceConfig(
uris=[img_path], channel_order=channel_order)
source = config.build(tmp_dir=self.tmp_dir)
with source.activate():
out_chip = source.get_image_array()
expected_out_chip = np.ones((2, 2, 3)).astype(np.uint8)
expected_out_chip[:, :, :] *= np.array([0, 1, 2]).astype(np.uint8)
np.testing.assert_equal(out_chip, expected_out_chip)
def test_non_geo(self):
# Check if non-georeferenced image files can be read and CRSTransformer
# implements the identity function.
img_path = join(self.tmp_dir, 'img.png')
chip = np.ones((2, 2, 3)).astype(np.uint8)
save_img(chip, img_path)
config = RasterioSourceConfig(uris=[img_path])
source = config.build(tmp_dir=self.tmp_dir)
with source.activate():
out_chip = source.get_image_array()
np.testing.assert_equal(out_chip, chip)
p = (3, 4)
out_p = source.get_crs_transformer().map_to_pixel(p)
np.testing.assert_equal(out_p, p)
out_p = source.get_crs_transformer().pixel_to_map(p)
np.testing.assert_equal(out_p, p)
def test_detects_alpha(self):
# Set first channel to alpha. Expectation is that when omitting channel_order,
# only the second and third channels will be in output.
img_path = join(self.tmp_dir, 'img.tif')
chip = np.ones((2, 2, 3)).astype(np.uint8)
chip[:, :, :] *= np.array([0, 1, 2]).astype(np.uint8)
save_img(chip, img_path)
ci = (ColorInterp.alpha, ColorInterp.blue, ColorInterp.green)
with rasterio.open(img_path, 'r+') as src:
src.colorinterp = ci
config = RasterioSourceConfig(uris=[img_path])
source = config.build(tmp_dir=self.tmp_dir)
with source.activate():
out_chip = source.get_image_array()
expected_out_chip = np.ones((2, 2, 2)).astype(np.uint8)
expected_out_chip[:, :, :] *= np.array([1, 2]).astype(np.uint8)
np.testing.assert_equal(out_chip, expected_out_chip)
def write_sample(self, sample: DataSample):
"""
This writes a training or validation sample to
(train|valid)/{class_name}/{scene_id}-{ind}.png
"""
class_id = sample.labels.get_cell_class_id(sample.window)
# If a chip is not associated with a class, don't
# use it in training data.
if class_id is None:
return
split_name = 'train' if sample.is_train else 'valid'
class_name = self.class_config.names[class_id]
class_dir = join(self.sample_dir, split_name, class_name)
make_dir(class_dir)
chip_path = join(class_dir, '{}-{}.png'.format(sample.scene_id,
self.sample_ind))
save_img(sample.chip, chip_path)
self.sample_ind += 1
def write_sample(self, sample: DataSample):
"""
This writes a training or validation sample to
(train|valid)/img/{scene_id}-{ind}.png and
(train|valid)/labels/{scene_id}-{ind}.png
"""
split_name = 'train' if sample.is_train else 'valid'
label_arr = sample.labels.get_label_arr(sample.window).astype(np.uint8)
img_dir = join(self.sample_dir, split_name, 'img')
labels_dir = join(self.sample_dir, split_name, 'labels')
make_dir(img_dir)
make_dir(labels_dir)
img_path = join(img_dir, '{}-{}.png'.format(sample.scene_id,
self.sample_ind))
labels_path = join(
labels_dir, '{}-{}.png'.format(sample.scene_id, self.sample_ind))
save_img(sample.chip, img_path)
save_img(label_arr, labels_path)
self.sample_ind += 1
def write_sample(self, sample: DataSample):
"""
This writes a training or validation sample to
(train|valid)/img/{scene_id}-{ind}.png and updates
some COCO data structures.
"""
split = 'train' if sample.is_train else 'valid'
split_dir = join(self.sample_dir, split)
img_dir = join(split_dir, 'img')
make_dir(img_dir)
img_fn = '{}-{}.png'.format(sample.scene_id, self.sample_ind)
img_path = join(img_dir, img_fn)
save_img(sample.chip, img_path)
images = self.splits[split]['images']
annotations = self.splits[split]['annotations']
images.append({
'file_name': img_fn,
'id': self.sample_ind,
'height': sample.chip.shape[0],
'width': sample.chip.shape[1]
})
npboxes = sample.labels.get_npboxes()
npboxes = ObjectDetectionLabels.global_to_local(npboxes, sample.window)
for box_ind, (box, class_id) in enumerate(
zip(npboxes, sample.labels.get_class_ids())):
bbox = [box[1], box[0], box[3] - box[1], box[2] - box[0]]
bbox = [int(i) for i in bbox]
annotations.append({
'id': '{}-{}'.format(self.sample_ind, box_ind),
'image_id': self.sample_ind,
'bbox': bbox,
'category_id': int(class_id)
})
self.sample_ind += 1
def write_sample(self, sample: DataSample):
"""
This writes a training or validation sample to
(train|valid)/img/{scene_id}-{ind}.png and
(train|valid)/labels/{scene_id}-{ind}.png
"""
split_name = 'train' if sample.is_train else 'valid'
sample_dir = Path(self.sample_dir)
img_dir = sample_dir / split_name / 'img'
label_dir = sample_dir / split_name / 'labels'
make_dir(img_dir)
make_dir(label_dir)
img = sample.chip
label_arr = sample.labels.get_label_arr(sample.window).astype(np.uint8)
img_fmt, label_fmt = self.img_format, self.label_format
sample_name = f'{sample.scene_id}-{self.sample_ind}'
# write image
img_filename = f'{sample_name}.{img_fmt}'
img_path = img_dir / img_filename
if img_fmt == 'npy':
np.save(img_path, img)
else:
save_img(img, img_path)
# write labels
label_filename = f'{sample_name}.{label_fmt}'
label_path = label_dir / label_filename
if label_fmt == 'npy':
np.save(label_path, label_arr)
else:
save_img(label_arr, label_path)
self.sample_ind += 1
