Permalink
Switch branches/tags
Nothing to show
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
372 lines (320 sloc) 14.5 KB
""" Simple RAW image processing module"""
import sys
import os
import scipy
from scipy import signal
import numpy as np
from numpy.lib.stride_tricks import as_strided
import rawpy
""" Process RAW file into a image file.
Example usage:
raw = read("sample.ARW")
rgb = process(raw)
write(rgb, "output.ARW")
"""
def read(filename):
"""
Read RAW data from specified file. Currently supported formats are
ARW (Sony RAW format)
JPEG with Raspberry Pi V2.1 camera RAW
:param filename: path to the target RAW file
"""
return rawpy.imread(filename)
def check_functions(filename):
""" Check what functions to be enabled based on filename"""
white_level = 1024
shading_enable = True
defect_correction_enable = True
noise_parameters = (8, 2, 246)
wbg_norm = 1
extension = os.path.splitext(filename)[1]
if extension in (".ARW", ".arw"):
shading_enable = False
defect_correction_enable = False
white_level = 8192
noise_parameters = (8, 0.2, 25)
wbg_norm = 1024
return (shading_enable, defect_correction_enable, white_level, wbg_norm, noise_parameters)
DEFALT_MATRIX = (1024, 0, 0, 0, 1024, 0, 0, 0, 1024)
DEFALT_TONE = ((0, 64, 128, 192, 256), (0, 64, 128, 192, 256))
def process(filename, output_filename, color_matrix=DEFALT_MATRIX, tone_curve=DEFALT_TONE):
"""
This processes RAW data that was read by read() method.
Must be called after read() operation. No error is checked.
"""
shading_enable, defect_crrection_enable, white_level, wbg_norm, noise_param = check_functions(filename)
raw = read(filename)
raw_array = get_raw_array(raw)
raw_array = black_level_correction(raw_array, raw.black_level_per_channel, raw.raw_pattern)
if defect_crrection_enable:
raw_array = defect_crrection(raw_array)
if shading_enable:
raw_array = lens_shading_correction(raw_array, LSC_DEFAULT)
raw_array = white_balance_Bayer(raw_array, raw.camera_whitebalance, wbg_norm, raw.raw_pattern)
rgb_array = advanced_demosaic(raw_array, raw.raw_pattern)
del raw_array, raw
rgb_array = noise_filter(rgb_array, noise_param[0], noise_param[1], noise_param[2])
rgb_array = color_correction_matrix(rgb_array, color_matrix)
rgb_array = gamma_correction(rgb_array/white_level, 2.2)
rgb_array = edge_correction(rgb_array, 2, 0.25, 1, 0.25)
rgb_array = tone_curve_correction(rgb_array, tone_curve[0], tone_curve[1])
write(rgb_array, output_filename)
def get_raw_array(raw):
""" convert raw_img into numpy array"""
h, w = raw.sizes.raw_height, raw.sizes.raw_width
raw_array = np.array(raw.raw_image).reshape((h, w)).astype('float')
return raw_array
def black_level_correction(raw_array, black_level_per_channel, bayer_pattern):
# rearrange black level
black_level = [0] * 4
black_level[bayer_pattern[0, 0]] = black_level_per_channel[bayer_pattern[0, 0]]
black_level[bayer_pattern[0, 1]] = black_level_per_channel[bayer_pattern[0, 1]]
black_level[bayer_pattern[1, 0]] = black_level_per_channel[bayer_pattern[1, 0]]
black_level[bayer_pattern[1, 1]] = black_level_per_channel[bayer_pattern[1, 1]]
blc_raw = raw_array.copy()
blc_raw[0::2, 0::2] -= black_level[0]
blc_raw[0::2, 1::2] -= black_level[1]
blc_raw[1::2, 0::2] -= black_level[2]
blc_raw[1::2, 1::2] -= black_level[3]
return blc_raw
def defect_crrection(raw_array):
dpc_raw = raw_array.copy()
footprint = np.ones((5, 5))
footprint[2, 2] = 0
for (yo, xo) in ((0, 0), (1, 0), (0, 1), (1, 1)):
single_channel = dpc_raw[yo::2, xo::2]
flt = np.array([[0, 1, 0], [1, 0, 1], [0, 1, 0]]) / 4
average = scipy.signal.convolve2d(single_channel, flt, mode='same')
local_max = scipy.ndimage.filters.maximum_filter(single_channel, footprint=footprint, mode='mirror')
local_min = scipy.ndimage.filters.minimum_filter(single_channel, footprint=footprint, mode='mirror')
threshold = 16
mask = (single_channel < local_min - threshold) + (single_channel > local_max + threshold)
single_channel[mask] = average[mask]
return dpc_raw
LSC_DEFAULT = [np.array([6.07106808e-07, 9.60556906e-01]),
np.array([6.32044369e-07, 9.70694361e-01]),
np.array([6.28455183e-07, 9.72493898e-01]),
np.array([9.58743579e-07, 9.29427169e-01])]
def lens_shading_correction(raw, coef):
"""
Apply lens shading correction to Bayer input (raw)
Parameter (coef) needs to be array type of coef[4][2]
coef[color][0] is coefficient on 2nd order term.
coef[color][1] is offset
"""
h, w = raw.shape
gain_map = np.zeros((h, w))
center_y, center_x = h // 2, w // 2
x = np.arange(0, w) - center_x
y = np.arange(0, h) - center_y
xs, ys = np.meshgrid(x, y, sparse=True)
r2 = ys * ys + xs * xs
gain_map[::2, ::2] = r2[::2, ::2] * coef[0][0] + coef[0][1]
gain_map[1::2, ::2] = r2[1::2, ::2] * coef[1][0] + coef[1][1]
gain_map[::2, 1::2] = r2[::2, 1::2] * coef[2][0] + coef[2][1]
gain_map[1::2, 1::2] = r2[1::2, 1::2] * coef[3][0] + coef[3][1]
return raw * gain_map
def preview_demosaic(raw_array, bayer_pattern):
""" Very simple demosaic with down sampling for preview purpose"""
h, w = raw_array.shape[0], raw_array.shape[1]
shuffle = np.zeros((h // 2, w // 2, 4))
shuffle[:, :, bayer_pattern[0, 0]] += raw_array[0::2, 0::2]
shuffle[:, :, bayer_pattern[0, 1]] += raw_array[0::2, 1::2]
shuffle[:, :, bayer_pattern[1, 0]] += raw_array[1::2, 0::2]
shuffle[:, :, bayer_pattern[1, 1]] += raw_array[1::2, 1::2]
dms_img = np.zeros((h // 2, w // 2, 3))
dms_img[:, :, 0] = shuffle[:, :, 0]
dms_img[:, :, 1] = (shuffle[:, :, 1] + shuffle[:, :, 3]) / 2
dms_img[:, :, 2] = shuffle[:, :, 2]
return dms_img
def simple_demosaic(raw, raw_array):
""" Simple demosaic algorithm with linear interpolation """
h, w = raw_array.shape
dms_img2 = np.zeros((h, w, 3))
green = raw_array.copy()
green[(raw.raw_colors == 0) | (raw.raw_colors == 2)] = 0
g_flt = np.array([[0, 1 / 4, 0], [1 / 4, 1, 1 / 4], [0, 1 / 4, 0]])
dms_img2[:, :, 1] = signal.convolve2d(green, g_flt, boundary='symm', mode='same')
red = raw_array.copy()
red[raw.raw_colors != 0] = 0
rb_flt = np.array([[1 / 4, 1 / 2, 1 / 4], [1 / 2, 1, 1 / 2], [1 / 4, 1 / 2, 1 / 4]])
dms_img2[:, :, 0] = signal.convolve2d(red, rb_flt, boundary='symm', mode='same')
blue = raw_array.copy()
blue[raw.raw_colors != 2] = 0
rb_flt = np.array([[1 / 4, 1 / 2, 1 / 4], [1 / 2, 1, 1 / 2], [1 / 4, 1 / 2, 1 / 4]])
dms_img2[:, :, 2] = signal.convolve2d(blue, rb_flt, boundary='symm', mode='same')
return dms_img2
def advanced_demosaic(dms_input, bayer_pattern):
""" Demosaic algorithm in frequency domain """
hlpf = np.array([[1, 2, 3, 4, 3, 2, 1]]) / 16
vlpf = np.transpose(hlpf)
hhpf = np.array([[-1, 2, -3, 4, -3, 2, -1]]) / 16
vhpf = np.transpose(hhpf)
identity_filter = np.zeros((7, 7))
identity_filter[3, 3] = 1
# generate FIR filters to extract necessary components
FC1 = np.matmul(vhpf, hhpf)
FC2H = np.matmul(vlpf, hhpf)
FC2V = np.matmul(vhpf, hlpf)
FL = identity_filter - FC1 - FC2V - FC2H
# f_C1 at 4 corners
c1_mod = signal.convolve2d(dms_input, FC1, boundary='symm', mode='same')
# f_C1^1 at wy = 0, wx = +Pi/-Pi
c2h_mod = signal.convolve2d(dms_input, FC2H, boundary='symm', mode='same')
# f_C1^1 at wy = +Pi/-Pi, wx = 0
c2v_mod = signal.convolve2d(dms_input, FC2V, boundary='symm', mode='same')
# f_L at center
f_L = signal.convolve2d(dms_input, FL, boundary='symm', mode='same')
# Move c1 to the center by shifting by Pi in both x and y direction
# f_c1 = c1 * (-1)^x * (-1)^y
f_c1 = c1_mod.copy()
f_c1[:, 1::2] *= -1
f_c1[1::2, :] *= -1
if bayer_pattern[0, 0] == 1 or bayer_pattern[0, 0] == 3:
f_c1 *= -1
# Move c2a to the center by shifting by Pi in x direction, same for c2b in y direction
c2h = c2h_mod.copy()
c2h[:, 1::2] *= -1
if bayer_pattern[0, 0] == 2 or bayer_pattern[1, 0] == 2:
c2h *= -1
c2v = c2v_mod.copy()
c2v[1::2, :] *= -1
if bayer_pattern[0, 0] == 2 or bayer_pattern[0, 1] == 2:
c2v *= -1
# f_c2 = (c2v_mod * x_mod + c2h_mod * y_mod) / 2
f_c2 = (c2v + c2h) / 2
# generate RGB channel using
# [R, G, B] = [[1, 1, 2], [1, -1, 0], [1, 1, - 2]] x [L, C1, C2]
height, width = dms_input.shape
dms_img = np.zeros((height, width, 3))
dms_img[:, :, 0] = f_L + f_c1 + 2 * f_c2
dms_img[:, :, 1] = f_L - f_c1
dms_img[:, :, 2] = f_L + f_c1 - 2 * f_c2
return dms_img
def white_balance_Bayer(raw_array, wbg, wbg_norm, bayer_pattern):
""" Apply white balance to bayer input"""
img_wb = raw_array.copy()
img_wb[0::2, 0::2] *= wbg[bayer_pattern[0, 0]] / wbg_norm
img_wb[0::2, 1::2] *= wbg[bayer_pattern[0, 1]] / wbg_norm
img_wb[1::2, 0::2] *= wbg[bayer_pattern[1, 0]] / wbg_norm
img_wb[1::2, 1::2] *= wbg[bayer_pattern[1, 1]] / wbg_norm
return img_wb
def noise_filter(rgb_array, coef=8, read_noise=2, shot_noise=246):
""" Apply bilateral noise filter to RGB image"""
h, w, _ = rgb_array.shape
luma_img = rgb_array[:, :, 0] + rgb_array[:, :, 1] + rgb_array[:, :, 2]
average = scipy.ndimage.filters.uniform_filter(luma_img, 5, mode='mirror')
sigma_map = average * shot_noise + read_noise
del average
sigma_map[sigma_map < 1] = 1
sy, sx = sigma_map.strides
sigma_tile = as_strided(sigma_map, strides=(sy, sx, 0, 0), shape=(h, w, 5, 5))
sigma_tile = sigma_tile[2:h-2, 2:w-2, :, :]
del sigma_map
sy, sx = luma_img.strides
luma_tile = as_strided(luma_img, strides=(sy, sx, 0, 0), shape=(h, w, 5, 5))
luma_tile = luma_tile[2:h-2, 2:w-2, :, :]
luma_box = as_strided(luma_img, strides=(sy, sx, sy, sx), shape=(h-4, w-4, 5, 5))
del luma_img
diff = luma_box - luma_tile
del luma_tile, luma_box
diff = diff * diff
weight = np.exp(-coef * diff / sigma_tile)
del diff, sigma_tile
weight_sum = weight.sum(axis=(2, 3))
sy, sx, sz, sw = weight.strides
weight_extend = as_strided(weight, strides=(sy, sx, 0, sz, sw), shape=(h-4, w-4, 3, 5, 5))
del weight
sy, sx = weight_sum.strides
weight_sum_extend = as_strided(weight_sum, strides=(sy, sx, 0), shape=(h-4, w-4, 3))
del weight_sum
sy, sx, sz = rgb_array.strides
img_boxes = as_strided(rgb_array, strides=(sy, sx, sz, sy, sx), shape=(h-4, w-4, 3, 5, 5))
img_flt = (weight_extend * img_boxes).sum(axis=(3, 4)) / weight_sum_extend
return img_flt
def color_correction_matrix(rgb_array, color_matrix):
""" Apply color correction matrix to RGB array"""
img_ccm = np.zeros_like(rgb_array)
ccm = np.array(color_matrix).reshape((3, 3))
norm = ccm.sum(axis=1).mean()
for c in (0, 1, 2):
img_ccm[:, :, c] = ccm[c, 0] * rgb_array[:, :, 0] + \
ccm[c, 1] * rgb_array[:, :, 1] + \
ccm[c, 2] * rgb_array[:, :, 2]
return img_ccm / norm
def gamma_correction(rgb_array, gamma_coef):
""" Apply gamma correction to RGB image"""
img_gamma = rgb_array.copy()
img_gamma[img_gamma < 0] = 0
img_gamma = np.power(img_gamma, 1/gamma_coef)
return img_gamma
def apply_matrix(input_array, matrix):
img_out = np.zeros_like(input_array)
for c in (0, 1, 2):
img_out[:, :, c] = matrix[c, 0] * input_array[:, :, 0] + \
matrix[c, 1] * input_array[:, :, 1] + \
matrix[c, 2] * input_array[:, :, 2]
return img_out
RGB_TO_YCBCR = np.array([[0.299, 0.587, 0.144],
[-0.168736, -0.331264, 0.5],
[0.5, -0.418688, -0.081312]])
def edge_correction(rgb_array, sigma1=2, coef1=0.25, sigma2=1, coef2=0.25):
""" Edge correction for RGB input"""
img_rgb = rgb_array.copy() * 256
img_rgb[img_rgb < 0] = 0
img_rgb[img_rgb > 255] = 255
img_ycbcr = apply_matrix(img_rgb, RGB_TO_YCBCR)
luma = img_ycbcr[:, :, 0]
unsharpen1 = scipy.ndimage.gaussian_filter(luma, sigma=sigma1)
unsharpen2 = scipy.ndimage.gaussian_filter(luma, sigma=sigma2)
sharpen = luma + coef1 * (luma - unsharpen1) + coef2 * (luma - unsharpen2)
img_ycbcr[:, :, 0] = sharpen
ycbcr2rgb = np.linalg.inv(RGB_TO_YCBCR)
img_shp_rgb = apply_matrix(img_ycbcr, ycbcr2rgb) / 256
img_shp_rgb[img_shp_rgb < 0] = 0
img_shp_rgb[img_shp_rgb > 1] = 1
return img_shp_rgb
def tone_curve_correction(img_rgb, xs=(0, 64, 128, 192, 256), ys=(0, 64, 128, 192, 256)):
func = scipy.interpolate.splrep(xs, ys)
img_ycbcr = apply_matrix(img_rgb * 256, RGB_TO_YCBCR)
img_ycbcr[:, :, 0] = scipy.interpolate.splev(img_ycbcr[:, :, 0], func)
ycbcr2rgb = np.linalg.inv(RGB_TO_YCBCR)
img_rgb_out = apply_matrix(img_ycbcr, ycbcr2rgb)
return img_rgb_out / 256
def write(rgb_image, output_filename):
"""
Write the processed RGB image to a specified file as PNG format.
Thsi must be called after process(). No error is checked.
:param output_filename: path to the output file. Extension must be png.
"""
import imageio
outimg = rgb_image.copy() * 256
outimg[outimg < 0] = 0
outimg[outimg > 255] = 255
imageio.imwrite(output_filename, outimg.astype('uint8'))
def main(argv):
""" main function """
if (len(argv) < 2):
print("Usage: {} input_filename [output_filename] \"color_matrix\" \"tone_x\" \"tone_y\"".format(argv[0]))
print("\tDefault output_filename is output.png")
print("\tDefault matrix is identity matrix \"1024, 0, 0, 0, 1024, 0, 0, 0, 1024\"")
print("\tDefault tone curve is identity function \"0, 128, 256] [0, 128, 256\"")
print("\tExample: python3 {} sample.ARW sample.png \"1141, -205, 88, -52, 1229, -154, 70, -225, 1179\" \"0, 72, 128, 200, 256\" \"0, 56, 128, 220, 256\"".format(argv[0]))
print("\tSupported RAW format is ARW (Sony RAW) and Raspberry Pi (embedded in JPEG)")
print("\tSupported output format is PNG only")
return
filename = argv[1]
output_filename = "output.png"
color_matrix = [1024, 0, 0, 0, 1024, 0, 0, 0, 1024]
tone_curve = [(0, 64, 128, 192, 256), (0, 64, 128, 192, 256)]
if len(argv) > 2:
output_filename = argv[2]
if len(argv) > 3:
color_matrix = [int(value) for value in (argv[3]).split(',')]
if len(argv) > 4:
tone_curve[0] = [int(value) for value in (argv[4]).split(',')]
if len(argv) > 5:
tone_curve[1] = [int(value) for value in (argv[5]).split(',')]
process(filename, output_filename, color_matrix, tone_curve)
if __name__ == "__main__":
main(sys.argv)