ToneMap UE4+GT Plot

plot_tonemap_ue4_gt.py

import matplotlib
matplotlib.use('TkAgg')

from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from matplotlib.figure import Figure
import tkinter as tk
import tkinter.messagebox as tkmsg
import numpy as np

# https://github.com/ampas/aces-dev/blob/master/transforms/ctl/README-MATRIX.md

# XYZ_2_AP0_MAT = np.array([
#     [1.0498110175, 0.0, -0.0000974845],
#     [-0.4959030231, 1.3733130458, 0.0982400361],
#     [0.0, 0.0, 0.9912520182]])

# AP0_2_XYZ_MAT = np.array([
#  [ 0.952552396,  0.00000000,  0.0000936785927],
#  [ 0.343966450,  0.728166097, -0.0721325464],
#  [ 0.00000000,  0.00000000,  1.00882518]])

XYZ_2_AP1_MAT = np.array([
    [1.6410233797, -0.3248032942, -0.2364246952],
    [-0.6636628587,  1.6153315917,  0.0167563477],
    [0.0117218943, -0.0082844420,  0.9883948585]])

AP1_2_XYZ_MAT = np.array([
 [ 0.66245418,  0.13400421,  0.15618769],
 [ 0.27222872,  0.67408177,  0.05368952],
 [-0.00557465,  0.00406073, 1.0103391 ]])

AP1_RGB2Y = AP1_2_XYZ_MAT[1]

RGB_2_XYZ_MAT = np.array([
	[0.412391, 0.357585, 0.180482],
    [0.212639, 0.71517, 0.0721926],
    [0.0193308, 0.119195, 0.950536]])

XYZ_2_RGB_MAT = np.array(
 [[ 3.24096769, -1.53738183, -0.49861209],
 [-0.96924115,  1.87596362,  0.04155539],
 [ 0.05562988, -0.20397614,  1.0569672 ]])

D65_2_D60_CAT = np.array([
    [1.01303, 0.00610531, -0.014971],
    [0.00769823, 0.998165, -0.00503202],
    [-0.00284131, 0.00468513, 0.924507]])

D60_2_D65_CAT = np.array(
 [[ 0.9872288,  -0.0061133,   0.01595341],
 [-0.0075984,   1.00185983,  0.00533   ],
 [ 0.00307258, -0.00509592,  1.08167959]])

# AP0_2_AP1_MAT = np.dot(XYZ_2_AP1_MAT, AP0_2_XYZ_MAT)
# AP1_2_AP0_MAT = np.dot(XYZ_2_AP0_MAT, AP1_2_XYZ_MAT)

# AP1_2_AP0_MAT = np.array(
#  [[ 1.45143932, -0.23651075, -0.21492857],
#  [-0.07655377,  1.1762297,  -0.09967593],
#  [ 0.00831615, -0.00603245,  0.9977163 ]])

# AP0_2_AP1_MAT = np.array(
#  [[ 0.69545224,  0.1406787,   0.16386906],
#  [ 0.04479456,  0.85967112,  0.09553432],
#  [-0.00552588,  0.00402521,  1.00150067]])

# sRGB -> XYZ -> AP1
# sRGB_2_AP1_MAT = np.dot(XYZ_2_AP1_MAT, RGB_2_XYZ_MAT)

# AP1 -> XYZ -> sRGB
# AP1_2_sRGB_MAT = np.dot(XYZ_2_RGB_MAT, AP1_2_XYZ_MAT)

# D65(sRGB) -> D60(sRGB) -> XYZ -> AP1
sRGB_2_AP1_MAT = np.dot(XYZ_2_AP1_MAT, np.dot(RGB_2_XYZ_MAT, D65_2_D60_CAT))

# AP1 -> XYZ -> sRGB(D60) -> sRGB(D65)
AP1_2_sRGB_MAT = np.dot(D60_2_D65_CAT, np.dot(XYZ_2_RGB_MAT, AP1_2_XYZ_MAT))

# sRGB_2_AP1_MAT = np.array(
#  [[0.61334146, 0.32964333, 0.03370195],
#  [0.07807693, 0.91871792, 0.00612098],
#  [0.02068772, 0.12040973, 0.86906566]])

# AP1_2_sRGB_MAT = np.array(
#  [[ 1.7095552,  -0.60527189, -0.0620327 ],
#  [-0.14514882,  1.14086934, -0.00240654],
#  [-0.02058472, -0.14366012,  1.15247114]])

def log10(x):
    return np.log10(x)

def lerp(a,b,x):
    return a+(b-a)*x

def step(edge, x):
    return 1 if x>=edge else 0

def smoothstep(edge0, edge1, x):
    t = np.clip((x-edge0) / (edge1-edge0), 0, 1)
    return t*t*(3-2*t)

