From 09081c9b7129d9ca44a29281cce501abaecb2404 Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Wed, 18 Dec 2024 16:17:09 -0700 Subject: [PATCH] DAS-2276: retain 3 and 4 band information in browse images (#39) --- CHANGELOG.md | 26 ++-- docker/service_version.txt | 2 +- hybig/browse.py | 109 +++++++++---- hybig/sizes.py | 2 +- tests/test_service/test_adapter.py | 29 ++-- tests/unit/test_browse.py | 240 ++++++++++++++++++++--------- 6 files changed, 285 insertions(+), 123 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 53a76e7..1ce6c2e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,10 +4,11 @@ HyBIG follows semantic versioning. All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). -## [unreleased] - 2024-12-10 +## [v2.1.0] - 2024-12-13 ### Changed +* Input GeoTIFF RGB[A] images are **no longer palettized** when converted to a PNG. The new resulting output browse images are now 3 or 4 band PNG retaining the color information of the input image.[#39](https://github.com/nasa/harmony-browse-image-generator/pull/39) * Changed pre-commit configuration to remove `black-jupyter` dependency [#38](https://github.com/nasa/harmony-browse-image-generator/pull/38) * Updates service image's python to 3.12 [#38](https://github.com/nasa/harmony-browse-image-generator/pull/38) * Simplifies test scripts to run with pytest and pytest plugins [#38](https://github.com/nasa/harmony-browse-image-generator/pull/38) @@ -90,14 +91,15 @@ outlined by the NASA open-source guidelines. For more information on internal releases prior to NASA open-source approval, see legacy-CHANGELOG.md. -[unreleased]:https://github.com/nasa/harmony-browse-image-generator/compare/2.0.2..HEAD -[v2.0.2]:https://github.com/nasa/harmony-browse-image-generator/compare/2.0.1..2.0.2 -[v2.0.1]:https://github.com/nasa/harmony-browse-image-generator/compare/2.0.0..2.0.1 -[v2.0.0]:https://github.com/nasa/harmony-browse-image-generator/compare/1.2.2..2.0.0 -[v1.2.2]: https://github.com/nasa/harmony-browse-image-generator/compare/1.2.1..1.2.2 -[v1.2.1]: https://github.com/nasa/harmony-browse-image-generator/compare/1.2.0..1.2.1 -[v1.2.0]: https://github.com/nasa/harmony-browse-image-generator/compare/1.1.0..1.2.0 -[v1.1.0]: https://github.com/nasa/harmony-browse-image-generator/compare/1.0.2..1.1.0 -[v1.0.2]: https://github.com/nasa/harmony-browse-image-generator/compare/1.0.1..1.0.2 -[v1.0.1]: https://github.com/nasa/harmony-browse-image-generator/compare/1.0.0..1.0.1 -[v1.0.0]: https://github.com/nasa/harmony-browse-image-generator/compare/0.0.11-legacy..1.0.0 +[unreleased]: https://github.com/nasa/harmony-browse-image-generator/ +[v2.1.0]: https://github.com/nasa/harmony-browse-image-generator/releases/tag/2.1.0 +[v2.0.2]: https://github.com/nasa/harmony-browse-image-generator/releases/tag/2.0.2 +[v2.0.1]: https://github.com/nasa/harmony-browse-image-generator/releases/tag/2.0.1 +[v2.0.0]: https://github.com/nasa/harmony-browse-image-generator/releases/tag/2.0.0 +[v1.2.2]: https://github.com/nasa/harmony-browse-image-generator/releases/tag/1.2.2 +[v1.2.1]: https://github.com/nasa/harmony-browse-image-generator/releases/tag/1.2.1 +[v1.2.0]: https://github.com/nasa/harmony-browse-image-generator/releases/tag/1.2.0 +[v1.1.0]: https://github.com/nasa/harmony-browse-image-generator/releases/tag/1.1.0 +[v1.0.2]: https://github.com/nasa/harmony-browse-image-generator/releases/tag/1.0.2 +[v1.0.1]: https://github.com/nasa/harmony-browse-image-generator/releases/tag/1.0.1 +[v1.0.0]: https://github.com/nasa/harmony-browse-image-generator/releases/tag/1.0.0 diff --git a/docker/service_version.txt b/docker/service_version.txt index e9307ca..7ec1d6d 100644 --- a/docker/service_version.txt +++ b/docker/service_version.txt @@ -1 +1 @@ -2.0.2 +2.1.0 diff --git a/hybig/browse.py b/hybig/browse.py index 9af788f..8c120fd 100644 --- a/hybig/browse.py +++ b/hybig/browse.py @@ -13,7 +13,7 @@ from harmony_service_lib.message import Source as HarmonySource from matplotlib.cm import ScalarMappable from matplotlib.colors import Normalize -from numpy import ndarray +from numpy import ndarray, uint8 from osgeo_utils.auxiliary.color_palette import ColorPalette from PIL import Image from rasterio.io import DatasetReader @@ -181,7 +181,9 @@ def create_browse_imagery( f'incorrect number of bands for image: {rio_in_array.rio.count}' ) - raster, color_map = prepare_raster_for_writing(raster, output_driver) + raster, color_map = standardize_raster_for_writing( + raster, output_driver, rio_in_array.rio.count + ) grid_parameters = get_target_grid_parameters(message, rio_in_array) grid_parameter_list, tile_locators = create_tiled_output_parameters( @@ -217,12 +219,14 @@ def create_browse_imagery( return processed_files -def convert_mulitband_to_raster(data_array: DataArray) -> ndarray: +def convert_mulitband_to_raster(data_array: DataArray) -> ndarray[uint8]: """Convert multiband to a raster image. - Reads the three or four bands from the file, then normalizes them to the range - 0 to 255. This assumes the input image is already in RGB or RGBA format and - just ensures that the output is 8bit. + Return a 4-band raster, where the alpha layer is presumed to be the missing + data mask. + + Convert 3-band data into a 4-band raster by generating an alpha layer from + any missing data in the RGB bands. """ if data_array.rio.count not in [3, 4]: @@ -233,26 +237,49 @@ def convert_mulitband_to_raster(data_array: DataArray) -> ndarray: bands = data_array.to_numpy() - # Create an alpha layer where input NaN values are transparent. + if data_array.rio.count == 4: + return convert_to_uint8(bands, original_dtype(data_array)) + + # Input NaNs in any of the RGB bands are made transparent. nan_mask = np.isnan(bands).any(axis=0) nan_alpha = np.where(nan_mask, TRANSPARENT, OPAQUE) - # grab any existing alpha layer - bands, image_alpha = remove_alpha(bands) + raster = convert_to_uint8(bands, original_dtype(data_array)) - norm = Normalize(vmin=np.nanmin(bands), vmax=np.nanmax(bands)) - raster = np.nan_to_num(np.around(norm(bands) * 255.0), copy=False, nan=0.0).astype( - 'uint8' - ) + return np.concatenate((raster, nan_alpha[None, ...]), axis=0) + + +def convert_to_uint8(bands: ndarray, dtype: str | None) -> ndarray[uint8]: + """Convert Banded data with NaNs (missing) into a uint8 data cube. + + Nearly all of the time this will simply pass through the data coercing it + back into unsigned ints and setting the missing values to 0 that will be + masked as transparent in the output png. - if image_alpha is not None: - # merge missing alpha with the image alpha band prefering transparency - # to opaqueness. - alpha = np.minimum(nan_alpha, image_alpha).astype(np.uint8) + There is a some small non-zero chance that the input RGB image was 16-bit + and if any of the values exceed 255, we must normalize all of input data to + the range 0-255. + + """ + + if dtype != 'uint8' and np.nanmax(bands) > 255: + norm = Normalize(vmin=np.nanmin(bands), vmax=np.nanmax(bands)) + scaled = np.around(norm(bands) * 255.0) + raster = scaled.filled(0).astype('uint8') else: - alpha = nan_alpha + raster = np.nan_to_num(bands).astype('uint8') + + return raster - return np.concatenate((raster, alpha[None, ...]), axis=0) + +def original_dtype(data_array: DataArray) -> str | None: + """Return the original input data's type. + + rastero_open retains the input dtype in the encoding dictionary and is used + to understand what kind of casts are safe. + + """ + return data_array.encoding.get('dtype') or data_array.encoding.get('rasterio_dtype') def convert_singleband_to_raster( @@ -330,16 +357,38 @@ def image_driver(mime: str) -> str: return 'PNG' -def prepare_raster_for_writing( - raster: ndarray, driver: str +def standardize_raster_for_writing( + raster: ndarray, + driver: str, + band_count: int, ) -> tuple[ndarray, dict | None]: - """Remove alpha layer if writing a jpeg.""" - if driver == 'JPEG': - if raster.shape[0] == 4: - raster = raster[0:3, :, :] - return raster, None + """Standardize raster data for writing to browse image. - return palettize_raster(raster) + Args: + raster: Input raster data array + driver: Output image format ('JPEG' or 'PNG') + band_count: Number of bands in original input data + + The function handles two special cases: + - JPEG output with 4-band data -> Drop alpha channel and return 3-band RGB + - PNG output with single-band data -> Convert to paletted format + + Returns: + tuple: (prepared_raster, color_map) where: + - prepared_raster is the processed ndarray + - color_map is either None or a dict mapping palette indices to RGBA values + + + """ + if driver == 'JPEG' and raster.shape[0] == 4: + return raster[0:3, :, :], None + + if driver == 'PNG' and band_count == 1: + # Only palettize single band input data that has been converted to an + # RGBA raster. + return palettize_raster(raster) + + return raster, None def palettize_raster(raster: ndarray) -> tuple[ndarray, dict]: @@ -476,9 +525,13 @@ def write_georaster_as_browse( """ n_bands = raster.shape[0] - dst_nodata = NODATA_IDX + if color_map is not None: + dst_nodata = NODATA_IDX color_map[dst_nodata] = NODATA_RGBA + else: + # for banded data set the each band's destination nodata to zero (TRANSPARENT). + dst_nodata = TRANSPARENT creation_options = { **grid_parameters, diff --git a/hybig/sizes.py b/hybig/sizes.py index f1d58f0..c2fd2e7 100644 --- a/hybig/sizes.py +++ b/hybig/sizes.py @@ -422,7 +422,7 @@ def find_closest_resolution( """ best_info = None - smallest_diff = np.Infinity + smallest_diff = np.inf for res in resolutions: for info in resolution_info: resolution_diff = np.abs(res - info.pixel_size) diff --git a/tests/test_service/test_adapter.py b/tests/test_service/test_adapter.py index 3e30b7b..eb11b65 100644 --- a/tests/test_service/test_adapter.py +++ b/tests/test_service/test_adapter.py @@ -19,7 +19,7 @@ from harmony_service.exceptions import HyBIGServiceError from hybig.browse import ( convert_mulitband_to_raster, - prepare_raster_for_writing, + standardize_raster_for_writing, ) from tests.utilities import Granule, create_stac @@ -270,10 +270,14 @@ def move_tif(*args, **kwargs): mock_reproject.call_args_list, expected_reproject_calls ): np.testing.assert_array_equal( - actual_call.kwargs['source'], expected_call.kwargs['source'] + actual_call.kwargs['source'], + expected_call.kwargs['source'], + strict=True, ) np.testing.assert_array_equal( - actual_call.kwargs['destination'], expected_call.kwargs['destination'] + actual_call.kwargs['destination'], + expected_call.kwargs['destination'], + strict=True, ) self.assertEqual( actual_call.kwargs['src_transform'], @@ -452,11 +456,11 @@ def move_tif(*args, **kwargs): 'transform': expected_transform, 'driver': 'PNG', 'dtype': 'uint8', - 'dst_nodata': 255, + 'dst_nodata': 0, 'count': 3, } raster = convert_mulitband_to_raster(rio_data_array) - raster, color_map = prepare_raster_for_writing(raster, 'PNG') + raster, color_map = standardize_raster_for_writing(raster, 'PNG', 3) dest = np.full( (expected_params['height'], expected_params['width']), @@ -466,26 +470,31 @@ def move_tif(*args, **kwargs): expected_reproject_calls = [ call( - source=raster[0, :, :], + source=raster[band, :, :], destination=dest, src_transform=rio_data_array.rio.transform(), src_crs=rio_data_array.rio.crs, dst_transform=expected_params['transform'], dst_crs=expected_params['crs'], - dst_nodata=255, + dst_nodata=expected_params['dst_nodata'], resampling=Resampling.nearest, ) + for band in range(4) ] - self.assertEqual(mock_reproject.call_count, 1) + self.assertEqual(mock_reproject.call_count, 4) for actual_call, expected_call in zip( mock_reproject.call_args_list, expected_reproject_calls ): np.testing.assert_array_equal( - actual_call.kwargs['source'], expected_call.kwargs['source'] + actual_call.kwargs['source'], + expected_call.kwargs['source'], + strict=True, ) np.testing.assert_array_equal( - actual_call.kwargs['destination'], expected_call.kwargs['destination'] + actual_call.kwargs['destination'], + expected_call.kwargs['destination'], + strict=True, ) self.assertEqual( actual_call.kwargs['src_transform'], diff --git a/tests/unit/test_browse.py b/tests/unit/test_browse.py index 6e25ac8..4f9c1c2 100644 --- a/tests/unit/test_browse.py +++ b/tests/unit/test_browse.py @@ -30,7 +30,7 @@ output_image_file, output_world_file, palettize_raster, - prepare_raster_for_writing, + standardize_raster_for_writing, validate_file_crs, validate_file_type, ) @@ -237,7 +237,7 @@ def test_create_browse_imagery_with_mocks( src_crs=da_mock.rio.crs, dst_transform=target_transform, dst_crs=CRS.from_string('EPSG:4326'), - dst_nodata=255, + dst_nodata=0, resampling=Resampling.nearest, ), call( @@ -247,7 +247,7 @@ def test_create_browse_imagery_with_mocks( src_crs=da_mock.rio.crs, dst_transform=target_transform, dst_crs=CRS.from_string('EPSG:4326'), - dst_nodata=255, + dst_nodata=0, resampling=Resampling.nearest, ), call( @@ -257,7 +257,7 @@ def test_create_browse_imagery_with_mocks( src_crs=da_mock.rio.crs, dst_transform=target_transform, dst_crs=CRS.from_string('EPSG:4326'), - dst_nodata=255, + dst_nodata=0, resampling=Resampling.nearest, ), ] @@ -266,10 +266,14 @@ def test_create_browse_imagery_with_mocks( reproject_mock.call_args_list, expected_calls ): np.testing.assert_array_equal( - actual_call.kwargs['source'], expected_call.kwargs['source'] + actual_call.kwargs['source'], + expected_call.kwargs['source'], + strict=True, ) np.testing.assert_array_equal( - actual_call.kwargs['destination'], expected_call.kwargs['destination'] + actual_call.kwargs['destination'], + expected_call.kwargs['destination'], + strict=True, ) self.assertEqual( actual_call.kwargs['src_transform'], @@ -339,7 +343,7 @@ def test_convert_singleband_to_raster_without_colortable(self): dtype='uint8', ) actual_raster = convert_singleband_to_raster(ds, None) - assert_array_equal(expected_raster, actual_raster) + assert_array_equal(expected_raster, actual_raster, strict=True) def test_convert_singleband_to_raster_with_colormap(self): ds = DataArray(self.data).expand_dims('band') @@ -376,7 +380,7 @@ def test_convert_singleband_to_raster_with_colormap(self): # Read down: red, yellow, green, blue image_palette = convert_colormap_to_palette(self.colormap) actual_raster = convert_singleband_to_raster(ds, image_palette) - assert_array_equal(expected_raster, actual_raster) + assert_array_equal(expected_raster, actual_raster, strict=True) def test_convert_singleband_to_raster_with_colormap_and_bad_data(self): data_array = np.array(self.data, dtype='float') @@ -419,9 +423,10 @@ def test_convert_singleband_to_raster_with_colormap_and_bad_data(self): image_palette = convert_colormap_to_palette(colormap) actual_raster = convert_singleband_to_raster(ds, image_palette) - assert_array_equal(expected_raster, actual_raster) + assert_array_equal(expected_raster, actual_raster, strict=True) - def test_convert_3_multiband_to_raster(self): + def test_convert_uint16_3_multiband_to_raster(self): + """Test that uint16 input scales the output.""" bad_data = np.copy(self.data).astype('float64') bad_data[1][1] = np.nan bad_data[1][2] = np.nan @@ -429,6 +434,7 @@ def test_convert_3_multiband_to_raster(self): np.stack([self.data, bad_data, self.data]), dims=('band', 'y', 'x'), ) + ds.encoding = {'dtype': 'uint16'} expected_raster = np.array( [ @@ -461,64 +467,110 @@ def test_convert_3_multiband_to_raster(self): ) actual_raster = convert_mulitband_to_raster(ds) - assert_array_equal(expected_raster, actual_raster.data) + assert_array_equal(expected_raster, actual_raster.data, strict=True) - def test_convert_4_multiband_to_raster(self): - """Input data has NaN _fillValue match in the red layer at [1,1] - and alpha channel also exists with a single transparent value at [0,0] + def test_convert_uint8_3_multiband_to_raster(self): + """Ensure valid data is unchanged when input is uint8.""" + scale_data = np.array( + [ + [10, 200, 30, 40], + [10, np.nan, np.nan, 40], + [10, 200, 30, 40], + [10, 200, 30, 40], + ] + ).astype('float32') - See that the expected output has transformed the missing data [nan] - into fully transparent at [1,1] and retained the transparent value of 1 - at [0,0] + ds = DataArray( + np.stack([scale_data, scale_data, scale_data]), + dims=('band', 'y', 'x'), + ) + ds.encoding = {'dtype': 'uint8'} - """ - ds = Mock(DataArray) - bad_data = np.copy(self.data).astype('float64') - bad_data[1, 1] = np.nan + expected_data = scale_data.copy() + expected_data[1][1] = 0 + expected_data[1][2] = 0 - alpha = np.ones_like(self.data) * 255 - alpha[0, 0] = 1 - ds.rio.count = 4 - ds.to_numpy.return_value = np.stack([bad_data, self.data, self.data, alpha]) expected_raster = np.array( [ + expected_data, + expected_data, + expected_data, [ - [0, 85, 170, 255], - [0, 0, 170, 255], - [0, 85, 170, 255], - [0, 85, 170, 255], - ], - [ - [0, 85, 170, 255], - [0, 85, 170, 255], - [0, 85, 170, 255], - [0, 85, 170, 255], - ], - [ - [0, 85, 170, 255], - [0, 85, 170, 255], - [0, 85, 170, 255], - [0, 85, 170, 255], - ], - [ - [1, 255, 255, 255], - [255, 0, 255, 255], - [255, 255, 255, 255], - [255, 255, 255, 255], + [OPAQUE, OPAQUE, OPAQUE, OPAQUE], + [OPAQUE, TRANSPARENT, TRANSPARENT, OPAQUE], + [OPAQUE, OPAQUE, OPAQUE, OPAQUE], + [OPAQUE, OPAQUE, OPAQUE, OPAQUE], ], ], dtype='uint8', ) actual_raster = convert_mulitband_to_raster(ds) - assert_array_equal(expected_raster, actual_raster.data) + assert_array_equal(expected_raster, actual_raster.data, strict=True) - def test_convert_4_multiband_masked_to_raster(self): - """Input data is selected from a subset of a real OPERA RTC input data - file that has masked the alpha layer and as a result as a datatype of - float32. + def test_convert_4_multiband_uint8_to_raster(self): + """4-band 'uint8' images are returned unchanged.""" + ds = Mock(DataArray) + ds.rio.count = 4 - """ + r_data = np.array( + [ + [10, 200, 30, 40], + [10, 200, 30, 40], + [10, 200, 30, 40], + [10, 200, 30, 40], + ] + ).astype('uint8') + + g_data = r_data.copy() + b_data = r_data.copy() + + a_data = np.ones_like(r_data) * 255 + a_data[0, 0] = 0 + + to_numpy_result = np.stack([r_data, g_data, b_data, a_data]) + + ds.to_numpy.return_value = to_numpy_result + + expected_raster = to_numpy_result + + actual_raster = convert_mulitband_to_raster(ds) + assert_array_equal(expected_raster, actual_raster.data, strict=True) + + def test_convert_4_multiband_uint16_to_raster(self): + """4-band 'uint16' images are scaled if their range exceeds 255.""" + ds = Mock(DataArray) + ds.rio.count = 4 + + r_data = np.array( + [ + [10, 200, 300, 400], + [10, 200, 300, 400], + [10, 200, 300, 400], + [10, 200, 300, 400], + ] + ).astype('uint16') + g_data = r_data.copy() + b_data = r_data.copy() + + a_data = np.ones_like(self.data) * OPAQUE + a_data[0, 0] = TRANSPARENT + + to_numpy_result = np.stack([r_data, g_data, b_data, a_data]) + + ds.to_numpy.return_value = to_numpy_result + + # expect the input data to have the data values from 0 to 400 to be + # scaled into the range 0 to 255. + expected_raster = np.around( + np.interp(to_numpy_result, (0, 400), (0.0, 1.0)) * 255.0 + ).astype('uint8') + + actual_raster = convert_mulitband_to_raster(ds) + assert_array_equal(expected_raster, actual_raster.data, strict=True) + + def test_convert_4_multiband_masked_to_raster(self): + """4-band images are returned with nan -> 0""" ds = Mock(DataArray) ds.rio.count = 4 nan = np.nan @@ -553,18 +605,38 @@ def test_convert_4_multiband_masked_to_raster(self): ) ds.to_numpy.return_value = input_array - expected_raster = np.ma.array( - data=[ - [[0, 0, 0, 121], [0, 0, 0, 64], [0, 0, 255, 0], [0, 0, 13, 255]], - [[0, 0, 0, 255], [0, 0, 0, 255], [0, 0, 255, 255], [0, 0, 255, 255]], - [[0, 0, 0, 121], [0, 0, 0, 64], [0, 0, 255, 0], [0, 0, 13, 255]], - [[0, 0, 0, 255], [0, 0, 0, 255], [0, 0, 255, 255], [0, 0, 255, 255]], + expected_raster = np.array( + [ + [ + [0, 0, 0, 234], + [0, 0, 0, 225], + [0, 0, 255, 215], + [0, 0, 217, 255], + ], + [ + [0, 0, 0, 255], + [0, 0, 0, 255], + [0, 0, 255, 255], + [0, 0, 255, 255], + ], + [ + [0, 0, 0, 234], + [0, 0, 0, 225], + [0, 0, 255, 215], + [0, 0, 217, 255], + ], + [ + [0, 0, 0, 255], + [0, 0, 0, 255], + [0, 0, 255, 255], + [0, 0, 255, 255], + ], ], dtype=np.uint8, ) actual_raster = convert_mulitband_to_raster(ds) - assert_array_equal(expected_raster.data, actual_raster.data) + assert_array_equal(expected_raster.data, actual_raster.data, strict=True) def test_convert_5_multiband_to_raster(self): ds = Mock(DataArray) @@ -581,33 +653,59 @@ def test_convert_5_multiband_to_raster(self): 'Cannot create image from 5 band image. Expecting 3 or 4 bands.', ) - def test_prepare_raster_for_writing_jpeg_3band(self): + def test_standardize_raster_for_writing_jpeg_3band(self): raster = self.random.integers(255, size=(3, 5, 6)) + count = 'irrelevant' driver = 'JPEG' expected_raster = np.copy(raster) expected_color_map = None - actual_raster, actual_color_map = prepare_raster_for_writing(raster, driver) + actual_raster, actual_color_map = standardize_raster_for_writing( + raster, driver, count + ) self.assertEqual(expected_color_map, actual_color_map) - np.testing.assert_array_equal(expected_raster, actual_raster) + np.testing.assert_array_equal(expected_raster, actual_raster, strict=True) - def test_prepare_raster_for_writing_jpeg_4band(self): + def test_standardize_raster_for_writing_jpeg_4band(self): raster = self.random.integers(255, size=(4, 7, 8)) driver = 'JPEG' + count = 'irrelevant' expected_raster = np.copy(raster[0:3, :, :]) expected_color_map = None - actual_raster, actual_color_map = prepare_raster_for_writing(raster, driver) + actual_raster, actual_color_map = standardize_raster_for_writing( + raster, driver, count + ) self.assertEqual(expected_color_map, actual_color_map) - np.testing.assert_array_equal(expected_raster, actual_raster) + np.testing.assert_array_equal(expected_raster, actual_raster, strict=True) @patch('hybig.browse.palettize_raster') - def test_prepare_raster_for_writing_png_4band(self, palettize_mock): + def test_standardize_raster_for_writing_png_4band(self, palettize_mock): raster = self.random.integers(255, size=(4, 7, 8)) driver = 'PNG' + count = 'not 1' + + expected, _ = standardize_raster_for_writing(raster, driver, count) + np.testing.assert_array_equal(raster, expected, strict=True) + palettize_mock.assert_not_called() + + @patch('hybig.browse.palettize_raster') + def test_standardize_raster_for_writing_png_3band(self, palettize_mock): + raster = self.random.integers(255, size=(3, 7, 8)) + driver = 'PNG' + count = 'not 1' - prepare_raster_for_writing(raster, driver) + expected, _ = standardize_raster_for_writing(raster, driver, count) + np.testing.assert_array_equal(raster, expected, strict=True) + palettize_mock.assert_not_called() - palettize_mock.assert_called_once_with(raster) + @patch('hybig.browse.palettize_raster') + def test_prepare_1band_raster_for_writing_png(self, palettize_mock): + raster = self.random.integers(255, size=(1, 7, 8)) + driver = 'PNG' + count = 1 + palettize_mock.return_value = (None, None) + expected, _ = standardize_raster_for_writing(raster, driver, count) + palettize_mock.assert_called_with(raster) @patch('hybig.browse.Image') @patch('hybig.browse.get_color_map_from_image') @@ -629,7 +727,7 @@ def test_palettize_raster_no_alpha_layer(self, get_color_map_mock, image_mock): multiband_image_mock.quantize.assert_called_once_with(colors=254) get_color_map_mock.assert_called_once_with(quantized_output) - np.testing.assert_array_equal(expected_out_raster, out_raster) + np.testing.assert_array_equal(expected_out_raster, out_raster, strict=True) @patch('hybig.browse.Image') @patch('hybig.browse.get_color_map_from_image') @@ -657,7 +755,7 @@ def test_palettize_raster_with_alpha_layer(self, get_color_map_mock, image_mock) multiband_image_mock.quantize.assert_called_once_with(colors=254) get_color_map_mock.assert_called_once_with(quantized_output) - np.testing.assert_array_equal(expected_out_raster, out_raster) + np.testing.assert_array_equal(expected_out_raster, out_raster, strict=True) def test_get_color_map_from_image(self): """PIL Image yields a color_map