Coverage for tests/test_rgb.py: 25%
186 statements
« prev ^ index » next coverage.py v6.4.4, created at 2022-09-15 02:49 -0700
« prev ^ index » next coverage.py v6.4.4, created at 2022-09-15 02:49 -0700
1# This file is part of afw.
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/>.
22import os
23import math
24import unittest
26import numpy as np
28import lsst.utils.tests
29import lsst.geom
30import lsst.afw.detection as afwDetect
31import lsst.afw.image as afwImage
32import lsst.afw.math as afwMath
33import lsst.afw.display as afwDisplay
34import lsst.afw.display.rgb as rgb
36ver1, ver2, ver3 = 1, 3, 1
37NO_MATPLOTLIB_STRING = "Requires matplotlib >= %d.%d.%d" % (ver1, ver2, ver3)
38try:
39 import matplotlib
40 mplVersion = matplotlib.__version__
41 # Split at + to check for development version (PEP 440)
42 mplVersion = mplVersion.split('+')
43 versionInfo = tuple(int(s.strip("rc")) for s in mplVersion[0].split("."))
44 HAVE_MATPLOTLIB = versionInfo >= (ver1, ver2, ver3)
45except ImportError:
46 HAVE_MATPLOTLIB = False
48try:
49 import scipy.misc
50 scipy.misc.imresize
51 HAVE_SCIPY_MISC = True
52except (ImportError, AttributeError):
53 HAVE_SCIPY_MISC = False
55try:
56 type(display)
57 afwDisplay.setDefaultMaskTransparency(75)
58except NameError:
59 display = False
62def saturate(image, satValue):
63 """Simulate saturation on an image, so we can test 'replaceSaturatedPixels'
65 Takes an Image, sets saturated pixels to NAN and masks them, returning
66 a MaskedImage.
67 """
68 image = afwImage.makeMaskedImage(image)
69 afwDetect.FootprintSet(image, afwDetect.Threshold(satValue), "SAT")
70 arr = image.getImage().getArray()
71 arr[np.where(arr >= satValue)] = np.nan
72 return image
75R, G, B = 2, 1, 0
78class RgbTestCase(unittest.TestCase):
79 """A test case for Rgb"""
81 def setUp(self):
82 self.min, self.range, self.Q = 0, 5, 20 # asinh
84 width, height = 85, 75
85 self.images = []
86 self.images.append(afwImage.ImageF(lsst.geom.ExtentI(width, height)))
87 self.images.append(afwImage.ImageF(lsst.geom.ExtentI(width, height)))
88 self.images.append(afwImage.ImageF(lsst.geom.ExtentI(width, height)))
90 for (x, y, A, g_r, r_i) in [(15, 15, 1000, 1.0, 2.0),
91 (50, 45, 5500, -1.0, -0.5),
92 (30, 30, 600, 1.0, 2.5),
93 (45, 15, 20000, 1.0, 1.0),
94 ]:
95 for i in range(len(self.images)):
96 if i == B:
97 amp = A
98 elif i == G:
99 amp = A*math.pow(10, 0.4*g_r)
100 elif i == R:
101 amp = A*math.pow(10, 0.4*r_i)
103 self.images[i][x, y, afwImage.LOCAL] = amp
105 psf = afwMath.AnalyticKernel(
106 15, 15, afwMath.GaussianFunction2D(2.5, 1.5, 0.5))
108 convolvedImage = type(self.images[0])(self.images[0].getDimensions())
109 randomImage = type(self.images[0])(self.images[0].getDimensions())
110 rand = afwMath.Random("MT19937", 666)
111 convolutionControl = afwMath.ConvolutionControl()
112 convolutionControl.setDoNormalize(True)
113 convolutionControl.setDoCopyEdge(True)
114 for i in range(len(self.images)):
115 afwMath.convolve(convolvedImage, self.images[i], psf, convolutionControl)
116 afwMath.randomGaussianImage(randomImage, rand)
117 randomImage *= 2
118 convolvedImage += randomImage
119 self.images[i][:] = convolvedImage
120 del convolvedImage
121 del randomImage
123 def tearDown(self):
124 for im in self.images:
125 del im
126 del self.images
128 def testStarsAsinh(self):
129 """Test creating an RGB image using an asinh stretch"""
130 asinhMap = rgb.AsinhMapping(self.min, self.range, self.Q)
131 rgbImage = asinhMap.makeRgbImage(
132 self.images[R], self.images[G], self.images[B])
134 if display:
135 rgb.displayRGB(rgbImage)
137 def testStarsAsinhZscale(self):
138 """Test creating an RGB image using an asinh stretch estimated using zscale"""
140 rgbImages = [self.images[R], self.images[G], self.images[B]]
142 map = rgb.AsinhZScaleMapping(rgbImages[0])
143 rgbImage = map.makeRgbImage(*rgbImages)
145 if display:
146 rgb.displayRGB(rgbImage)
148 def testStarsAsinhZscaleIntensity(self):
149 """Test creating an RGB image using an asinh stretch estimated using zscale on the intensity"""
151 rgbImages = [self.images[R], self.images[G], self.images[B]]
153 map = rgb.AsinhZScaleMapping(rgbImages)
154 rgbImage = map.makeRgbImage(*rgbImages)
156 if display:
157 rgb.displayRGB(rgbImage)
159 def testStarsAsinhZscaleIntensityPedestal(self):
160 """Test creating an RGB image using an asinh stretch estimated using zscale on the intensity
161 where the images each have a pedestal added"""
163 rgbImages = [self.images[R], self.images[G], self.images[B]]
165 pedestal = [100, 400, -400]
166 for i, ped in enumerate(pedestal):
167 rgbImages[i] += ped
169 map = rgb.AsinhZScaleMapping(rgbImages, pedestal=pedestal)
170 rgbImage = map.makeRgbImage(*rgbImages)
172 if display:
173 rgb.displayRGB(rgbImage)
175 def testStarsAsinhZscaleIntensityBW(self):
176 """Test creating a black-and-white image using an asinh stretch estimated
177 using zscale on the intensity"""
179 rgbImage = rgb.AsinhZScaleMapping(self.images[R]).makeRgbImage()
181 if display:
182 rgb.displayRGB(rgbImage)
184 @unittest.skipUnless(HAVE_MATPLOTLIB, NO_MATPLOTLIB_STRING)
185 def testMakeRGB(self):
186 """Test the function that does it all"""
187 satValue = 1000.0
188 with lsst.utils.tests.getTempFilePath(".png") as fileName:
189 red = saturate(self.images[R], satValue)
190 green = saturate(self.images[G], satValue)
191 blue = saturate(self.images[B], satValue)
192 rgb.makeRGB(red, green, blue, self.min, self.range, self.Q, fileName=fileName,
193 saturatedBorderWidth=1, saturatedPixelValue=2000)
194 self.assertTrue(os.path.exists(fileName))
196 def testLinear(self):
197 """Test using a specified linear stretch"""
199 rgbImage = rgb.LinearMapping(-8.45, 13.44).makeRgbImage(self.images[R])
201 if display:
202 rgb.displayRGB(rgbImage)
204 def testLinearMinMax(self):
205 """Test using a min/max linear stretch
207 N.b. also checks that an image passed to the ctor is used as the default in makeRgbImage()
208 """
210 rgbImage = rgb.LinearMapping(image=self.images[R]).makeRgbImage()
212 if display:
213 rgb.displayRGB(rgbImage)
215 def testZScale(self):
216 """Test using a zscale stretch"""
218 rgbImage = rgb.ZScaleMapping(self.images[R]).makeRgbImage()
220 if display:
221 plt = rgb.displayRGB(rgbImage, False)
222 plt.title("zscale")
223 plt.show()
225 @unittest.skipUnless(HAVE_MATPLOTLIB, NO_MATPLOTLIB_STRING)
226 def testWriteStars(self):
227 """Test writing RGB files to disk"""
228 asinhMap = rgb.AsinhMapping(self.min, self.range, self.Q)
229 rgbImage = asinhMap.makeRgbImage(
230 self.images[R], self.images[G], self.images[B])
231 with lsst.utils.tests.getTempFilePath(".png") as fileName:
232 rgb.writeRGB(fileName, rgbImage)
233 self.assertTrue(os.path.exists(fileName))
235 def testSaturated(self):
236 """Test interpolating saturated pixels"""
238 satValue = 1000.0
239 for f in [R, G, B]:
240 self.images[f] = saturate(self.images[f], satValue)
242 rgb.replaceSaturatedPixels(
243 self.images[R], self.images[G], self.images[B], 1, 2000)
244 #
245 # Check that we replaced those NaNs with some reasonable value
246 #
247 for f in [R, G, B]:
248 self.assertTrue(np.isfinite(
249 self.images[f].getImage().getArray()).all())
251 if display > 1:
252 afwDisplay.Display(frame=0).mtv(self.images[B], title="B: SAT-interpolated")
253 afwDisplay.Display(frame=1).mtv(self.images[G], title="G: SAT-interpolated")
254 afwDisplay.Display(frame=2).mtv(self.images[R], title="R: SAT-interpolated")
255 #
256 # Prepare for generating an output file
257 #
258 for f in [R, G, B]:
259 self.images[f] = self.images[f].getImage()
261 asinhMap = rgb.AsinhMapping(self.min, self.range, self.Q)
262 rgbImage = asinhMap.makeRgbImage(
263 self.images[R], self.images[G], self.images[B])
265 if display:
266 rgb.displayRGB(rgbImage)
268 @unittest.skipUnless(HAVE_SCIPY_MISC, "Resizing images requires scipy.misc")
269 def testStarsResizeToSize(self):
270 """Test creating an RGB image of a specified size"""
272 xSize = self.images[R].getWidth()//2
273 ySize = self.images[R].getHeight()//2
274 for rgbImages in ([self.images[R], self.images[G], self.images[B]],
275 [afwImage.ImageU(_.getArray().astype('uint16')) for _ in [
276 self.images[R], self.images[G], self.images[B]]]):
277 rgbImage = rgb.AsinhZScaleMapping(rgbImages[0]).makeRgbImage(*rgbImages,
278 xSize=xSize, ySize=ySize)
280 if display:
281 rgb.displayRGB(rgbImage)
283 @unittest.skipUnless(HAVE_SCIPY_MISC, "Resizing images requires scipy.misc")
284 def testStarsResizeSpecifications(self):
285 """Test creating an RGB image changing the output """
287 rgbImages = [self.images[R], self.images[G], self.images[B]]
288 map = rgb.AsinhZScaleMapping(rgbImages[0])
290 for xSize, ySize, frac in [(self.images[R].getWidth()//2, self.images[R].getHeight()//2, None),
291 (2*self.images[R].getWidth(), None, None),
292 (self.images[R].getWidth()//2, None, None),
293 (None, self.images[R].getHeight()//2, None),
294 (None, None, 0.5),
295 (None, None, 2),
296 ]:
297 rgbImage = map.makeRgbImage(
298 *rgbImages, xSize=xSize, ySize=ySize, rescaleFactor=frac)
300 h, w = rgbImage.shape[0:2]
301 self.assertTrue(xSize is None or xSize == w)
302 self.assertTrue(ySize is None or ySize == h)
303 self.assertTrue(frac is None or w == int(frac*self.images[R].getWidth()),
304 "%g == %g" % (w, int((frac if frac else 1)*self.images[R].getWidth())))
306 if display:
307 rgb.displayRGB(rgbImage)
309 @unittest.skipUnless(HAVE_SCIPY_MISC, "Resizing images requires scipy.misc")
310 @unittest.skipUnless(HAVE_MATPLOTLIB, NO_MATPLOTLIB_STRING)
311 def testMakeRGBResize(self):
312 """Test the function that does it all, including rescaling"""
313 rgb.makeRGB(self.images[R], self.images[G],
314 self.images[B], xSize=40, ySize=60)
316 with lsst.utils.tests.getTempFilePath(".png") as fileName:
317 rgb.makeRGB(self.images[R], self.images[G],
318 self.images[B], fileName=fileName, rescaleFactor=0.5)
319 self.assertTrue(os.path.exists(fileName))
322class TestMemory(lsst.utils.tests.MemoryTestCase):
323 pass
326def setup_module(module):
327 lsst.utils.tests.init()
330if __name__ == "__main__": 330 ↛ 331line 330 didn't jump to line 331, because the condition on line 330 was never true
331 lsst.utils.tests.init()
332 unittest.main()