1313from harmony_service_lib .message import Source as HarmonySource
1414from matplotlib .cm import ScalarMappable
1515from matplotlib .colors import Normalize
16- from numpy import ndarray
16+ from numpy import ndarray , uint8
1717from osgeo_utils .auxiliary .color_palette import ColorPalette
1818from PIL import Image
1919from rasterio .io import DatasetReader
@@ -181,7 +181,9 @@ def create_browse_imagery(
181181 f'incorrect number of bands for image: { rio_in_array .rio .count }
182182 )
183183
184- raster , color_map = prepare_raster_for_writing (raster , output_driver )
184+ raster , color_map = standardize_raster_for_writing (
185+ raster , output_driver , rio_in_array .rio .count
186+ )
185187
186188 grid_parameters = get_target_grid_parameters (message , rio_in_array )
187189 grid_parameter_list , tile_locators = create_tiled_output_parameters (
@@ -217,12 +219,14 @@ def create_browse_imagery(
217219 return processed_files
218220
219221
220- def convert_mulitband_to_raster (data_array : DataArray ) -> ndarray :
222+ def convert_mulitband_to_raster (data_array : DataArray ) -> ndarray [ uint8 ] :
221223 """Convert multiband to a raster image.
222224
223- Reads the three or four bands from the file, then normalizes them to the range
224- 0 to 255. This assumes the input image is already in RGB or RGBA format and
225- just ensures that the output is 8bit.
225+ Return a 4-band raster, where the alpha layer is presumed to be the missing
226+ data mask.
227+
228+ Convert 3-band data into a 4-band raster by generating an alpha layer from
229+ any missing data in the RGB bands.
226230
227231 """
228232 if data_array .rio .count not in 3 , 4 ]:
@@ -233,26 +237,49 @@ def convert_mulitband_to_raster(data_array: DataArray) -> ndarray:
233237
234238 bands = data_array .to_numpy ()
235239
236- # Create an alpha layer where input NaN values are transparent.
240+ if data_array .rio .count == 4 :
241+ return convert_to_uint8 (bands , original_dtype (data_array ))
242+
243+ # Input NaNs in any of the RGB bands are made transparent.
237244 nan_mask = np .isnan (bands ).any (axis = 0 )
238245 nan_alpha = np .where (nan_mask , TRANSPARENT , OPAQUE )
239246
240- # grab any existing alpha layer
241- bands , image_alpha = remove_alpha (bands )
247+ raster = convert_to_uint8 (bands , original_dtype (data_array ))
242248
243- norm = Normalize (vmin = np .nanmin (bands ), vmax = np .nanmax (bands ))
244- raster = np .nan_to_num (np .around (norm (bands ) * 255.0 ), copy = False , nan = 0.0 ).astype (
245- 'uint8'
246- )
249+ return np .concatenate ((raster , nan_alpha [None , ...]), axis = 0 )
250+
251+
252+ def convert_to_uint8 (bands : ndarray , dtype : str | None ) -> ndarray [uint8 ]:
253+ """Convert Banded data with NaNs (missing) into a uint8 data cube.
254+
255+ Nearly all of the time this will simply pass through the data coercing it
256+ back into unsigned ints and setting the missing values to 0 that will be
257+ masked as transparent in the output png.
247258
248- if image_alpha is not None :
249- # merge missing alpha with the image alpha band prefering transparency
250- # to opaqueness.
251- alpha = np .minimum (nan_alpha , image_alpha ).astype (np .uint8 )
259+ There is a some small non-zero chance that the input RGB image was 16-bit
260+ and if any of the values exceed 255, we must normalize all of input data to
261+ the range 0-255.
262+
263+ """
264+
265+ if dtype != 'uint8' and np .nanmax (bands ) > 255 :
266+ norm = Normalize (vmin = np .nanmin (bands ), vmax = np .nanmax (bands ))
267+ scaled = np .around (norm (bands ) * 255.0 )
268+ raster = scaled .filled (0 ).astype ('uint8' )
252269 else :
253- alpha = nan_alpha
270+ raster = np .nan_to_num (bands ).astype ('uint8' )
271+
272+ return raster
254273
255- return np .concatenate ((raster , alpha [None , ...]), axis = 0 )
274+
275+ def original_dtype (data_array : DataArray ) -> str | None :
276+ """Return the original input data's type.
277+
278+ rastero_open retains the input dtype in the encoding dictionary and is used
279+ to understand what kind of casts are safe.
280+
281+ """
282+ return data_array .encoding .get ('dtype' ) or data_array .encoding .get ('rasterio_dtype' )
256283
257284
258285def convert_singleband_to_raster (
@@ -330,16 +357,38 @@ def image_driver(mime: str) -> str:
330357 return 'PNG'
331358
332359
333- def prepare_raster_for_writing (
334- raster : ndarray , driver : str
360+ def standardize_raster_for_writing (
361+ raster : ndarray ,
362+ driver : str ,
363+ band_count : int ,
335364) -> tuple [ndarray , dict | None ]:
336- """Remove alpha layer if writing a jpeg."""
337- if driver == 'JPEG' :
338- if raster .shape [0 ] == 4 :
339- raster = raster [0 :3 , :, :]
340- return raster , None
365+ """Standardize raster data for writing to browse image.
341366
342- return palettize_raster (raster )
367+ Args:
368+ raster: Input raster data array
369+ driver: Output image format ('JPEG' or 'PNG')
370+ band_count: Number of bands in original input data
371+
372+ The function handles two special cases:
373+ - JPEG output with 4-band data -> Drop alpha channel and return 3-band RGB
374+ - PNG output with single-band data -> Convert to paletted format
375+
376+ Returns:
377+ tuple: (prepared_raster, color_map) where:
378+ - prepared_raster is the processed ndarray
379+ - color_map is either None or a dict mapping palette indices to RGBA values
380+
381+
382+ """
383+ if driver == 'JPEG' and raster .shape [0 ] == 4 :
384+ return raster [0 :3 , :, :], None
385+
386+ if driver == 'PNG' and band_count == 1 :
387+ # Only palettize single band input data that has been converted to an
388+ # RGBA raster.
389+ return palettize_raster (raster )
390+
391+ return raster , None
343392
344393
345394def palettize_raster (raster : ndarray ) -> tuple [ndarray , dict ]:
@@ -476,9 +525,13 @@ def write_georaster_as_browse(
476525
477526 """
478527 n_bands = raster .shape [0 ]
479- dst_nodata = NODATA_IDX
528+
480529 if color_map is not None :
530+ dst_nodata = NODATA_IDX
481531 color_map [dst_nodata ] = NODATA_RGBA
532+ else :
533+ # for banded data set the each band's destination nodata to zero (TRANSPARENT).
534+ dst_nodata = TRANSPARENT
482535
483536 creation_options = {
484537 ** grid_parameters ,
0 commit comments