fix(fits): read/write color FITS image using the correct convention

src/doc/builtinplugins.rst

@@ -522,10 +522,33 @@ found at: http://fits.gsfc.nasa.gov/
 OpenImageIO supports multiple images in FITS files, and supports the
 following pixel data types: UINT8, UINT16, UINT32, FLOAT, DOUBLE.
 
-FITS files can store various kinds of arbitrary data arrays, but
-OpenImageIO's support of FITS is mostly limited using FITS for image
-storage.  Currently, OpenImageIO only supports 2D FITS data (images), not 3D
-(volume) data, nor 1-D or higher-dimensional arrays.
+**Inferring dimensions and color channels**
+
+FITS files can store arbitrary N-dimensional data arrays (``NAXIS`` /
+``NAXISn``), but say nothing about which axes, if any, are "color."
+OpenImageIO infers a role for each axis from ``NAXIS`` (trailing axes of
+length 1 are dropped first):
+
+* ``NAXIS`` = 1: a single row (width = NAXIS1, height = 1).
+* ``NAXIS`` = 2: an ordinary 2D grayscale image (width = NAXIS1, height = NAXIS2).
+* ``NAXIS`` = 3: if NAXIS3 <= 4, a color image, stored as one full-resolution
+  plane per channel (width = NAXIS1, height = NAXIS2, NAXIS3 = nchannels).
+  Otherwise, if NAXIS3 is more than 3, the file is interpreted as a grayscale
+  volume (width = NAXIS1, height = NAXIS2, depth = NAXIS3).
+* ``NAXIS`` = 4: a color volume (width = NAXIS1, height = NAXIS2, depth =
+  NAXIS3, NAXIS4 = nchannels, again one full-resolution volume per channel).
+
+This is only a heuristic -- FITS has no way to formally declare an axis as
+"color" -- so a genuinely non-color NAXIS3 that happens to be <= 4 long can be
+misclassified as color.
+
+Channel names default to the usual OpenImageIO conventions (``R,G,B``,
+``R,G,B,A``, ``Y``, etc.), improved when the header says more: if the
+channel axis is a FITS WCS ``STOKES`` axis, channels are named from the
+standard Stokes codes (``I``, ``Q``, ``U``, ``V``, ``RR``, ``LL``, ``RL``,
+``LR``, ``pp``, ``qq``, ``pq``, ``qp``); otherwise, if every plane has a
+``FILTERn`` or ``BANDn`` keyword (an informal convention, not part of the
+FITS standard), those values are used instead.
 
 
 

src/fits.imageio/fits_pvt.h

@@ -114,6 +114,21 @@ class FitsInput final : public ImageInput {
     // converts date in FITS format (YYYY-MM-DD or DD/MM/YY)
     // to DateTime format
     std::string convert_date(const std::string& date);
+
+    // Give m_spec.channelnames sensible values: generic defaults (R,G,B,...),
+    // improved with real per-plane names when the header tells us what the
+    // channel axis actually is (a STOKES polarization axis, or informal
+    // FILTERn/BANDn keywords).
+    void assign_channel_names();
+
+    // True if multi-channel pixel data is stored as separate per-channel
+    // blocks -- one full width x height (x depth, for NAXIS=4) block per
+    // channel -- rather than interleaved. This is the only multi-channel
+    // layout we recognize; see the comments in read_fits_header().
+    bool planar_channels() const
+    {
+        return (m_naxes == 3 || m_naxes == 4) && m_spec.nchannels > 1;
+    }
 };
 
 

src/fits.imageio/fitsinput.cpp

@@ -4,11 +4,13 @@
 
 
 #include <cctype>
+#include <cmath>
 #include <cstdlib>
 
 #include "fits_pvt.h"
 
 #include <OpenImageIO/fmath.h>
+#include <OpenImageIO/strutil.h>
 
 
 OIIO_PLUGIN_NAMESPACE_BEGIN
