Initial port/rewrite from afw_extensions_rgb

With changes from code reviews squashed in

python/lsst/afw/display/Rgb.h

@@ -0,0 +1,17 @@
+#if !defined(LSST_AFW_DISPLAY_RGB_H)
+#define LSST_AFW_DISPLAY_RGB_H 1
+
+namespace lsst { namespace afw { namespace display {
+
+template<typename ImageT>
+void
+replaceSaturatedPixels(ImageT & rim,    //< R image (e.g. i)
+                       ImageT & gim,    //< G image (e.g. r)
+                       ImageT & bim,    //< B image (e.g. g)
+                       int borderWidth = 2, //< width of border used to estimate colour of saturated regions
+                       float saturatedPixelValue = 65535 //< the brightness of a saturated pixel, once fixed
+                      );
+
+}}}
+
+#endif

python/lsst/afw/display/SConscript

@@ -1,3 +1,4 @@
 # -*- python -*-
 from lsst.sconsUtils import scripts
-scripts.BasicSConscript.python(['displayLib', 'xpa'], swigSrc={"displayLib": ["simpleFits.cc"]})
+scripts.BasicSConscript.python(['displayLib', 'xpa'],
+                               swigSrc={"displayLib": ["saturated.cc", "simpleFits.cc"]})

python/lsst/afw/display/displayLib.i

@@ -41,6 +41,7 @@ Basic routines to talk to ds9
 #   include "lsst/afw/image.h"
 
 #   include "simpleFits.h"
+#   include "Rgb.h"
 %}
 
 %include "lsst/p_lsstSwig.i"
@@ -58,3 +59,7 @@ Basic routines to talk to ds9
 %template(writeFitsImage) lsst::afw::display::writeBasicFits<lsst::afw::image::Image<double> >;
 %template(writeFitsImage) lsst::afw::display::writeBasicFits<lsst::afw::image::Mask<boost::uint16_t> >;
 
+%include "Rgb.h"
+
+%template(replaceSaturatedPixels)
+        lsst::afw::display::replaceSaturatedPixels<lsst::afw::image::MaskedImage<float> >;

python/lsst/afw/display/rgb.py

@@ -0,0 +1,145 @@
+import numpy as np
+
+import lsst.pex.exceptions as pexExceptions
+import lsst.afw.detection as afwDetect
+import lsst.afw.image as afwImage
+import lsst.afw.geom as afwGeom
+import lsst.afw.math as afwMath
+import lsst.afw.display.ds9 as ds9
+from lsst.afw.display.displayLib import replaceSaturatedPixels
+
+class Mapping(object):
+    """!Baseclass to map red, blue, green intensities into uint8 values"""
+
+    def __init__(self, min):
+        self._uint8Max = float(np.iinfo(np.uint8).max)
+
+        try:
+            len(min)
+        except:
+            min = 3*[min]
+        assert len(min) == 3, "Please provide 1 or 3 values for min"
+
+        self._min = min
+
+    def makeRgbImage(self, imageR, imageG, imageB):
+        """!Convert 3 arrays, imageR, imageG, and imageB into a numpy RGB image
+
+N.b. images may be afwImages or numpy arrays
+        """
+        imageRGB = [imageR, imageG, imageB]
+        for i, c in enumerate(imageRGB):
+            if hasattr(c, "getArray"):
+                imageRGB[i] = c.getArray()
+
+        return np.flipud(np.dstack(self._convertImagesToUint8(*imageRGB)).astype(np.uint8))
+
+    def intensity(self, imageR, imageG, imageB):
+        """!Return the total intensity from the red, blue, and green intensities"""
+        return (imageR + imageG + imageB)/float(3);
+
+    def mapIntensityToUint8(self, I):
+        """Map an intensity into the range of a uint8, [0, 255] (but not converted to uint8)"""
+        return np.where(I <= 0, 0, np.where(I < self._uint8Max, I, self._uint8Max))
+
+    def _convertImagesToUint8(self, imageR, imageG, imageB):
+        """Use the mapping to convert images imageR, imageG, and imageB to a triplet of uint8 images"""
+        imageR = imageR - self._min[0]  # n.b. makes copy
+        imageG = imageG - self._min[1]
+        imageB = imageB - self._min[2]
+
+        fac = self.mapIntensityToUint8(self.intensity(imageR, imageG, imageB))
+
+        imageRGB = [imageR, imageG, imageB]
+
+        for c in imageRGB:
+            c *= fac
+            c[c <= 0] = 0
+
+        pixmax = self._uint8Max
+        r0, g0, b0 = imageRGB           # copies -- could work row by row to minimise memory usage
+
+        with np.errstate(invalid='ignore', divide='ignore'): # n.b. np.where doesn't (and can't) short-circuit
+            for i, c in enumerate(imageRGB):
+                c = np.where(r0 > g0,
+                             np.where(r0 > b0,
+                                      np.where(r0 >= pixmax, c*pixmax/r0, c),
+                                      np.where(b0 >= pixmax, c*pixmax/b0, c)),
+                             np.where(g0 > b0,
+                                      np.where(g0 >= pixmax, c*pixmax/g0, c),
+                                      np.where(b0 >= pixmax, c*pixmax/b0, c))).astype(np.uint8)
+                c[c > pixmax] = pixmax
+
+                imageRGB[i] = c
+
+        return imageRGB
+
+class AsinhMapping(Mapping):
+    """!A mapping for an asinh stretch (preserving colours independent of brightness)
+
+x = asinh(Q (I - min)/range)/Q
+
+This reduces to a linear stretch if Q == 0
+"""
+
+    def __init__(self, min, range, Q):
+        Mapping.__init__(self, min)
+
+        epsilon = 1.0/2**23            # 32bit floating point machine epsilon; sys.float_info.epsilon is 64bit
+        if abs(Q) < epsilon:
+            Q = 0.1
+        else:
+            Qmax = 1e10
+            if Q > Qmax:
+                Q = Qmax
+
+        if False:
+            self._slope = self._uint8Max/Q # gradient at origin is self._slope
+        else:
+            frac = 0.1                  # gradient estimated using frac*range is _slope
+            self._slope = frac*self._uint8Max/np.arcsinh(frac*Q)
+
+        self._soften = Q/float(range);
+
+    def mapIntensityToUint8(self, I):
+        return np.where(I <= 0, 0, np.arcsinh(I*self._soften)*self._slope/I)
+
+def makeRGB(imageR, imageG, imageB, min=0, range=5, Q=20, fileName=None):
+    """Make a set of three images into an RGB image using an asinh stretch and optionally write it to disk"""
+    asinhMap = AsinhMapping(min, range, Q)
+    rgb = asinhMap.makeRgbImage(imageR, imageG, imageB)
+    if fileName:
+        writeRGB(fileName, rgb)
+
+    return rgb
+
+def displayRGB(rgb):
+    """Display an rgb image using matplotlib"""
+    import matplotlib.pyplot as plt
+    plt.imshow(rgb, interpolation='nearest')
+    plt.show()
+
+def writeRGB(fileName, rgbImage):
+    import matplotlib.image
+    matplotlib.image.imsave(fileName, rgbImage)
+
+#-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
+#
+# Support the legacy API
+#
+class asinhMappingF(object):
+    def __init__(self, min, range, Q):
+        self.min = min
+        self.range = range
+        self.Q = Q
+
+class _RgbImageF(object):
+    def __init__(self, imageR, imageG, imageB, mapping):
+        asinh = AsinhMapping(mapping.min, mapping.range, mapping.Q)
+        self.rgb = asinh.makeRgbImage(imageR, imageG, imageB)
+
+    def write(self, fileName):
+        writeRGB(fileName, self.rgb)
+
+def RgbImageF(imageR, imageG, imageB, mapping):
+    return _RgbImageF(imageR, imageG, imageB, mapping)

