# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# Package that provides access to Fits files 
from astropy.io import fits     
import numpy
import OpenEXR
import Imath
from PIL import Image


# +
# input_exr - input EXR file
# output_fits - output FITS file
# return - EXR data as array
def exr_to_fits(input_exr, output_fits):
    
    # Open the EXR file. 
    exr = OpenEXR.InputFile(input_exr)
    
#     print (exr.header().keys())
    
#     for element in exr.header().keys():
#         print (exr.header()[element])
    
    # Access EXR image data by specifying the channel. Returns EXR data in byte form. 
    exr_bytes = exr.channel("R", Imath.PixelType(Imath.PixelType.FLOAT))
    
    # Convert EXR byte to Image object
    exr_image = Image.frombytes("F", (2048, 1489), exr_bytes)
    
    # Convert Image to array. 
    exr_array = numpy.array(exr_image)
        
    # Create FITS HDU with EXR data, copy header, add HDU to HDUList, and write FITS output file. 
    hdu = fits.PrimaryHDU(exr_array)
#     hdu.header = exr.header()['header']
    hdul = fits.HDUList([hdu])
    hdul.writeto(output_fits)
    
    return exr_array


# +
# input_fits - input FITS file
# output_exr - output EXR file\
# return - exr data as array
def fits_to_exr(input_fits, output_exr, hdu_num = -1, float_format = 32):
    
    # Returns HDUList object which is a list-like collection of HDU objects.
    # Each HDU contains a header and data. Header contains keyword, value, and comment. Keyword / value
    # can be manipulated like a python dict. Data is a numpy ndarray (n-dimensional array). 
    hdu_list = fits.open(input_fits)
    hdu = hdu_list[0]
    
    # Choose either the specified hdu or the first hdu with an image. 
    if hdu_num != -1:
        hdu = hdu_list[hdu_num]
    else:
        for element in hdu_list:
            if element.is_image:
                hdu = element
                break
            
    fits_header = hdu.header 
    fits_data = hdu.data
    
    hdu_list.close()
    
    # Get the dimension of the image. 
    width, height = fits_data.shape
    
    # Create the EXR header with appropriate height and width.
    # Copy the Fits header to EXR. Probably need to do more of a deep copy later. 
    header = OpenEXR.Header(height, width)
    header['fits_h'] = fits_header

#     header['owner'] = 'Copyright 2007 <NAME>'
# #     header['dummy1'] = {'G': 2}
# #     header['dummy2'] = {'fits': 4}
    
    # Convert the elements in fitz_data to float format.
    # This can be done analogously for 32, or 16 bit. 
    fits_array = []
        
    if float_format == 32:
        fits_array = numpy.float32(fits_data)
        header['float_t'] = 32
        header['channels'] = { 'R' : Imath.Channel(Imath.PixelType(OpenEXR.FLOAT))}
        
    elif float_format == 16:
        fits_array = numpy.float16(fits_data)
        header['float_t'] = 16
        header['channels'] = { 'R' : Imath.Channel(Imath.PixelType(OpenEXR.HALF))}
        
    else:
        print ("INVALID float type.")
        return []
    
#     print (header.keys())
    
#     for element in header.keys():
#         print (header[element])
        
    # Create EXR output file with header. Write Fits data to EXR red channel. 
    exr = OpenEXR.OutputFile(output_exr, header)

    
    exr.writePixels({"R": fits_array})
    
    return fits_data
    

# +
print ("Start")
fits_input = "frame-r-006793-3-0127.fits"
exr_output = "exr_demo.exr"
fits_output = "fits_demo6.fits"

print ("Processing FITS->EXR")
one = fits_to_exr(fits_input, exr_output, float_format = 32)

print ("Processing EXR->FITS")
two = exr_to_fits(exr_output, fits_output)

for i in range(len(one)):
    for j in range(len(one[i])):
        if one[i][j] != two[i][j]:
            print (one[i][j], two[i][j])
            print ("Ouch.")
            break
print ("The images are the same!")

# +
# Current Issues: No support for 64 bit. FITS header is not copying over. 
# -





StarcoderdataJupyterScriptsDedupFiltered