Merge pull request #181 from gerritholl/bugfix-fill-value-palette

Respect fill value when saving palette images while keeping the palette.

trollimage/tests/test_image.py

@@ -1,6 +1,4 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-# Copyright (c) 2009-2021 trollimage developers
+# Copyright (c) 2009-2024 trollimage developers
 #
 # This file is part of trollimage.
 #
@@ -1384,6 +1382,25 @@ def test_save_geotiff_int_with_bad_cmap(self):
                 img.save(tmp.name, keep_palette=True, cmap=t_cmap,
                          dtype='uint16')
 
+    def test_save_geotiff_with_cmap_and_fill_value(self, tmp_path):
+        """Test saving GeoTIFF with colormap and fill value."""
+        import rasterio
+        test_file = tmp_path / "test.tif"
+        fv = np.uint8(42)
+        arr = np.ones((1, 5, 5), dtype="uint8")
+        arr[0, 2, 2] = 255
+        data = xr.DataArray(
+                arr,
+                dims=["bands", 'y', 'x'],
+                attrs={"_FillValue": 255},
+                coords={"bands": ["P"]})
+        img = xrimage.XRImage(data)
+        img.save(test_file, keep_palette=True, cmap=brbg,
+                 fill_value=fv)
+        with rasterio.open(test_file) as f:
+            cont = f.read()
+            assert cont[0, 2, 2] == fv
+
     @pytest.mark.skipif(sys.platform.startswith('win'),
                         reason="'NamedTemporaryFile' not supported on Windows")
     def test_save_geotiff_closed_file(self):

trollimage/xrimage.py

@@ -1,13 +1,6 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
+# Copyright (c) 2017-2024 trollimage developers
 #
-# Copyright (c) 2017-2018
-#
-# Author(s):
-#
-#   Martin Raspaud <[email protected]>
-#   Adam Dybbroe <[email protected]>
-#   Esben S. Nielsen <[email protected]>
+# This file is part of trollimage.
 #
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU General Public License as published by
@@ -830,20 +823,25 @@ def _get_input_fill_value(self, data):
             return data.attrs.get('_FillValue')
         return None
 
-    def _scale_and_replace_fill_value(self, data, input_fill_value, fill_value, dtype):
+    def _scale_and_replace_fill_value(self, data, input_fill_value, fill_value,
+                                      dtype, scale, replace_fill_value):
         # scale float data to the proper dtype
         # this method doesn't cast yet so that we can keep track of NULL values
-        data = self._scale_to_dtype(data, dtype, fill_value)
-        data = self._replace_fill_value(data, input_fill_value, fill_value, dtype)
+        if scale:
+            data = self._scale_to_dtype(data, dtype, fill_value)
+        if replace_fill_value:
+            data = self._replace_fill_value(data, input_fill_value, fill_value, dtype)
         return data
 
-    def _scale_alpha_or_fill_data(self, data, fill_value, dtype):
+    def _scale_alpha_or_fill_data(self, data, fill_value, dtype, scale, alpha, replace_fill_value):
         input_fill_value = self._get_input_fill_value(data)
-        needs_alpha = fill_value is None and not self.mode.endswith('A')
+        needs_alpha = alpha and fill_value is None and not self.mode.endswith('A')
         if needs_alpha:
             # We don't have a fill value or an alpha, let's add an alpha
             return self._add_alpha_and_scale(data, input_fill_value, dtype)
-        return self._scale_and_replace_fill_value(data, input_fill_value, fill_value, dtype)
+        return self._scale_and_replace_fill_value(data, input_fill_value,
+                                                  fill_value, dtype, scale,
+                                                  replace_fill_value)
 
     def finalize(self, fill_value=None, dtype=np.uint8, keep_palette=False):
         """Finalize the image to be written to an output file.
@@ -908,8 +906,11 @@ def finalize(self, fill_value=None, dtype=np.uint8, keep_palette=False):
             pass
         with xr.set_options(keep_attrs=True):
             attrs = final_data.attrs
-            if not keep_palette:
-                final_data = self._scale_alpha_or_fill_data(final_data, fill_value, dtype)
+            final_data = self._scale_alpha_or_fill_data(
+                    final_data, fill_value, dtype,
+                    scale=not keep_palette,
+                    alpha=not keep_palette,
+                    replace_fill_value=True)
             final_data = final_data.astype(dtype)
             final_data.attrs = attrs