def tonemap_UE4(x, filmSlope=0.88, filmToe=0.55, filmShoulder=0.26, filmBlackClip=0.0, filmWhiteClip=0.04):
    """
    UE4 ToneMap.
    https://docs.unrealengine.com/en-us/Engine/Rendering/PostProcessEffects/ColorGrading

    Parameters
    ----------
    flimSlope : float
        This will adjust the steepness of the S-curve used for the tone mapper, where larger values will make the slope steeper (darker) and lower values will make the slope less steep (lighter). Value is in the range of [0.0, 1.0]. [0.0, 1.0]. 
    filmToe : float
        This will adjust the dark color in the tone mapper. Value is in the range of [0.0, 1.0]. [0.0, 1.0]. 
    filmShoulder : float
        This will adjust the bright color in the tone mapper. Value is in the range of [0.0, 1.0] 
    filmBlackClip : float
        This will set where the crossover happens where black's start to cut off their value. In general, this value should NOT be adjusted. Value is in the range of [0.0, 1.0] 
    flimWhiteClip : float
        This will set where the crossover happens where white's start to cut off their values. This will appear as a subtle change in most cases. Value is in the range of [0.0, 1.0]
    """
    if x <= 0:
        return 0
    
    toeScale = 1.0 + filmBlackClip - filmToe
    shoulderScale = 1.0 + filmWhiteClip - filmShoulder
    inMatch = 0.18
    outMatch = 0.18

    toeMatch = 0.0
    if filmToe > 0.8:
        # 0.18 will be on straight segment
        toeMatch = (1.0 - filmToe - outMatch) / filmSlope + log10(inMatch)
    else:
        # 0.18 will be on toe segment
        # Solve for toeMatch such that input of inMatch gives output of outMatch
        bt = (outMatch + filmBlackClip) / toeScale - 1.0
        toeMatch = log10(inMatch) - 0.5*np.log((1.0+bt)/(1.0-bt+1e-05))*(toeScale / filmSlope)
    
    straightMatch = (1.0 - filmToe) / filmSlope - toeMatch
    shoulderMatch = filmShoulder / filmSlope - straightMatch

    logColor = log10(x)
    straightColor = filmSlope * (logColor + straightMatch)

    toeColor = (-filmBlackClip) + (2.0*toeScale) / (1.0+np.exp((-2.0*filmSlope/toeScale)*(logColor - toeMatch)))
    shoulderColor = (1.0+filmWhiteClip) - (2.0*shoulderScale) / (1.0+np.exp(( 2.0*filmSlope/shoulderScale)*(logColor - shoulderMatch)))
    
    toeColor = toeColor if logColor < toeMatch else straightColor
    shoulderColor = shoulderColor if logColor > shoulderMatch else straightColor

    t = np.clip((logColor - toeMatch) / (shoulderMatch - toeMatch + 1e-05), 0, 1)
    t = 1.0-t if shoulderMatch < toeMatch else t
    t = (3.0-2.0*t)*t*t
    toneColor = lerp(toeColor, shoulderColor, t)

    return toneColor

def tonemap_GT(x, P=1.0, a=1.0, m=0.22, l=0.4, c=1.33, b=0.0):
    """
    GT Tonemap.
    Uchimura 2017, "HDR theory and practice"
    JP: https://www.desmos.com/calculator/mbkwnuihbd
    EN: https://www.desmos.com/calculator/gslcdxvipg

    Parameters
    ----------
    P : float
        max display brightness
    a : float
        contrast
    m : float
        linear section start
    l : float
        linear section length
    c : float
        black
    b : float
        pedestal
    """
    l0 = ((P-m)*l)/a
    # L0 = m-m/a
    # L1 = m+(1.0-m)/a
    S0 = m+l0
    S1 = m+a*l0
    C2 = (a*P)/(P-S1+1e-05)
    CP = -C2/P

    w0 = 1.0 - smoothstep(0.0, m, x)
    w2 = step(m+l0,x)
    w1 = 1.0 - w0 - w2

    T = m*((x/m)**c)+b
    S = P-(P-S1)*np.exp(CP*(x-S0))
    L = m+a*(x-m)

    return T*w0 + L*w1 + S*w2


