Anaglyph to SBS toy converter

a2sbs.py

#!/usr/bin/env python3

#################################
# Anaglyph to SBS toy converter #
#################################

import cv2
import numpy as np
import argparse

# Argument parsing helper
def fminmax(a, b):
    def mm(x):
        x = float(x)
        if x < a:
            raise argparse.ArgumentTypeError('Minimum is ' + a)
        if x > b:
            raise argparse.ArgumentTypeError('Maximym is ' + b)
        return x
    return mm

# Blur helper
def colors_blur(img, nb):
    wf = int(w * o.colors_blur) & ~1
    for i in range(nb):
        img = cv2.GaussianBlur(img, (wf+1, 3), wf, None, 3)
    return (img * o.colors_dim).astype(np.uint8)

# Parse command line arguments
parser = argparse.ArgumentParser()
parser.add_argument('image', help='Path of anaglyph image to convert')
parser.add_argument('output', default=None, help='Path of result image', nargs='?')
parser.add_argument('-colors_flow', '-cf', action='store_true', help='Restore colors with flow')
parser.add_argument('-colors', '-c', action='store_true', help='Hack colors')
parser.add_argument('-colors_dim', '-cd', default=1, type=fminmax(0,1), help='Dim colors')
parser.add_argument('-colors_blur', '-cb', default=0.1, type=fminmax(0,1), help='Blur colors percentage size')
parser.add_argument('-cross', '-x', '-rl', action='store_true', help='Output as cross side-by-side format')
parser.add_argument('-xtalk', '-ghosting', default=0, type=fminmax(0,100), help='Try to reduce crosstalk (ghosting)')
parser.add_argument('-xtalk_delta', '-ghosting_delta', default=0, type=fminmax(0,100), help='Try to reduce crosstalk (ghosting) asynmetry')
o = parser.parse_args()

# Separate image hack
image = cv2.imread(o.image)
b, g, r = cv2.split(image)
il = cv2.cvtColor(cv2.merge((r, r, r)), cv2.COLOR_BGR2LAB)[:,:,0]
ir = cv2.cvtColor(cv2.merge((b, g, g)), cv2.COLOR_BGR2LAB)[:,:,0]
h, w, _ = image.shape

# Manage crosstalk
xtalk = o.xtalk/100.0
xtalkd = o.xtalk_delta/100.0
if xtalk > 0:
    i = (il * (1 - xtalk - xtalkd) + ir * (xtalk + xtalkd)).astype(np.uint8)
    ir = (ir * (1 - xtalk + xtalkd) + il * (xtalk - xtalkd)).astype(np.uint8)
    il = i

# Colors hack
# (could crosstalk be taken into account here too?)
if o.colors and not o.colors_flow:
    colors = colors_blur(image, 5)
    colors = cv2.cvtColor(colors, cv2.COLOR_BGR2LAB)

    colors[:,:,0] = il
    il = cv2.cvtColor(colors, cv2.COLOR_LAB2BGR)

    colors[:,:,0] = ir
    ir = cv2.cvtColor(colors, cv2.COLOR_LAB2BGR)

# Restore colors with flow
if o.colors_flow:
    DIS = cv2.DISOpticalFlow_create()
    DIS.setFinestScale(0)
    DIS.setPatchStride(1)
    DIS.setPatchSize(2)
    DIS.setGradientDescentIterations(10)
    DIS.setVariationalRefinementIterations(15)
    DIS.setVariationalRefinementDelta(1)
    DIS.setVariationalRefinementGamma(0)
    DIS.setVariationalRefinementAlpha(3)

    flow = DIS.calc(ir, il, None)
    s = flow.shape[:2]
    mapx = np.fromfunction(lambda i, j: j, s, dtype=np.float32) + flow[:,:,0]
    mapy = np.fromfunction(lambda i, j: i, s, dtype=np.float32) + flow[:,:,1]
    rw = cv2.remap(r, mapx, mapy, cv2.INTER_CUBIC, borderValue=0, borderMode=cv2.BORDER_CONSTANT)
    colors = colors_blur(cv2.merge((b, g, rw)), 5)
    colors = cv2.cvtColor(colors, cv2.COLOR_BGR2LAB)
    colors[:,:,0] = il
    ilo = il
    il = cv2.cvtColor(colors, cv2.COLOR_LAB2BGR)

    flow = DIS.calc(ilo, ir, None)
    s = flow.shape[:2]
    mapx = np.fromfunction(lambda i, j: j, s, dtype=np.float32) + flow[:,:,0]
    mapy = np.fromfunction(lambda i, j: i, s, dtype=np.float32) + flow[:,:,1]
    bw = cv2.remap(b, mapx, mapy, cv2.INTER_CUBIC, borderValue=0, borderMode=cv2.BORDER_CONSTANT)
    gw = cv2.remap(g, mapx, mapy, cv2.INTER_CUBIC, borderValue=0, borderMode=cv2.BORDER_CONSTANT)
    colors = colors_blur(cv2.merge((bw, gw, r)), 5)
    colors = cv2.cvtColor(colors, cv2.COLOR_BGR2LAB)
    colors[:,:,0] = ir
    ir = cv2.cvtColor(colors, cv2.COLOR_LAB2BGR)

# Merge into SBS
result = np.hstack((ir, il)) if o.cross else np.hstack((il, ir))

# Save
if o.output:
    cv2.imwrite(o.output, result)
else:
    cv2.imshow('a2sbs', result)
    cv2.waitKey(-1)
    cv2.destroyAllWindows()