Refactor conditional image download function code

Author: drewbo

label_maker/images.py

@@ -6,7 +6,7 @@
 
 import numpy as np
 
-from label_maker.utils import download_tile_tms, get_tile_tif, is_tif, is_wms, get_tile_wms
+from label_maker.utils import get_image_function
 
 def download_images(dest_folder, classes, imagery, ml_type, background_ratio, imagery_offset=False, **kwargs):
     """Download satellite images specified by a URL and a label.npz file
@@ -69,11 +69,7 @@ def class_test(value):
     print('Downloading {} tiles to {}'.format(len(tiles), tiles_dir))
 
     # get image acquisition function based on imagery string
-    image_function = download_tile_tms
-    if is_tif(imagery):
-        image_function = get_tile_tif
-    if is_wms(imagery):
-        image_function = get_tile_wms
+    image_function = get_image_function(imagery)
 
     for tile in tiles:
         image_function(tile, imagery, tiles_dir, imagery_offset)

label_maker/preview.py

@@ -7,7 +7,7 @@
 import numpy as np
 from PIL import Image, ImageDraw
 
-from label_maker.utils import class_match, download_tile_tms, get_tile_tif, is_tif, get_tile_wms, is_wms
+from label_maker.utils import class_match, get_image_function
 
 def preview(dest_folder, number, classes, imagery, ml_type, imagery_offset=False, **kwargs):
     """Produce imagery examples for specified classes
@@ -48,11 +48,7 @@ def preview(dest_folder, number, classes, imagery, ml_type, imagery_offset=False
     print('Writing example images to {}'.format(examples_dir))
 
     # get image acquisition function based on imagery string
-    image_function = download_tile_tms
-    if is_tif(imagery):
-        image_function = get_tile_tif
-    if is_wms(imagery):
-        image_function = get_tile_wms
+    image_function = get_image_function(imagery)
 
     for i, cl in enumerate(classes):
         # create class directory

label_maker/utils.py

@@ -122,3 +122,11 @@ def is_tif(imagery):
 def is_wms(imagery):
     """Determine if an imagery path is a WMS endpoint"""
     return '{bbox}' in imagery
+
+def get_image_function(imagery):
+    """Return the correct image downloading function based on the imagery string"""
+    if is_tif(imagery):
+        return get_tile_tif
+    if is_wms(imagery):
+        return get_tile_wms
+    return download_tile_tms