adding uint8 and tests

requirements.txt

@@ -1,3 +1,4 @@
 click
 rasterio
 rio-mucho>=0.2.1
+rio-color>=0.4

rio_toa/scripts/cli.py

@@ -46,7 +46,7 @@ def radiance(ctx, src_path, src_mtl, dst_path,
 @click.argument('src_path', type=click.Path(exists=True))
 @click.argument('src_mtl', type=click.Path(exists=True))
 @click.argument('dst_path', type=click.Path(exists=False))
-@click.option('--dst-dtype', type=click.Choice(['float32', 'uint16']), default='float32')
+@click.option('--dst-dtype', type=click.Choice(['float32', 'float64', 'uint16', 'uint8']), default='float32')
 @click.option('--rescale-factor', '-r', type=click.Choice([55000.0/2**16, 1.0]), default=55000.0/2**16)
 @click.option('--workers', '-j', type=int, default=4)
 @click.option('--l8-bidx', default=0,

rio_toa/toa_utils.py

@@ -116,6 +116,11 @@ def _get_bounds_from_metadata(product_metadata):
 def rescale(arr, rescale_factor, dtype):
     """Convert an array from 0..1 to dtype, scaling up linearly
     """
-    arr *= rescale_factor * np.iinfo(np.uint16).max
-    return np.clip(arr, 1, np.iinfo(np.uint16).max).astype(dtype)
+    if dtype == np.__dict__['uint8']:
+        arr *= rescale_factor * np.iinfo(np.uint8).max
+        return np.clip(arr, 1, np.iinfo(np.uint8).max).astype(dtype)
+
+    else:
+        arr *= rescale_factor * np.iinfo(np.uint16).max
+        return np.clip(arr, 1, np.iinfo(np.uint16).max).astype(dtype)
 

tests/test_toa_utils.py

@@ -59,12 +59,25 @@ def test_load_mtl_key():
 
 def test_rescale():
 	arr = np.array(np.linspace(0.0, 1.5, num=9).reshape(3,3))
-	dtype = 'uint16'
+	dtype = np.__dict__['uint16']
 	rescale_factor = 1.0
 	rescaled_arr = rescale(arr, rescale_factor, dtype)
 	mask = (rescaled_arr != np.iinfo(np.uint16).max) & (rescaled_arr != 1.0)
 
 	assert np.all(rescaled_arr) <= np.iinfo(np.uint16).max
 	assert np.all(rescaled_arr) >= 0.0
 	assert np.array_equal(rescaled_arr[mask], arr[mask].astype(int))
+	assert rescaled_arr.dtype == 'uint16'
+
+def test_rescale2():
+	arr = np.array(np.linspace(0.0, 1.5, num=9).reshape(3,3))
+	dtype = np.__dict__['uint8']
+	rescale_factor = 1.0
+	rescaled_arr = rescale(arr, rescale_factor, dtype)
+	mask = (rescaled_arr != np.iinfo(dtype).max) & (rescaled_arr != 1.0)
+
+	assert np.all(rescaled_arr) <= np.iinfo(dtype).max
+	assert np.all(rescaled_arr) >= 0.0
+	assert np.array_equal(rescaled_arr[mask], arr[mask].astype(int))
+	assert rescaled_arr.dtype == 'uint8'