raw2rgb.py

import numpy as np
import cv2
import matplotlib.image

def gamma_compression(image):
    """Converts from linear to gamma space."""
    return np.maximum(image, 1e-8) ** (1.0 / 2.2)

def tonemap(image):
    """Simple S-curved global tonemap"""
    return (3*(image**2)) - (2*(image**3))

def postprocess_raw(raw):
    """Simple post-processing to visualize demosaic RAW imgaes
    Input:  (h,w,3) RAW image normalized
    Output: (h,w,3) post-processed RAW image
    
    Source: https://github.com/mv-lab/AISP
    """
    raw = gamma_compression(raw)
    raw = tonemap(raw)  # tonemap cause weird color (Official sRGB version does gamma correction only)
    raw = np.clip(raw, 0, 1)
    return raw

def demosaic(raw):
    """Simple demosaicing to visualize RAW images
    Inputs:
     - raw: (h,w,4) RAW RGGB image normalized [0..1] as float32
    Returns: 
     - Simple Avg. Green Demosaiced RAW image with shape (h*2, w*2, 3)
     
    source: https://github.com/mv-lab/AISP
    """
    assert raw.shape[-1] == 4
    shape = raw.shape
    
    red        = raw[:,:,0]
    green_red  = raw[:,:,1]
    green_blue = raw[:,:,2]
    blue       = raw[:,:,3]
    avg_green  = (green_red + green_blue) / 2
    image      = np.stack((red, avg_green, blue), axis=-1)
    image      = cv2.resize(image, (shape[1]*2, shape[0]*2))
    return image

def raw2rgb(img, fileName):
    rgb = postprocess_raw(demosaic(img.squeeze(0).permute(1, 2, 0).cpu().numpy()))
    matplotlib.image.imsave(fileName, rgb)

if __name__ == "__main__":
    print("Visualize RAW image")
    low_img, high_img = val_loader.dataset.__getitem__(0)
    with torch.no_grad():
        enhanced_img = model(low_img.cuda())
    raw2rgb(low_img, 'low_resolution_raw.png')
    raw2rgb(high_img, 'high_resolution_raw.png')
    raw2rgb(enhanced_img, 'enhanced_raw.png')