@@ -100,16 +102,46 @@ FitsInput::read_native_scanline(int subimage, int miplevel, int y, int /*z*/,
         return true;
 
     std::vector<unsigned char> data_tmp(m_spec.scanline_bytes());
-    long scanline_off = (m_spec.height - y) * m_spec.scanline_bytes();
-    fseek(m_fd, scanline_off, SEEK_CUR);
-    size_t n = fread(&data_tmp[0], 1, m_spec.scanline_bytes(), m_fd);
-    if (n != m_spec.scanline_bytes()) {
-        if (feof(m_fd))
-            errorfmt("Hit end of file unexpectedly (offset={}, scanline {})",
-                     ftell(m_fd), y);
-        else
-            errorfmt("read error");
-        return false;  // Read failed
+
+    if (!planar_channels()) {
+        long scanline_off = (m_spec.height - y) * m_spec.scanline_bytes();
+        fseek(m_fd, scanline_off, SEEK_CUR);
+        size_t n = fread(&data_tmp[0], 1, m_spec.scanline_bytes(), m_fd);
+        if (n != m_spec.scanline_bytes()) {
+            if (feof(m_fd))
+                errorfmt("Hit end of file unexpectedly (offset={}, scanline {})",
+                         ftell(m_fd), y);
+            else
+                errorfmt("read error");
+            return false;  // Read failed
+        }
+    } else {
+        // Channels are stored as separate contiguous width x height (x depth,
+        // for NAXIS=4) blocks rather than interleaved, so we must gather one
+        // row from each block and interleave them into the scanline buffer.
+        size_t comp_size   = m_spec.format.size();
+        size_t row_bytes   = size_t(m_spec.width) * comp_size;
+        size_t plane_bytes = row_bytes * size_t(m_spec.height)
+                             * size_t(m_spec.depth);
+        long row_off = (m_spec.height - y) * long(row_bytes);
+        std::vector<unsigned char> chan_row(row_bytes);
+        for (int c = 0; c < m_spec.nchannels; ++c) {
+            fsetpos(m_fd, &m_filepos);
+            fseek(m_fd, long(c * plane_bytes) + row_off, SEEK_CUR);
+            size_t n = fread(&chan_row[0], 1, row_bytes, m_fd);
+            if (n != row_bytes) {
+                if (feof(m_fd))
+                    errorfmt(
+                        "Hit end of file unexpectedly (offset={}, scanline {})",
+                        ftell(m_fd), y);
+                else
+                    errorfmt("read error");
+                return false;  // Read failed
+            }
+            for (int x = 0; x < m_spec.width; ++x)
+                memcpy(&data_tmp[(size_t(x) * m_spec.nchannels + c) * comp_size],
+                       &chan_row[size_t(x) * comp_size], comp_size);
+        }
     }
 
     // in FITS image data is stored in big-endian so we have to switch to
@@ -199,7 +231,7 @@ FitsInput::set_spec_info()
             errorfmt("Unsupported FITS BITPIX value {}", m_bitpix);
             return false;
         }
-        if (!check_open(m_spec, { 0, 1 << 20, 0, 1 << 20, 0, 1 << 16, 0, 4 }))
+        if (!check_open(m_spec, { 0, 1 << 20, 0, 1 << 20, 0, 1 << 16, 0, 12 }))
             return false;
         if (!check_compression_ratio(m_spec, Filesystem::file_size(m_filename)))
             return false;
@@ -344,22 +376,29 @@ FitsInput::read_fits_header(void)
         } else if (m_naxes == 2) {
             m_spec.width  = m_naxis[0];
             m_spec.height = m_naxis[1];
-        } else if (m_naxes == 3 && m_naxis[0] <= 4) {
-            // 3D, small number of most-rapidly changing dimension: color image?
-            m_spec.nchannels = m_naxis[0];
-            m_spec.width     = m_naxis[1];
-            m_spec.height    = m_naxis[2];
+        } else if (m_naxes == 3 && m_naxis[2] <= 4) {
+            // 3D, small number of most-slowly changing dimension: color
+            // image, one full width x height plane per channel (NAXIS3 =
+            // nchannels). This is the only real-world FITS color-cube
+            // convention -- see the git history for how we know that.
+            m_spec.width     = m_naxis[0];
+            m_spec.height    = m_naxis[1];
+            m_spec.nchannels = m_naxis[2];
         } else if (m_naxes == 3) {
-            // 3D, large number of most-rapidly changing dimension: volume?
+            // 3D, no small axis: volume
             m_spec.width  = m_naxis[0];
             m_spec.height = m_naxis[1];
             m_spec.depth  = m_naxis[2];
         } else if (m_naxes == 4) {
-            // 4D... volume + color?
-            m_spec.nchannels = m_naxis[0];
-            m_spec.width     = m_naxis[1];
-            m_spec.height    = m_naxis[2];
-            m_spec.depth     = m_naxis[3];
+            // 4D volume+color: NAXIS1-3 are x,y,z and NAXIS4 = nchannels,
+            // one full-resolution volume per channel. No real-world
+            // convention to confirm this against (unlike the 3D color case
+            // above), but it's the natural extension of the same pattern:
+            // the slowest-varying axis is the channel axis.
+            m_spec.width     = m_naxis[0];
+            m_spec.height    = m_naxis[1];
+            m_spec.depth     = m_naxis[2];
+            m_spec.nchannels = m_naxis[3];
         } else {
             errorfmt("Don't know now to read {}-channel FITS image", m_naxes);
             return false;
@@ -369,6 +408,7 @@ FitsInput::read_fits_header(void)
         m_spec.full_depth  = m_spec.depth;
 
         m_spec.attribute("oiio:subimages", (int)m_subimages.size());
+        assign_channel_names();
 
         // if (m_spec.width < 1 || m_spec.height < 1 || m_spec.depth < 1 ||
         //     m_spec.nchannels < 1) {
@@ -416,6 +456,89 @@ FitsInput::add_to_spec(const std::string& keyname, const std::string& value)
 
 
 
+void
+FitsInput::assign_channel_names()
+{
+    m_spec.default_channel_names();
+    if (m_spec.nchannels <= 1)
+        return;
+
+    // Which physical FITS axis (1-based) is the channel axis, so we know
+    // which CTYPEn/FILTERn-style keywords describe it?
+    int axis = 0;
+    if (m_naxes == 3)
+        axis = 3;
+    else if (m_naxes == 4)
+        axis = 4;  // volume+color: channel is the slowest-varying axis
+    if (!axis)
+        return;
+
+    // The FITS WCS standard defines a STOKES axis (polarization) with fixed
+    // integer codes per plane: 1=I, 2=Q, 3=U, 4=V, -1=RR, -2=LL, -3=RL,
+    // -4=LR, -5=pp, -6=qq, -7=pq, -8=qp. If that's what this axis is, decode
+    // the per-plane values via the standard linear WCS mapping and use the
+    // Stokes names as channel names.
+    std::string ctype = m_spec.get_string_attribute(
+        Strutil::format("Ctype{}", axis));
+    if (Strutil::iequals(ctype, "STOKES")) {
+        // Standard FITS WCS linear axis mapping and Stokes codes, see:
+        // https://fits.gsfc.nasa.gov/standard40/fits_standard40aa.pdf
+        static const std::map<int, const char*> stokes_names
+            = { { 1, "I" },   { 2, "Q" },   { 3, "U" },   { 4, "V" },
+                { -1, "RR" }, { -2, "LL" }, { -3, "RL" }, { -4, "LR" },
+                { -5, "pp" }, { -6, "qq" }, { -7, "pq" }, { -8, "qp" } };
+        float crval
+            = m_spec.get_float_attribute(Strutil::format("Crval{}", axis),
+                                         1.0f);
+        float crpix
+            = m_spec.get_float_attribute(Strutil::format("Crpix{}", axis),
+                                         1.0f);
+        float cdelt
+            = m_spec.get_float_attribute(Strutil::format("Cdelt{}", axis),
+                                         1.0f);
+        bool all_found = true;
+        std::vector<std::string> names(m_spec.nchannels);
+        for (int c = 0; c < m_spec.nchannels && all_found; ++c) {
+            int code  = int(std::lround(crval + (c + 1 - crpix) * cdelt));
+            auto find = stokes_names.find(code);
+            if (find == stokes_names.end())
+                all_found = false;
+            else
+                names[c] = find->second;
+        }
+        if (all_found) {
+            for (int c = 0; c < m_spec.nchannels; ++c)
+                m_spec.channelnames[c] = names[c];
+            return;
+        }
+    }
+
+    // No standard WCS description of the axis. Some pipelines instead tag
+    // each plane with an informal FILTERn or BANDn keyword (e.g. FILTER1,
+    // FILTER2, FILTER3 giving the name of the filter used for that plane).
+    // This isn't part of the FITS standard, just a convention we've seen in
+    // the wild, so only trust it if every plane has one.
+    for (const char* prefix : { "Filter", "Band" }) {
+        std::vector<std::string> names(m_spec.nchannels);
+        bool all_found = true;
+        for (int c = 0; c < m_spec.nchannels && all_found; ++c) {
+            std::string key = Strutil::format("{}{}", prefix, c + 1);
+            auto* param     = m_spec.find_attribute(key, TypeDesc::STRING);
+            if (!param)
+                all_found = false;
+            else
+                names[c] = *(const char**)param->data();
+        }
+        if (all_found) {
+            for (int c = 0; c < m_spec.nchannels; ++c)
+                m_spec.channelnames[c] = names[c];
+            return;
+        }
+    }
+}
+
+
+
 void
 FitsInput::subimage_search()
 {

src/fits.imageio/fitsoutput.cpp

@@ -3,6 +3,7 @@
 // https://github.com/AcademySoftwareFoundation/OpenImageIO
 
 #include <OpenImageIO/fmath.h>
+#include <OpenImageIO/strutil.h>
 
 #include "fits_pvt.h"
 
@@ -95,10 +96,6 @@ FitsOutput::write_scanline(int y, int /*z*/, TypeDesc format, const void* data,
     std::vector<unsigned char> data_tmp(m_spec.scanline_bytes(), 0);
     memcpy(&data_tmp[0], data, m_spec.scanline_bytes());
 
-    // computing scanline offset
-    long scanline_off = (m_spec.height - y) * m_spec.scanline_bytes();
-    fseek(m_fd, scanline_off, SEEK_CUR);
-
     // in FITS image data is stored in big-endian so we have to switch to
     // big-endian on little-endian machines
     if (littleendian()) {
@@ -115,12 +112,37 @@ FitsOutput::write_scanline(int y, int /*z*/, TypeDesc format, const void* data,
                         data_tmp.size() / sizeof(double));
     }
 
-    size_t byte_count = fwrite(&data_tmp[0], 1, data_tmp.size(), m_fd);
+    if (m_spec.nchannels == 1) {
+        // computing scanline offset
+        long scanline_off = (m_spec.height - y) * m_spec.scanline_bytes();
+        fseek(m_fd, scanline_off, SEEK_CUR);
+        size_t byte_count = fwrite(&data_tmp[0], 1, data_tmp.size(), m_fd);
+        fsetpos(m_fd, &m_filepos);
+        return byte_count == data_tmp.size();
+    }
 
+    // Channels are stored as separate contiguous width x height planes
+    // (NAXIS3 = nchannels), the only real-world FITS color-cube convention,
+    // so de-interleave this scanline and write each channel's row into its
+    // own plane.
+    size_t comp_size   = m_spec.format.size();
+    size_t row_bytes   = size_t(m_spec.width) * comp_size;
+    size_t plane_bytes = row_bytes * size_t(m_spec.height);
+    long row_off       = (m_spec.height - y) * long(row_bytes);
+    std::vector<unsigned char> chan_row(row_bytes);
+    bool ok = true;
+    for (int c = 0; c < m_spec.nchannels; ++c) {
+        for (int x = 0; x < m_spec.width; ++x)
+            memcpy(&chan_row[size_t(x) * comp_size],
+                   &data_tmp[(size_t(x) * m_spec.nchannels + c) * comp_size],
+                   comp_size);
+        fsetpos(m_fd, &m_filepos);
+        fseek(m_fd, long(c * plane_bytes) + row_off, SEEK_CUR);
+        size_t byte_count = fwrite(&chan_row[0], 1, row_bytes, m_fd);
+        ok &= (byte_count == row_bytes);
+    }
     fsetpos(m_fd, &m_filepos);
-
-    //byte_count == data.size --> all written
-    return byte_count == data_tmp.size();
+    return ok;
 }
 
 
@@ -264,16 +286,16 @@ FitsOutput::create_basic_header(std::string& header)
         axes += 1;
     header += create_card("NAXIS", num2str(axes));
 
-    // now we save NAXIS1 and NAXIS2
-    // this keywords represents width and height
+    // now we save NAXIS1, NAXIS2 (and NAXIS3 for multi-channel images)
     if (m_spec.nchannels == 1) {
         header += create_card("NAXIS1", num2str(m_spec.width));
         header += create_card("NAXIS2", num2str(m_spec.height));
     } else {
-        // 3D image for color
-        header += create_card("NAXIS1", num2str(m_spec.nchannels));
-        header += create_card("NAXIS2", num2str(m_spec.width));
-        header += create_card("NAXIS3", num2str(m_spec.height));
+        // One full-resolution plane per channel (NAXIS3 = nchannels), the
+        // only real-world FITS color-cube convention.
+        header += create_card("NAXIS1", num2str(m_spec.width));
+        header += create_card("NAXIS2", num2str(m_spec.height));
+        header += create_card("NAXIS3", num2str(m_spec.nchannels));
     }
 }