Bayer Reconstruction of 12-bit PGM in RGGB configuration with color/contrast balancing

reconstruct.py

#!/usr/bin/python
#
# python reconstruct.py filename.pgm
#	takes in a raw image having a Bayer Color Formatted Array (RGGB)
#	and outputs two JPEG files. The first file is simply the 
#	Bayer Reconstruction as RGB. The second file has been
# 	processed to enhance color and contrast using combinations of:
#	1) independent color channel scaling 
#	2) simplest color balancing
#	3) white balancing
#	4) contrast limited adaptive histogram equalization
#	5) gamma adjustment
#
#  author: Dr. Bradley Matthew Battista
#  
#  depends on opencv (pip install opencv-python)
#
#

import sys
import os
import numpy
import math
import cv2

# Gamma adjustment
def adjust_gamma(img, gamma=1.0):
  invGamma = 1.0 / gamma
  table = numpy.array([((i / 255.0) ** invGamma) * 255
    for i in numpy.arange(0, 256)]).astype("uint8")
  return cv2.LUT(img, table)

# Contrast Limited Adaptive Histogram Equalization (CLAHE)
def apply_clahe(img):
  clahe = cv2.createCLAHE(clipLimit=1, tileGridSize=(8,8))
  lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
  l, a, b = cv2.split(lab)  # split on 3 different channels
  l2 = clahe.apply(l)  # apply CLAHE to the L-channel
  lab = cv2.merge((l2,a,b))  # merge channels
  result = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
  return result

# White Balance
def white_balance(img):
  result = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
  avg_a = numpy.average(result[:, :, 1])
  avg_b = numpy.average(result[:, :, 2])
  result[:, :, 1] = result[:, :, 1] - ((avg_a - 128) * (result[:, :, 0] / 255.0) * 1.1)
  result[:, :, 2] = result[:, :, 2] - ((avg_b - 128) * (result[:, :, 0] / 255.0) * 1.1)
  result = cv2.cvtColor(result, cv2.COLOR_LAB2BGR)
  return result

# Mask for simplest color balance
def apply_mask(matrix, mask, fill_value):
  masked = numpy.ma.array(matrix, mask=mask, fill_value=fill_value)
  return masked.filled()

# Threshold for simplest color balance
def apply_threshold(matrix, low_value, high_value):
  low_mask = matrix < low_value
  matrix = apply_mask(matrix, low_mask, low_value)
  high_mask = matrix > high_value
  matrix = apply_mask(matrix, high_mask, high_value)
  return matrix

# Simplest Color Balance Algorithm
def simplest_cb(img, percent):
  assert img.shape[2] == 3
  assert percent > 0 and percent < 100
  half_percent = percent / 200.0
  channels = cv2.split(img)
  out_channels = []
  for channel in channels:
    assert len(channel.shape) == 2
    height, width = channel.shape
    vec_size = width * height
    flat = channel.reshape(vec_size)
    assert len(flat.shape) == 1
    flat = numpy.sort(flat)
    n_cols = flat.shape[0]
    low_val  = flat[int(math.floor(n_cols * half_percent))]
    high_val = flat[int(math.ceil( n_cols * (1.0 - half_percent)))]
    thresholded = apply_threshold(channel, low_val, high_val)
    normalized = cv2.normalize(thresholded, thresholded.copy(), 0, 255, cv2.NORM_MINMAX)
    out_channels.append(normalized)
  return cv2.merge(out_channels)

# Channel scaling
def scale_chn(img, chn, factor):
  img = img.astype('uint16')
  scale = int(2**(factor-1))
  if chn == 'r':
    b = img[:,:,0]
    g = img[:,:,1]
    r = numpy.clip(img[:,:,2]+scale-1, 0, 255)
  if chn == 'g':
    b = img[:,:,0]
    g = numpy.clip(img[:,:,1]+scale-1, 0, 255)
    r = img[:,:,2]
  if chn == 'b':
    b = numpy.clip(img[:,:,0]+scale-1, 0, 255)
    g = img[:,:,1]
    r = img[:,:,2]
 
  return cv2.merge((b.astype('uint8'), g.astype('uint8'), r.astype('uint8'))) 

if __name__ == '__main__':

  infile = sys.argv[1]
  file,ext = os.path.splitext(infile)

  ### READ RAW PGM AND CONVERT TO JPEG WITHOUT PROCESSING ###

  # Read in file and shift 4 bits (get 12-bit values from 16-bit precision)
  bayer = numpy.right_shift(cv2.imread(infile,-1),4)

  # Perform a Bayer Reconstruction
  demosaic = cv2.cvtColor(bayer, cv2.COLOR_BAYER_BG2BGR)

  # Save the demosaic file as a JPEG
  cv2.imwrite(file+'.jpg',demosaic)


  ### READ NEW JPEG AND APPLY PROCESSING ###

  # Process the JPEG to improve quality
  img = cv2.imread(file+'.jpg',1)

  # Uncomment this section if you want to compare individual steps in a image panel
  #im_gamma = adjust_gamma(img,1.1)
  #im_scb = simplest_cb(img,0.25)
  #im_clahe = apply_clahe(img)
  #im_wb = white_balance(img)
  #panels = numpy.vstack(( numpy.hstack((im_gamma, im_wb)), numpy.hstack((im_clahe, im_scb)) ))
  #cv2.imwrite(file+'_panel.jpg',panels)

  # Apply the processing workflow
  img = scale_chn(img, 'r', 5)
  img = scale_chn(img, 'g', 3)
  img = scale_chn(img, 'b', 1)
  img = simplest_cb(img,0.25)
  #img = white_balance(img)
  img = apply_clahe(img)
  #img = adjust_gamma(img,1.2)

  # Save the processed file as a JPEG
  cv2.imwrite(file+'_proc.jpg',img)