diff --git a/label_maker/images.py b/label_maker/images.py index d0d258f..2907fe4 100644 --- a/label_maker/images.py +++ b/label_maker/images.py @@ -66,7 +66,7 @@ def class_test(value): # download tiles tiles = class_tiles + background_tiles - print('Downloading {} tiles to {}'.format(len(tiles), op.join(dest_folder, 'tiles'))) + print('Downloading {} tiles to {}'.format(len(tiles), tiles_dir)) # get image acquisition function based on imagery string image_function = download_tile_tms @@ -74,4 +74,4 @@ def class_test(value): image_function = get_tile_tif for tile in tiles: - image_function(tile, imagery, dest_folder, imagery_offset) + image_function(tile, imagery, tiles_dir, imagery_offset) diff --git a/label_maker/preview.py b/label_maker/preview.py index f185fed..aa0e595 100644 --- a/label_maker/preview.py +++ b/label_maker/preview.py @@ -9,7 +9,7 @@ from label_maker.utils import class_match, download_tile_tms, get_tile_tif, is_tif -def preview(dest_folder, number, classes, imagery, ml_type, imagery_offset, **kwargs): +def preview(dest_folder, number, classes, imagery, ml_type, imagery_offset=False, **kwargs): """Produce imagery examples for specified classes Parameters @@ -55,6 +55,7 @@ def preview(dest_folder, number, classes, imagery, ml_type, imagery_offset, **kw for i, cl in enumerate(classes): # create class directory class_dir = op.join(dest_folder, 'examples', cl.get('name')) + if not op.isdir(class_dir): makedirs(class_dir) diff --git a/label_maker/utils.py b/label_maker/utils.py index 9b93a90..804a217 100644 --- a/label_maker/utils.py +++ b/label_maker/utils.py @@ -24,16 +24,16 @@ def class_match(ml_type, label, i): return np.count_nonzero(label == i) return None -def download_tile_tms(tile, imagery, dest_folder, *args): +def download_tile_tms(tile, imagery, folder, *args): """Download a satellite image tile from a tms endpoint""" o = urlparse(imagery) _, image_format = op.splitext(o.path) r = requests.get(url(tile.split('-'), imagery)) - tile_img = op.join(dest_folder, 'tiles', '{}{}'.format(tile, image_format)) + tile_img = op.join(folder, '{}{}'.format(tile, image_format)) open(tile_img, 'wb').write(r.content) return tile_img -def get_tile_tif(tile, imagery, dest_folder, imagery_offset): +def get_tile_tif(tile, imagery, folder, imagery_offset): """ Read a GeoTIFF with a window corresponding to a TMS tile @@ -80,7 +80,7 @@ def get_tile_tif(tile, imagery, dest_folder, imagery_offset): src.read(k, window=window, out=data[k - 1], boundless=True) # save - tile_img = op.join(dest_folder, 'tiles', '{}{}'.format(tile, '.jpg')) + tile_img = op.join(folder, '{}{}'.format(tile, '.jpg')) img = Image.fromarray(np.moveaxis(data, 0, -1), mode='RGB') img.save(tile_img) diff --git a/test/unit/test_utils.py b/test/unit/test_utils.py index 99005d5..73d8635 100644 --- a/test/unit/test_utils.py +++ b/test/unit/test_utils.py @@ -52,7 +52,7 @@ def test_get_tile_tif(self): if not op.isdir(tiles_dir): makedirs(tiles_dir) - get_tile_tif(tile, 'test/fixtures/drone.tif', dest_folder, None) + get_tile_tif(tile, 'test/fixtures/drone.tif', tiles_dir, None) test_tile = Image.open('test/tiles/{}.jpg'.format(tile)) fixture_tile = Image.open('test/fixtures/{}.jpg'.format(tile)) self.assertEqual(test_tile, fixture_tile) @@ -66,7 +66,7 @@ def test_get_tile_tif_offset(self): if not op.isdir(tiles_dir): makedirs(tiles_dir) - get_tile_tif(tile, 'test/fixtures/drone.tif', dest_folder, [128, 64]) + get_tile_tif(tile, 'test/fixtures/drone.tif', tiles_dir, [128, 64]) test_tile = Image.open('test/tiles/{}.jpg'.format(tile)) fixture_tile = Image.open('test/fixtures/{}_offset.jpg'.format(tile)) self.assertEqual(test_tile, fixture_tile) @@ -80,7 +80,7 @@ def test_get_tile_vrt(self): if not op.isdir(tiles_dir): makedirs(tiles_dir) - get_tile_tif(tile, 'test/fixtures/drone.vrt', dest_folder, None) + get_tile_tif(tile, 'test/fixtures/drone.vrt', tiles_dir, None) test_tile = Image.open('test/tiles/{}.jpg'.format(tile)) fixture_tile = Image.open('test/fixtures/{}.jpg'.format(tile)) self.assertEqual(test_tile, fixture_tile)