class App:
    def __init__(self):
        pass
    
    def makeGraph(self, root):
        self.f = Figure(figsize=(20,6), dpi=100)
        self.a = self.f.add_subplot(111)
        self.canvas = FigureCanvasTkAgg(self.f, master=root)
        self.canvas.draw()
        self.canvas.get_tk_widget().pack(side=tk.TOP, fill=tk.BOTH, expand=1)
        self.canvas._tkcanvas.pack(side=tk.TOP, fill=tk.BOTH, expand=1)

    def onChanged(self, n):

        exposure  = self.uiExposure.get()
        slope     = self.uiSlope.get()
        toe       = self.uiToe.get()
        shoulder  = self.uiShoulder.get()
        blackClip = self.uiBlackClip.get()
        whiteClip = self.uiWhiteClip.get()
        GT_P      = self.uiMaximumBrightness.get()
        GT_a      = self.uiContrast.get()
        GT_m      = self.uiLinearSectionStart.get()
        GT_l      = self.uiLinearSectionLength.get()
        GT_c      = self.uiBlack.get()
        GT_b      = self.uiPedestal.get()

        self.a.cla()
        x = np.linspace(0, 2, 100)
        yUE4 = np.zeros(100)
        yGT = np.zeros(100)
        for i in range(100):
            linearColor = np.array([x[i]]*3) * exposure

            #------------------------------------------
            # UE4 ToneMap
            #------------------------------------------

            # sRGB -> AP1(ACEScg)
            t1 = np.max(np.dot(sRGB_2_AP1_MAT, linearColor),0)
            # Pre desaturate
            t2 = lerp(np.array([np.dot(linearColor, AP1_RGB2Y)]*3), t1, 0.96)
            # Tone Mapping
            t3 = np.array([
                tonemap_UE4(t2[0], slope, toe, shoulder, blackClip, whiteClip),
                tonemap_UE4(t2[1], slope, toe, shoulder, blackClip, whiteClip),
                tonemap_UE4(t2[2], slope, toe, shoulder, blackClip, whiteClip)])
            # Post desaturate
            toneColor = t3
            t4 = lerp(np.array([np.dot(toneColor, AP1_RGB2Y)]*3), t3, 0.93)
            # AP1(ACEScg) -> sRGB
            t5 = np.dot(AP1_2_sRGB_MAT, t4)
            yUE4[i] = t5[0]

            #------------------------------------------
            # GT ToneMap
            #------------------------------------------

            yGT[i] = tonemap_GT(linearColor[0], GT_P, GT_a, GT_m, GT_l, GT_c, GT_b)

        self.a.plot(x,yUE4, label="UE4")
        self.a.plot(x,yGT, label="GT")
        self.a.set_xlim(0, 2.0)
        self.a.set_ylim(0, 1.1)
        self.a.hlines(y=1,xmin=0,xmax=2,linewidths=0.5,linestyles='dashed')
        self.a.set_title('ToneMapping')
        self.a.legend(loc="upper left")
        self.canvas.draw()

    def makeUIScale(self, parent, label, row, from_, to, value):
        uiLabel = tk.Label(parent, text=label)
        uiLabel.grid(row=row, column=0, sticky=tk.E)
        uiScale = tk.Scale(parent, orient=tk.HORIZONTAL, from_=from_, to=to, resolution=0.01, command=self.onChanged)
        uiScale.grid(row=row, column=1)
        uiScale.set(value)
        return (uiLabel, uiScale)

    def makeUI(self, parent):
        row = 0

        self.uiExposureLabel, self.uiExposure = self.makeUIScale(parent, ' Exposure: ', row, 0.0, 10.0, 1.0)
        row += 1

        uiLabel = tk.Label(parent, text="--- UE4 ---")
        uiLabel.grid(row=row, column=0, columnspan=2)
        row += 1
        self.uiSlopeLabel, self.uiSlope = self.makeUIScale(parent, ' Slope: ', row, 0.0, 1.0, 0.88)
        row += 1
        self.uiToeLabel, self.uiToe = self.makeUIScale(parent, ' Toe: ', row, 0.0, 1.0, 0.55)
        row += 1
        self.uiShoulderLabel, self.uiShoulder = self.makeUIScale(parent, ' Shoulder: ', row, 0.0, 1.0, 0.26)
        row += 1
        self.uiBlackClipLabel, self.uiBlackClip = self.makeUIScale(parent, ' BlackClip: ', row, 0.0, 1.0, 0.0)
        row += 1
        self.uiWhiteLabel, self.uiWhiteClip = self.makeUIScale(parent, ' WhiteClip: ', row, 0.0, 1.0, 0.02)
        row += 1

        uiLabel = tk.Label(parent, text="--- GT ---")
        uiLabel.grid(row=row, column=0, columnspan=2)
        row += 1

        self.uiMaximumBrightnessLabel, self.uiMaximumBrightness = self.makeUIScale(parent, ' P: ', row, 1.0, 2.0, 1.0)
        row += 1
        self.uiContrastLabel, self.uiContrast = self.makeUIScale(parent, ' a: ', row, 0.0, 5.0, 1.0)
        row += 1
        self.uiLinearSectionStartLabel, self.uiLinearSectionStart = self.makeUIScale(parent, ' m: ', row, 0.0, 1.0, 0.22)
        row += 1
        self.uiLinearSectionLengthLabel, self.uiLinearSectionLength = self.makeUIScale(parent, ' l: ', row, 0.0, 1.0, 0.4)
        row += 1
        self.uiBlackLabel, self.uiBlack = self.makeUIScale(parent, ' c: ', row, 1.0, 3.0, 1.33)
        row += 1
        self.uiPedestalLabel, self.uiPedestal = self.makeUIScale(parent, ' b: ', row, 0.0, 1.0, 0.0)
        row += 1
    
    def run(self):
        self.root = tk.Tk()
        self.root.title('python + tkinter + matplotlib')
        self.root.geometry("1000x600")
        self.lpanel = tk.Frame(self.root)
        self.lpanel.pack(side=tk.LEFT, fill=tk.BOTH)
        self.makeUI(self.lpanel)
        self.rpanel = tk.Frame(self.root)
        self.rpanel.pack(side=tk.RIGHT, fill=tk.BOTH)
        self.makeGraph(self.rpanel)
        self.root.mainloop()

if __name__ == "__main__":
    app = App()
    app.run()