fits2vrt.py

"""
This file is part of fits2vrt
A FITS to GDAL Virtual Header conversion tool
31/01/2018
Authors : Chiara Marmo (chiara.marmo@u-psud.fr)
          Trent Hare (thare@usgs.gov)
Copyright : CNRS, Universite Paris-Sud - USGS
"""

import os, math
import os.path
import re
from astropy.io import fits
from astropy import wcs
from osgeo import gdal
from osgeo import osr

"""
The FITS Image metadata
"""
class fitskeys(object):

    # Image object initialization #
    def __init__(self,imname):

        # the image name
        self.__name = imname

        if os.path.isfile(imname):
            hdulist = fits.open(imname)

            # FITS header
            self.__header = hdulist[0].header

            # the WCS structure
            self.__wcs = wcs.WCS(self.__header)

            # alternate linear WCS
            ctype1 = re.compile("CTYPE1")
            px = re.compile("PX-")
            for key in self.__header.keys():
                alt = re.search(ctype1,key)
                if re.search(ctype1,key):
                    ctype = str(self.__header[key])
                    if re.search(px,ctype):
                        altkey = key[alt.end(0):]
                        self.__altkey = altkey

            # byte offset to the raster (points into FITS image)
            lenheader = (len(hdulist[0].header)+1)*80
            [blocks,remainder] = divmod(lenheader,2880)
            self.__offset = 2880 * (blocks+1)

            hdulist.close()

        else:
            print('ERROR!! Cannot find image {0:s}!\n'.format(imname))

    """
    The FITS metadata conversion to GDAL VRT Header
    """
    def fits2vrt(self):

        # Virtual header initialization
        driver = gdal.GetDriverByName('vrt')
        vrtname, fitsext = os.path.splitext(self.__name)

        # Dataset dimensions
        dimx = self.__header['NAXIS1']
        dimy = self.__header['NAXIS2']
        if (self.__header['NAXIS'] > 2):
            dimz = self.__header['NAXIS3']
        else:
            dimz = 1

        # Scale and Offset
        try:
            bzero = self.__header['BZERO']
            bscale = self.__header['BSCALE']
        except:
            bzero = 0.0
            bscale = 1.0

        """
        # Dataset bit type
        # will define
        # - nodata: nodata value if not defined elsewhere
        # - pxoffset: times offset to point to FITS raster
        # Following http://www.gdal.org/gdal_8h.html#a22e22ce0a55036a96f652765793fb7a4 (GDAL)
        # and https://heasarc.gsfc.nasa.gov/docs/software/fitsio/c/c_user/node20.html (FITS)
        """

        fbittype = self.__header['BITPIX']
        if (fbittype == 8):
            gbittype = gdal.GDT_Byte
            nodata = 0
            pxoffset = 1
        elif (fbittype == 16):
            pxoffset = 2
            if (bzero <= 0):
              gbittype = gdal.GDT_Int16
              nodata = -32768
            elif (bzero > 0):
              gbittype = gdal.GDT_UInt16
              nodata = 0
        elif (fbittype == 32):
            pxoffset = 4
            if (bzero <= 0):
              gbittype = gdal.GDT_Int32
            elif (bzero > 0):
              gbittype = gdal.GDT_UInt32
        elif (fbittype == -32):
            pxoffset = 4
            gbittype = gdal.GDT_Float32
            nodata = -3.40282e+38
        elif (fbittype == -64):
            pxoffset = 8
            gbittype = gdal.GDT_Float64
        else:
            print ("Bit Type not supported")
            print (fbittype)
            return 'fail' # need better error text or method

        # The next lines only work if gdal has cfitsio properly configured
        #src_ds = gdal.Open(self.__name)
        #dst_ds = driver.CreateCopy( vrtname + '.vrt', src_ds, 0)

        # Addressing FITS as raw raster: this will work without CFITSIO GDAL dependence.
        dst_ds = driver.Create( vrtname + '.vrt', dimx, dimy, 0 )

        # lineoffset only needed for raw VRT type
        lnoffset = dimx * pxoffset
        src_filename_opt = 'SourceFileName=' + self.__name
        px_offset_opt = 'PixelOffset=' + str(pxoffset)
        ln_offset_opt = 'LineOffset=' + str(lnoffset)

        for i in range(dimz):
          im_offset_opt = 'ImageOffset=' + str(self.__offset + i*dimx*dimy*pxoffset)
          options = [
            'subClass=VRTRawRasterBand',
            src_filename_opt,
            'relativeToVRT=1',
            im_offset_opt,
            px_offset_opt,
            ln_offset_opt,
            'ByteOrder=MSB' # FITS is always MSB
          ]
          result = dst_ds.AddBand( gbittype, options )
          if result != gdal.CE_None:
            print ('AddBand() returned error code')
            return 'fail'

        # Setting all non mandatory FITS keywords as metadata
        # COMMENT and HISTORY keywords are excluded too
        metadata = {}
        header = self.__header
        nometadata = ['SIMPLE','EXTEND','BITPIX','BZERO','BSCALE','COMMENT','NAXIS','HISTORY']
        for i in range(1, header['NAXIS']+1):
            nometadata.append('NAXIS'+str(i))
        for key in header.keys():
            if key not in nometadata:
                metadata[key] = str(header[key])
        dst_ds.SetMetadata( metadata )

        # Defining Target object
        try:
            target = header['OBJECT'] 
        except:
            print ("OBJECT keyword is missing")
            target = 'Undefined'

        # GDAL Spatial Coordinate System initialization
        srs = osr.SpatialReference()

        try:
            # Get radius values (maybe add an external method (e.g. string or URI)
            # new FITS keywords A_RADIUS, C_RADIUS 
            # note B_RADIUS (to define a triaxial) not generally used for mapping applications
            semiMajor = header['A_RADIUS']
            semiMinor = header['C_RADIUS']
            if ((semiMajor - semiMinor) > 0.00001):
                invFlattening= semiMajor / ( semiMajor - semiMinor)
            else:
                invFlattening= 0.0

            gcsName   = 'GCS_' + target
            datumName = 'D_' + target
            srsGeoCS  = 'GEOGCS["'+gcsName+'",DATUM["'+datumName+'",SPHEROID["'+ \
                        target+'",' + str(semiMajor) + ',' + str(invFlattening) + \
                        ']], PRIMEM["Reference_Meridian",0],' + \
                        'UNIT["degree",0.0174532925199433]]'

            srs.ImportFromWkt(srsGeoCS)

        except:
            print ("WARNING! No Radii keyword available, metadata will not contain DATUM information.")

        # Defining Geotransform: if linear WCS is defined 
        # GeoTransform[1] = CD1_1a
        # GeoTransform[2] = CD1_2a
        # GeoTransform[4] = CD2_1a
        # GeoTransform[5] = CD2_2a
        # GeoTransform[0] and GeoTransform[3] (topleftx, toplefty) must be computed.
        try:
            altkey = self.__altkey
            geot1 = header['CD1_1'+altkey]
            geot2 = header['CD1_2'+altkey]
            geot4 = header['CD2_1'+altkey]
            geot5 = header['CD2_2'+altkey] # FITS is upside-down

            # We are reading FITS rasters as raw matrix
            topleftx = header['CRVAL1'+altkey] - geot1 * (header['CRPIX1'+altkey]-0.5)
            toplefty = header['CRVAL2'+altkey] - geot5 * (header['CRPIX2'+altkey]-0.5)
            dst_ds.SetGeoTransform( [ topleftx, geot1, geot2, toplefty, geot4, geot5] )
        except:
            print ("WARNING! No linear keyword available, geotransformation matrix will be calculated using equatorial radius.")
            degtorad = math.pi / 180
            radfac = degtorad * semiMajor
            try:
                degtorad = math.pi / 180
                radfac = degtorad * semiMajor
                try:
                    geot1 = header['CD1_1'] * radfac
                    geot2 = header['CD1_2'] * radfac
                    geot4 = header['CD2_1'] * radfac
                    geot5 = header['CD2_2'] * radfac # FITS is upside-down
                except:
                    print ("WARNING! No CD keywords, looking for CDELT.")
                    try:
                        geot1 = header['CDELT1'] * radfac
                        geot2 = 0.
                        geot4 = 0.
                        geot5 = header['CDELT2'] * radfac # FITS is upside-down
                    except:
                        print ("WARNING! No WCS matrix available, geotransformation matrix cannot be calculated.")

                # We are reading FITS rasters as raw matrix
                topleftx = header['CRVAL1'] * radfac - geot1 * (header['CRPIX1']-0.5)
                toplefty = header['CRVAL2'] * radfac - geot5 * (header['CRPIX2']-0.5)
                dst_ds.SetGeoTransform( [ topleftx, geot1, geot2, toplefty, geot4, geot5] )

            except:
                print ("WARNING! No Radii keyword available, geotransformation matrix cannot be calculated.")


        # Defining projection type
        # Following http://www.gdal.org/ogr__srs__api_8h.html (GDAL)
        # and http://www.aanda.org/component/article?access=bibcode&bibcode=&bibcode=2002A%2526A...395.1077CFUL (FITS)

        wcsproj = (self.__header['CTYPE1'])[-3:]
        # Sinusoidal / SFL projection
        if ( wcsproj == 'SFL' ):
            gdalproj = 'Sinusoidal'
            srs.SetProjection(gdalproj)
            clong = self.__header['CRVAL1']
            srs.SetProjParm('longitude_of_center',clong)

	# Mercator, Oblique (Hotine) Mercator, Transverse Mercator
        # Mercator / MER projection
        elif ( wcsproj == 'MER' ): 
            gdalproj = 'Mercator'
            srs.SetProjection(gdalproj)
            cmer = self.__header['CRVAL1']
            srs.SetProjParm('central_meridian',cmer)
            olat = self.__header['CRVAL2']
            srs.SetProjParm('latitude_of_origin',olat)
            # Is scale factor a reference point or a ratio between pixel scales?
            if olat is not None:
                srs.SetProjParm('scale_factor',olat)
            else: #set default of 0.0
                srs.SetProjParm('scale_factor',0.0)

        # Equirectangular / CAR projection
        elif ( wcsproj == 'CAR' ):
            gdalproj = 'Equirectangular'
            srs.SetProjection(gdalproj)
            cmer = self.__header['CRVAL1']
            srs.SetProjParm('central_meridian',cmer)
            # The standard_parallel_1 defines where the local radius is calculated
            # not the center of Y Cartesian system (which is latitude_of_origin)
            # But FITS WCS only supports projections on the sphere
            # we assume here that the local radius is the one computed at the projection center
            spar = self.__header['CRVAL2']
            srs.SetProjParm('standard_parallel_1',spar)
            olat = self.__header['CRVAL2']
            srs.SetProjParm('latitude_of_origin',olat)

        # Lambert Azimuthal Equal Area / ZEA projection
        elif ( wcsproj == 'ZEA' ):
            gdalproj = 'Lambert_Azimuthal_Equal_Area'
            srs.SetProjection(gdalproj)
            clong = self.__header['CRVAL1']
            srs.SetProjParm('longitude_of_center',clong)
            clat = self.__header['CRVAL2']
            srs.SetProjParm('latitude_of_center',clat)

        # Lambert Conformal Conic 1SP / COO projection
        elif ( wcsproj == 'COO' ):
            gdalproj = 'Lambert_Conformal_Conic_1SP'
            srs.SetProjection(gdalproj)
            clong = self.__header['CRVAL1']
            srs.SetProjParm('longitude_of_center',clong)
            #clat = self.__header['XXXXX']
            #srs.SetProjParm('latitude_of_center',clat)
            scale = self.__header['XXXXX']
            if scale is not None:
                srs.SetProjParm('scale_factor',scale)
            else: #set default of 1.0
                srs.SetProjParm('scale_factor',1.0)

        # Orthographic / SIN projection
        elif ( wcsproj == 'SIN' ):
            gdalproj = 'Orthographic'
            srs.SetProjection(gdalproj)
            cmer = self.__header['CRVAL1']
            srs.SetProjParm('central_meridian',cmer)
            #olat = self.__header['XXXXX']
            #srs.SetProjParm('latitude_of_origin',olat)

        # Point Perspective / AZP projection
        elif ( wcsproj == 'AZP' ):
            gdalproj = 'perspective_point_height'
            srs.SetProjection(gdalproj)
            # appears to need height... maybe center lon/lat

        # Polar Stereographic / STG projection
        elif ( wcsproj == 'STG' ):
            gdalproj = 'Polar_Stereographic'
            srs.SetProjection(gdalproj)
            cmer = self.__header['CRVAL1']
            srs.SetProjParm('central_meridian',cmer)
            olat = self.__header['CRVAL2']
            srs.SetProjParm('latitude_of_origin',olat)

        else:
            print ("Unknown projection")
            return 'fail'

        projname = gdalproj + '_' + target
        srs.SetAttrValue('projcs',projname)
        # false easting and northing not used for planetary
        srs.SetProjParm('false_easting',0.0)
        srs.SetProjParm('false_northing',0.0)

        wkt = srs.ExportToWkt()
        print ('projection:\n'+wkt)
        dst_ds.SetProjection(wkt)

        # Adding SimpleSource info
        band = dst_ds.GetRasterBand(1)
        band.SetNoDataValue(nodata)
        band.SetScale(bscale)
        band.SetOffset(bzero)

        # Close properly the dataset
        dst_ds = None
        #src_ds = None