#!/usr/bin/env python3                                                                                                                                                                                         
 
'''                                                                                                                                                                                                            
  Fetch RGGB data from FITS file, convert it to an RGB numpy                                                                                                                                                   
  array, with and without interpolation, and store it as                                                                                                                                                       
  uint16 PNG. This code serves as an example.                                                                                                                                                                  
 
  Variant "no interpolation" based on Osmo Systems' picamraw.                                                                                                                                                  
 
  Requires Python >= 3.5                                                                                                                                                                                       
 
  Markus Wildi, 2019                                                                                                                                                                                           
 
'''                                                                                                                                                                                                            
 
import cv2                                                                                                                                                                                                     
from astropy.io import fits                                                                                                                                                                                    
 
import numpy as np                                                                                                                                                                                             
from enum import Enum                                                                                                                                                                                          
 
# source: adapted from Osmo Systems' picamraw (https://github.com/osmosystems/picamraw).                                                                                                                       
 
class BayerOrder(Enum):                                                                                                                                                                                        
    ''' There are four supported arrangements of the R, G, G, and B pixels:                                                                                                                                    
        RGGB:                                                                                                                                                                                                  
              RG                                                                                                                                                                                               
              GB                                                                                                                                                                                               
        GBRG:                                                                                                                                                                                                  
              GB                                                                                                                                                                                               
              RG                                                                                                                                                                                               
        BGGR:                                                                                                                                                                                                  
              BG                                                                                                                                                                                               
              GR                                                                                                                                                                                               
        GRBG:                                                                                                                                                                                                  
              GR                                                                                                                                                                                               
              BG                                                                                                                                                                                               
    '''                                                                                                                                                                                                        
 
    RGGB = 'RGGB'                                                                                                                                                                                              
    GBRG = 'GBRG'                                                                                                                                                                                              
    BGGR = 'BGGR'                                                                                                                                                                                              
    GRBG = 'GRBG'                                                                                                                                                                                              
 
R_CHANNEL_INDEX, G_CHANNEL_INDEX, B_CHANNEL_INDEX = [0, 1, 2]                                                                                                                                                  
 
BAYER_ORDER_TO_RGB_CHANNEL_COORDINATES = {                                                                                                                                                                     
    # (ry, rx), (gy, gx), (Gy, Gx), (by, bx)                                                                                                                                                                   
    BayerOrder.RGGB: ((0, 0), (1, 0), (0, 1), (1, 1)),                                                                                                                                                         
    BayerOrder.GBRG: ((1, 0), (0, 0), (1, 1), (0, 1)),                                                                                                                                                         
    BayerOrder.BGGR: ((1, 1), (0, 1), (1, 0), (0, 0)),                                                                                                                                                         
    BayerOrder.GRBG: ((0, 1), (1, 1), (0, 0), (1, 0)),                                                                                                                                                         
}                                                                                                                                                                                                              
 
def _guard_attribute_is_a_multiple_of(attribute_name, attribute_value, multiple):                                                                                                                              
    if not attribute_value % multiple == 0:                                                                                                                                                                    
        raise ValueError(                                                                                                                                                                                      
            'Incoming data is the wrong shape: {attribute_name} ({attribute_value}) is not a multiple of {multiple}'                                                                                           
            .format(**locals())                                                                                                                                                                                
        )                                                                                                                                                                                                      
 
 
bggr = fits.getdata('4sec.fits')                                                                                                                                                                               
#                                                                                                                                                                                                              
# variant without interpolation                                                                                                                                                                                
#                                                                                                                                                                                                              
_guard_attribute_is_a_multiple_of('width', bggr.shape[0], 2)                                                                                                                                                   
_guard_attribute_is_a_multiple_of('height', bggr.shape[1], 2)                                                                                                                                                  
 
rgb_no_int = np.zeros((int(bggr.shape[0]/2), int(bggr.shape[1]/2), 3))                                                                                                                                         
 
((ry, rx), (gy, gx), (Gy, Gx), (by, bx)) = BAYER_ORDER_TO_RGB_CHANNEL_COORDINATES[BayerOrder.RGGB]                                                                                                             
 
rgb_no_int[:, :, R_CHANNEL_INDEX] = bggr[ry::2, rx::2]                                                                                                                                                         
rgb_no_int[:, :, G_CHANNEL_INDEX] = (bggr[gy::2, gx::2] + bggr[Gy::2, Gx::2])/2.                                                                                                                               
rgb_no_int[:, :, B_CHANNEL_INDEX] = bggr[by::2, bx::2]                                                                                                                                                         
cv2.imwrite( '4sec_no_int.png', rgb_no_int.astype(np.uint16))                                                                                                                                                  
#                                                                                                                                                                                                              
# variant with interpolation                                                                                                                                                                                   
#                                                                                                                                                                                                              
rgb = cv2.cvtColor(bggr, cv2.COLOR_BAYER_BG2RGB)                                                                                                                                                               
 
cv2.imwrite('4sec_int.png', rgb)