diving_color_correction.py

# A script to improve color channels in diving videos
# By Marek Lukasiewicz 2018
# http://lukasiewicz.tech

# Requires python openCV with ffmpeg installed to open all video types
# Requires ffmpeg installed command ffmpeg used to copy sound from original video

from os import system, remove
from sys import argv
import numpy as np
import cv2

filename = None
if len(argv) != 2:
    print('Run with filename to process as the argument')
    quit()
else:
    filename = str(argv[1])
    print('Correcting ' + filename)
    print("Press 'q' in a window with image to abort")

filename_base = filename.split('.')[0]  # Assumes there is only one dot in filename, just before extension
res = (1920,1080)  # TODO: Determine resolution from input file

# give function that maps floats in range 0, 1
def lut_maker(foo):
    values = np.array([foo(x/256)*256 for x in range(256)])
    lut = np.array(np.clip(values, 0, 255), dtype = np.dtype('uint8'))
    return lut

# Splines as calculated by https://www.desmos.com/calculator/yvgfpq2prf where point 1 (0,0) and point 3 (1,1)
def spline_desmos(x, x2, y2):
    a1 = (x2 - y2)/(-2*x2*x2*(x2-1))
    b1 = (x2*x2 + x2*y2 - 2*y2)/(2*(-x2)*(-1)*(x2-1))

    if x <= x2:

        return a1 * (x ** 3) + b1 * x
    else:
        a2 = (x2/(x2-1))*a1
        b2 = (x2*(y2-3)+x2*x2+y2)/(2*(-x2)*(-1)*(x2-1))

        return a2 * ( (x-1) ** 3 ) + b2 * (x - 1) + 1

# The following were obtained with GIMP color curves tool
def spline_blue(x):
    return spline_desmos(
        x,
        0.69,
        0.63)

def spline_green(x):
    return spline_desmos(
        x,
        0.44,
        0.38)

def spline_red(x):
    return spline_desmos(
        x,
        0.38,
        0.63)

cap = cv2.VideoCapture(filename)
fourcc = cv2.VideoWriter_fourcc(*'DIVX')  # For Windows
out = cv2.VideoWriter(filename_base+'_nosound.avi', fourcc, 20.0, res)

identity = np.arange(256, dtype = np.dtype('uint8'))
zeros = np.zeros(256, np.dtype('uint8'))

blue_lut = lut_maker(spline_blue)
green_lut = lut_maker(spline_green)
red_lut = lut_maker(spline_red)

lut = np.dstack((blue_lut, green_lut, red_lut))

cv2.namedWindow('input')
cv2.moveWindow('input', 0, 0)
cv2.namedWindow('output')
cv2.moveWindow('output', int(res[0]/2), 0)

while True:
    ret, frame = cap.read()

    try:
        corrected = cv2.LUT(frame, lut)
    except:
        break

    out.write(corrected)

    # Display the resulting frame
    corrected = cv2.pyrDown(corrected)
    frame = cv2.pyrDown(frame)
    
    cv2.imshow('input',frame)
    cv2.imshow('output',corrected)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
out.release()
cv2.destroyAllWindows()

# Add sound from original video
system('ffmpeg -y -i {} -f wav -ab 192000 -vn {}.wav'.format(filename, filename_base))
system('ffmpeg -y -i {}_nosound.avi -i {}.wav -shortest -c:v copy -c:a aac -b:a 256k {}_corrected.avi'.format(filename_base, filename_base, filename_base))
remove('{}.wav'.format(filename_base))
remove('{}_nosound.avi'.format(filename_base))