DAS-2276: Fix browse code for jpeg images.

nasa · Dec 13, 2024 · b17e020 · b17e020
1 parent 97086e5
commit b17e020
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 14 deletions.
diff --git a/hybig/browse.py b/hybig/browse.py
@@ -181,11 +181,11 @@ def create_browse_imagery(
                     f'incorrect number of bands for image: {rio_in_array.rio.count}'
                 )
 
-            if rio_in_array.rio.count == 1:
+            raster, color_map = prepare_raster_for_writing(raster, output_driver)
+
+            if rio_in_array.rio.count == 1 and output_driver == 'PNG':
                 # we only paletize single band input data
-                raster, color_map = prepare_raster_for_writing(raster, output_driver)
-            else:
-                color_map = None
+                raster, color_map = palettize_raster(raster)
 
             grid_parameters = get_target_grid_parameters(message, rio_in_array)
             grid_parameter_list, tile_locators = create_tiled_output_parameters(
@@ -363,12 +363,9 @@ def prepare_raster_for_writing(
     raster: ndarray, driver: str
 ) -> 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
-
-    return palettize_raster(raster)
+    if driver == 'JPEG' and raster.shape[0] == 4:
+        raster = raster[0:3, :, :]
+    return raster, None
 
 
 def palettize_raster(raster: ndarray) -> tuple[ndarray, dict]:

diff --git a/tests/test_service/test_adapter.py b/tests/test_service/test_adapter.py
@@ -470,7 +470,7 @@ 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,
@@ -479,9 +479,10 @@ def move_tif(*args, **kwargs):
                 dst_nodata=255,
                 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
         ):

diff --git a/tests/unit/test_browse.py b/tests/unit/test_browse.py
@@ -677,9 +677,16 @@ def test_prepare_raster_for_writing_png_4band(self, palettize_mock):
         raster = self.random.integers(255, size=(4, 7, 8))
         driver = 'PNG'
 
-        prepare_raster_for_writing(raster, driver)
+        expected, _ = prepare_raster_for_writing(raster, driver)
+        np.testing.assert_array_equal(raster, expected, strict=True)
 
-        palettize_mock.assert_called_once_with(raster)
+    @patch('hybig.browse.palettize_raster')
+    def test_prepare_raster_for_writing_png_3band(self, palettize_mock):
+        raster = self.random.integers(255, size=(3, 7, 8))
+        driver = 'PNG'
+
+        expected, _ = prepare_raster_for_writing(raster, driver)
+        np.testing.assert_array_equal(raster, expected, strict=True)
 
     @patch('hybig.browse.Image')
     @patch('hybig.browse.get_color_map_from_image')