Coverage for python / lsst / pipe / tasks / prettyPictureMaker / _functors / _gamut_fixer.py: 18%
68 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-22 09:17 +0000
1# This file is part of pipe_tasks.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <https://www.gnu.org/licenses/>.
22from __future__ import annotations
24__all__ = ("GamutFixer",)
26import skimage
27import numpy as np
28import logging
31from lsst.cpputils import fixGamutOK
32from lsst.pipe.tasks.prettyPictureMaker.types import LABImage, RGBImage
33from lsst.pex.config.configurableActions import ConfigurableAction
34from lsst.pex.config import ChoiceField, Field
35from lsst.rubinoxide import rgb
36from scipy.ndimage import label, find_objects, binary_dilation
39def heal_gamut(
40 lab_image: LABImage,
41 mask: np.ndarray[tuple[int, int], np.bool],
42 xyz_whitepoint: tuple[float, float],
43 max_size: int = 500,
44 dilation_iterations: int = 3,
45) -> RGBImage:
46 """Heal out-of-gamut regions in a Lab image using diffusion-based inpainting.
48 Parameters
49 ----------
50 lab_image : `LABImage`
51 A NxMx3 array in the Lab colorspace. Modified in-place by healing out-of-gamut regions.
52 For each region, a copy is made before modification, then written back.
53 mask : `numpy.ndarray` of `bool`
54 A boolean mask indicating which pixels are out of gamut.
55 xyz_whitepoint : `tuple` of `float`, `float`
56 Sets the white point of the xyz colorspace in xy coordinates.
57 max_size : `int`, optional
58 Maximum size of regions to heal. Larger regions are skipped. Default is 500.
59 dilation_iterations : `int`, optional
60 Number of iterations for mask dilation. Default is 3.
62 Returns
63 -------
64 result : `RGBImage`
65 The healed image converted to RGB colorspace.
67 Raises
68 ------
69 ValueError
70 Raised if the shapes of lab_image and mask are incompatible.
72 Notes
73 -----
74 The healing algorithm works by:
76 1. Labeling connected regions in the mask
77 2. For each region smaller than max_size:
79 - Dilating the mask to create an annulus around the region
80 - Computing average a,b color values from the annulus
81 - Using rgb.inpaint_mask to interpolate the L, a, and b channels
82 - Filling the masked region with the interpolated values
84 3. Regions larger than max_size are skipped.
85 """
86 # Need to split all the regions of the mask
87 labels = label(binary_dilation(mask, iterations=3))[0]
88 places = find_objects(labels)
89 # then grow the slices by 20% of the max size
90 new_places = []
91 for place in places:
92 size = int(3 * np.min([sl.stop - sl.start for sl in place]))
93 new_y = slice(np.max((0, place[0].start - size)), np.min((mask.shape[0], place[0].stop + size)), None)
94 new_x = slice(np.max((0, place[1].start - size)), np.min((mask.shape[1], place[1].stop + size)), None)
95 new_places.append((new_y, new_x))
96 # for each slice, dilate the mask by n-pixels, and then diff the mask to make anulus
97 # get the color data in the anulus
98 # set the values in the mask for the a,b channels to the color data
99 # use the rgb diffusion to fix the lum channel
100 # assign the fixed image back into the
101 for (place_y, place_x), (old_y, old_x) in zip(new_places, places):
102 sub_labels = labels[old_y, old_x]
103 if np.sum(sub_labels > 0) >= max_size:
104 continue
105 label_number = np.max(sub_labels)
106 # copy to ensure contiguous array, this is faster than operating on view
107 sub_mask = labels[place_y, place_x] == label_number
108 sub_lab = np.copy(lab_image[place_y, place_x])
109 new_lum = rgb.inpaint_mask(
110 np.ascontiguousarray(sub_lab[..., 0]),
111 sub_mask,
112 iterations=62,
113 radius=26,
114 radius_center=5,
115 )
116 new_a = rgb.inpaint_mask(
117 np.ascontiguousarray(sub_lab[..., 1]),
118 sub_mask,
119 iterations=62,
120 radius=26,
121 radius_center=5,
122 )
123 new_b = rgb.inpaint_mask(
124 np.ascontiguousarray(sub_lab[..., 2]),
125 sub_mask,
126 iterations=62,
127 radius=26,
128 radius_center=5,
129 )
130 sub_lab[..., 0] = new_lum
131 sub_lab[..., 1] = new_a
132 sub_lab[..., 2] = new_b
134 lab_image[place_y, place_x] = sub_lab
136 return rgb.Oklab_to_RGB(np.ascontiguousarray(lab_image), xyz_whitepoint)
139class GamutFixer(ConfigurableAction):
140 """Fix out-of-gamut colors in images by remapping them back into the RGB gamut.
142 This class provides multiple methods for handling colors that fall outside
143 the standard RGB color gamut, including inpainting, mapping functions,
144 and diffusion-based healing.
145 """
147 gamutMethod = ChoiceField[str](
148 doc="If doRemapGamut is True this determines the method",
149 default="inpaint",
150 allowed={
151 "mapping": "Use a mapping function",
152 "inpaint": "Use surrounding pixels to determine likely value",
153 "heal": "Heal regions with reverse diffusion",
154 "none": "Don't fix out of gamut colors",
155 },
156 )
157 max_size = Field[int](doc="Maximum number of contiguous pixels that will be fixed", default=10000)
159 def __call__(self, Lab: LABImage, xyz_whitepoint: tuple[float, float]) -> RGBImage:
160 """Remap colors that fall outside an RGB color gamut back into it.
162 When gamutMethod is 'heal', regions larger than the configured max_size
163 are skipped.
165 Parameters
166 ----------
167 Lab : `LABImage`
168 A NxMx3 array in the Lab colorspace containing the image data.
169 xyz_whitepoint : `tuple` of `float`, `float`
170 Sets the white point of the xyz colorspace in xy coordinates.
172 Returns
173 -------
174 result : `RGBImage`
175 The image with out-of-gamut colors remapped to RGB colorspace.
177 Raises
178 ------
179 ValueError
180 Raise if the gamut method is not one of the supported methods.
181 """
182 rgb_prime = rgb.Oklab_to_RGB(np.ascontiguousarray(Lab), xyz_whitepoint)
184 if self.gamutMethod == "none":
185 return rgb_prime
187 # Determine if there are any out of bounds pixels
188 outOfBounds = np.bitwise_or(
189 np.bitwise_or(rgb_prime[:, :, 0] > 1, rgb_prime[:, :, 0] < 0),
190 np.bitwise_or(rgb_prime[:, :, 1] > 1, rgb_prime[:, :, 1] < 0),
191 )
192 outOfBounds = np.bitwise_or(
193 outOfBounds, np.bitwise_or(rgb_prime[:, :, 2] > 1, rgb_prime[:, :, 2] < 0)
194 )
196 # If all pixels are in bounds, return immediately.
197 if not np.any(outOfBounds):
198 logging.info("There are no out of gamut pixels.")
199 return rgb_prime
201 logging.info("There are out of gamut pixels, remapping colors")
202 match self.gamutMethod:
203 case "inpaint":
204 logging.debug("Running inpaint")
205 labels, num_features = label(outOfBounds)
206 label_counts = np.bincount(labels.ravel())[0:]
207 for index, amount in zip(range(num_features), label_counts):
208 if amount > self.max_size:
209 logging.debug(f"Eliminating {amount} pixels")
210 # ignore areas that are too large
211 outOfBounds[labels == index] = 0
213 results = skimage.restoration.inpaint_biharmonic(rgb_prime, outOfBounds, channel_axis=-1)
214 case "mapping":
215 results = fixGamutOK(Lab[outOfBounds])
216 Lab[outOfBounds] = results
217 results = rgb.Oklab_to_RGB(np.ascontiguousarray(Lab), xyz_whitepoint)
218 case "heal":
219 results = heal_gamut(Lab, outOfBounds, xyz_whitepoint)
220 case _:
221 raise ValueError(f"gamut correction {self.gamutMethod} is not supported")
223 logging.debug(f"The total number of remapped pixels is: {np.sum(outOfBounds)}")
224 return results