python/lsst/afw/display/saturated.cc

@@ -0,0 +1,175 @@
+/**
+ * \file
+ *
+ * Handle saturated pixels when making colour images
+ */
+#include "boost/format.hpp"
+#include "lsst/utils/ieee.h"
+#include "lsst/afw/detection.h"
+#include "lsst/afw/image/MaskedImage.h"
+#include "Rgb.h"
+
+namespace lsst { namespace afw { namespace display {
+
+namespace {
+    template <typename ImageT>
+    class SetPixels : public detection::FootprintFunctor<ImageT> {
+    public:
+        explicit SetPixels(ImageT const& img     // The image the source lives in
+                          ) : detection::FootprintFunctor<ImageT>(img), _value(0) {}
+
+        void setValue(float value) { _value = value; }
+
+        // method called for each pixel by apply()
+        void operator()(typename ImageT::xy_locator loc,        // locator pointing at the pixel
+                        int,                                    // column-position of pixel
+                        int                                     // row-position of pixel
+                       ) {
+            *loc = _value;
+        }
+    private:
+        float _value;
+    };
+}
+
+template<typename ImageT>
+void
+replaceSaturatedPixels(ImageT & rim,    // R image (e.g. i)
+                       ImageT & gim,    // G image (e.g. r)
+                       ImageT & bim,    // B image (e.g. g)
+                       int borderWidth,	// width of border used to estimate colour of saturated regions
+                       float saturatedPixelValue // the brightness of a saturated pixel, once fixed
+                      )
+{
+    int const width = rim.getWidth(), height = rim.getHeight();
+    int const x0 = rim.getX0(), y0 = rim.getY0();
+
+    if (width != gim.getWidth() || height != gim.getHeight() || x0 != gim.getX0() || y0 != gim.getY0()) {
+        throw LSST_EXCEPT(pex::exceptions::InvalidParameterError,
+                          str(boost::format("R image has different size/origin from G image "
+                                            "(%dx%d+%d+%d v. %dx%d+%d+%d") %
+                              width % height % x0 % y0 %
+                              gim.getWidth() % gim.getHeight() % gim.getX0() % gim.getY0()));
+
+    }
+    if (width != bim.getWidth() || height != bim.getHeight() || x0 != bim.getX0() || y0 != bim.getY0()) {
+        throw LSST_EXCEPT(pex::exceptions::InvalidParameterError,
+                          str(boost::format("R image has different size/origin from B image "
+                                            "(%dx%d+%d+%d v. %dx%d+%d+%d") %
+                              width % height % x0 % y0 %
+                              bim.getWidth() % bim.getHeight() % bim.getX0() % bim.getY0()));
+
+    }
+
+    bool const useMaxPixel = !utils::isfinite(saturatedPixelValue);
+
+    SetPixels<typename ImageT::Image>
+        setR(*rim.getImage()),
+        setG(*gim.getImage()),
+        setB(*bim.getImage()); // functors used to set pixel values
+
+    // Find all the saturated pixels in any of the three image
+    int const npixMin = 1;              // minimum number of pixels in an object
+    afw::image::MaskPixel const SAT = rim.getMask()->getPlaneBitMask("SAT");
+    detection::Threshold const satThresh(SAT, detection::Threshold::BITMASK);
+
+    detection::FootprintSet       sat(*rim.getMask(), satThresh, npixMin);
+    sat.merge(detection::FootprintSet(*gim.getMask(), satThresh, npixMin));
+    sat.merge(detection::FootprintSet(*bim.getMask(), satThresh, npixMin));
+    // go through the list of saturated regions, determining the mean colour of the surrounding pixels
+    typedef detection::FootprintSet::FootprintList FootprintList;
+    PTR(FootprintList) feet = sat.getFootprints();
+    for (FootprintList::const_iterator ptr = feet->begin(), end = feet->end(); ptr != end; ++ptr) {
+        PTR(detection::Footprint) const foot = *ptr;
+        PTR(detection::Footprint) const bigFoot = growFootprint(*foot, borderWidth);
+
+        double sumR = 0, sumG = 0, sumB = 0; // sum of all non-saturated adjoining pixels
+        double maxR = 0, maxG = 0, maxB = 0; // maximum of non-saturated adjoining pixels
+
+        for (detection::Footprint::SpanList::const_iterator sptr = bigFoot->getSpans().begin(),
+                 send = bigFoot->getSpans().end(); sptr != send; ++sptr) {
+            PTR(detection::Span) const span = *sptr;
+
+            int const y = span->getY() - y0;
+            if (y < 0 || y >= height) {
+                continue;
+            }
+            int sx0 = span->getX0() - x0;
+            if (sx0 < 0) {
+                sx0 = 0;
+            }
+            int sx1 = span->getX1() - x0;
+            if (sx1 >= width) {
+                sx1 = width - 1;
+            }
+
+            for (typename ImageT::iterator
+                     rptr = rim.at(sx0, y),
+                     rend = rim.at(sx1 + 1, y),
+                     gptr = gim.at(sx0, y),
+                     bptr = bim.at(sx0, y); rptr != rend; ++rptr, ++gptr, ++bptr) {
+                if (!((rptr.mask() | gptr.mask() | bptr.mask()) & SAT)) {
+                    float val = rptr.image();
+                    sumR += val;
+                    if (val > maxR) {
+                        maxR = val;
+                    }
+
+                    val = gptr.image();
+                    sumG += val;
+                    if (val > maxG) {
+                        maxG = val;
+                    }
+
+                    val = bptr.image();
+                    sumB += val;
+                    if (val > maxB) {
+                        maxB = val;
+                    }
+                }
+            }
+        }
+        // OK, we have the mean fluxes for the pixels surrounding this set of saturated pixels
+        // so we can figure out the proper values to use for the saturated ones
+        float R = 0, G = 0, B = 0;      // mean intensities
+        if (sumR + sumB + sumG > 0) {
+            if (sumR > sumG) {
+                if (sumR > sumB) {
+                    R = useMaxPixel ? maxR : saturatedPixelValue;
+
+                    G = (R*sumG)/sumR;
+                    B = (R*sumB)/sumR;
+                } else {
+                    B = useMaxPixel ? maxB : saturatedPixelValue;
+                    R = (B*sumR)/sumB;
+                    G = (B*sumG)/sumB;
+                }
+            } else {
+                if (sumG > sumB) {
+                    G = useMaxPixel ? maxG : saturatedPixelValue;
+                    R = (G*sumR)/sumG;
+                    B = (G*sumB)/sumG;
+                } else {
+                    B = useMaxPixel ? maxB : saturatedPixelValue;
+                    R = (B*sumR)/sumB;
+                    G = (B*sumG)/sumB;
+                }
+            }
+        }
+        // Now that we know R, G, and B we can fix the values
+        setR.setValue(R); setR.apply(*foot);
+        setG.setValue(G); setG.apply(*foot);
+        setB.setValue(B); setB.apply(*foot);
+    }
+}
+
+template
+void
+replaceSaturatedPixels(image::MaskedImage<float> & rim,
+                       image::MaskedImage<float> & gim,
+                       image::MaskedImage<float> & bim,
+                       int borderWidth,
+                       float saturatedPixelValue
+                      );
+
+}}}