#####################################################################
# #
# Gray_Balance v.1.0, by Kisa_AG, 2013 #
# #
# A gray balance script, which is based on Black/White/Gray #
# Balance picker. #
# #
#####################################################################

# --- LATEST CHANGES ---
#
# v1.0
# - First Release

# --- REQUIREMENTS ---
#
# Input colorspaces: YV12, YUY2, RGB24
# Ouput colorspace: YV12 in working mode, RGB24 in TestMode.
# Avisynth 2.5.8 ot 2.6
#
# Core plugins:
# SmoothAdjust.dll (2.6 or above)
#

# --- GETTING STARTED ---
#
# Gray_Balance(clip, Float "RG", Float "GG", Float "BG", Float "R0", Float "G0", Float "B0", Float "RW", Float "GW",
# Float "BW", bool "Interlaced", string "Matrix_In", string "Matrix_Out", bool "TestMode")
# Clip - The clip to be processed
# RG, GG, BG - RGB values of the object in clip, which should be neutral gray.
# After processing of clip, this point will have equal RGB values R=G=B=(RG+GG+BG)/3 (but corrected according
# to custom black and white point settings in order to maintain relative contrast).
# Default values: RG=GG=BG=128
#
# R0, G0, B0 - custom RGB values of black point of the clip.
# After processing of clip, this point will have equal RGB values R=G=B=16 (if TVRange=True), or R=G=B=0 (if TVRange=False).
# Default values: R0=G0=B0=0
#
# RW, GW, BW - custom RGB values of white point of the clip.
# After processing of clip, this point will have equal RGB values R=G=B=235 (if TVRange=True), or R=G=B=255 (if TVRange=False).
# Default values: RW=GW=BW=255
#
# Interlaced - to perform interlaced color convertion.
# Default value: Interlaced=False
#
# Matrix_In - input matrix for convertion from YV12 to RGB24. Valid values are PC.709, Rec709, PC.601 and Rec601.
# When using PC.709 and PC.601, it's assumed that source has full range [0,255].
# When using Rec709 and Rec601, it's assumed that source has TV range [16,235].
# Default value: Matrix_In=PC.709
#
# Matrix_Out - output matrix for convertion from RGB24 to YV12. Valid values are PC.709, Rec709, PC.601 and Rec601.
# When using PC.709 and PC.601, then output will use full range [0,255].
# When using Rec709 and Rec601, then output will be SCALED to TV range [16,235] (no clipping).
# Default value: Matrix_Out=PC.709
#
# TestMode - is the mode for finding RGB values of the object in clip, which should be neutral gray.
# Ouput colorspace is switched to RGB24, any other variables are ignored.
# Default value: TestMode=False

# Usig of Gray_Balance script has two stages.
########### Stage 1 ############
# write a script like this:
# YourSource("yourfile")
# Gray_Balance(TestMode=True)
# Save this script with an ".avs" extension.
#
# You can even use Gray_Balance(131,127,113,TestMode=True) or any other string, because everything is ignored while TestMode=True.
# Then open this .AVS script in VirtualDub, find a frame with an object which should be gray. Copy this frame to clipboard (Ctrl+1).
# Go to any image edidting application (Photoshop, GIMP, whatever),
# paste the frame, choose the picker tool and grab the RGB values for the object, which should be gray.
# (It's better to use the object in semitones, i.e. it's RGB values should be, let's say, somewhere in between 80-180).
# Another approach is to use any screen picker like ColorCop to grab RGB values directly from VirtualDub.
# It seems that the resualt should be absolutely the same.

########### Stage 2 ############
# Open script in Notepad again, remove TestMode=True and enter RGB values grabbed in the first stage
# as RG, GG and BG. E.g. Gray_Balance(131,127,113).
#
# You can do the same to set Black and White points of the clip, if you want.
# Now you can open your .AVS script in VirtualDub to check what happend. If everything was done correctly,
# then you should see your gray object really perfectly gray. The overall gray balance should also become more correct.
# Curves approximation is based on Quadratic Polynomials.

######################################################
function Gray_Balance(clip clp, Float "RG", Float "GG", Float "BG", Float "R0", Float "G0", Float "B0", Float "RW", Float "GW", Float "BW", bool "Interlaced", string "Matrix_In", string "Matrix_Out", bool "TestMode")
{

TestMode = default(TestMode, false) # TestMode Default
Interlaced = default(Interlaced, False) # Interlaced Default
Matrix_In = default(Matrix_In, "PC.709") # Matrix_In Default
Matrix_Out = default(Matrix_Out, "PC.709") # Matrix_Out Default
RG = default(RG, 128) # RG Default
GG = default(GG, 128) # GG Default
BG = default(BG, 128) # BG Default
R0 = default(R0, 0) # R0 Default
G0 = default(G0, 0) # G0 Default
B0 = default(B0, 0) # B0 Default
RW = default(RW, 255) # RW Default
GW = default(GW, 255) # GW Default
BW = default(BW, 255) # BW Default

(Interlaced == True && Matrix_In == "Rec709") ? Clp.ConvertToRGB24(matrix="Rec709", Interlaced=True) : NOP
(Interlaced == True && Matrix_In == "PC.709") ? Clp.ConvertToRGB24(matrix="PC.709", Interlaced=True) : NOP
(Interlaced == True && Matrix_In == "Rec601") ? Clp.ConvertToRGB24(matrix="Rec601", Interlaced=True) : NOP
(Interlaced == True && Matrix_In == "PC.601") ? Clp.ConvertToRGB24(matrix="PC.601", Interlaced=True) : NOP
(Interlaced == False && Matrix_In == "Rec709") ? Clp.ConvertToRGB24(matrix="Rec709") : NOP
(Interlaced == False && Matrix_In == "PC.709") ? Clp.ConvertToRGB24(matrix="PC.709") : NOP
(Interlaced == False && Matrix_In == "Rec601") ? Clp.ConvertToRGB24(matrix="Rec601") : NOP
(Interlaced == False && Matrix_In == "PC.601") ? Clp.ConvertToRGB24(matrix="PC.601") : NOP

(TestMode == true) ? Eval("""NOP()
""") : Eval("""

# New RGB for Black Point
R0n=0
G0n=0
B0n=

# New RGB for White Point
RWn=255
GWn=255
BWn=255

RGB_W_AVG=(RW+GW+BW)/3
RGB_0_AVG=(R0+G0+B0)/3
RGB_Wn_AVG=(RWn+GWn+BWn)/3
RGB_0n_AVG=(R0n+G0n+B0n)/3

# New RGB for Gray Object
#RGB_Gray_AVG=(RG+GG+BG)/3 #Old Style calc

RGB_Gray_AVG=(RG+GG+BG)/3

RGB_Gray_AVG=RGB_0n_AVG+(RGB_Gray_AVG-RGB_0_AVG)*(RGB_Wn_AVG-RGB_0n_AVG)/(RGB_W_AVG-RGB_0_AVG)

RGn=RGB_Gray_AVG
GGn=RGB_Gray_AVG
BGn=RGB_Gray_AVG

# Coefficients should be calculated in the following sequence ONLY: B, A, C!!!

B_R=Find_B_1(R0,RG,RW,R0n,RGn,RWn)
A_R=Find_A_2(R0,RG,RW,R0n,RGn,RWn,B_R)
C_R=Find_C_3(R0,RG,RW,R0n,RGn,RWn,A_R,B_R)

B_G=Find_B_1(G0,GG,GW,G0n,GGn,GWn)
A_G=Find_A_2(G0,GG,GW,G0n,GGn,GWn,B_G)
C_G=Find_C_3(G0,GG,GW,G0n,GGn,GWn,A_G,B_G)

B_B=Find_B_1(B0,BG,BW,B0n,BGn,BWn)
A_B=Find_A_2(B0,BG,BW,B0n,BGn,BWn,B_B)
C_B=Find_C_3(B0,BG,BW,B0n,BGn,BWn,A_B,B_B)

# Applying correction
red = showred("YV12")
green = showgreen("YV12")
blue = showblue("YV12")

red2 = SmoothCustom(red, Yexpr="Y*Y*"+string(A_R)+"+Y*"+string(B_R)+"+"+string(C_R), smooth=100, dither=50)
green2 = SmoothCustom(green, Yexpr="Y*Y*"+string(A_G)+"+Y*"+string(B_G)+"+"+string(C_G), smooth=100, dither=50)
blue2 = SmoothCustom(blue, Yexpr="Y*Y*"+string(A_B)+"+Y*"+string(B_B)+"+"+string(C_B), smooth=100, dither=50)

mergergb(red2,green2,blue2)

(Interlaced == True && Matrix_Out == "Rec709") ? ConvertToYV12(matrix="Rec709", Interlaced=True) : NOP
(Interlaced == True && Matrix_Out == "PC.709") ? ConvertToYV12(matrix="PC.709", Interlaced=True) : NOP
(Interlaced == True && Matrix_Out == "Rec601") ? ConvertToYV12(matrix="Rec601", Interlaced=True) : NOP
(Interlaced == True && Matrix_Out == "PC.601") ? ConvertToYV12(matrix="PC.601", Interlaced=True) : NOP
(Interlaced == False && Matrix_Out == "Rec709") ? ConvertToYV12(matrix="Rec709") : NOP
(Interlaced == False && Matrix_Out == "PC.709") ? ConvertToYV12(matrix="PC.709") : NOP
(Interlaced == False && Matrix_Out == "Rec601") ? ConvertToYV12(matrix="Rec601") : NOP
(Interlaced == False && Matrix_Out == "PC.601") ? ConvertToYV12(matrix="PC.601") : NOP

function Find_B_1(float X0, float X1, float X2, float Y0, float Y1, float Y2)
{
# Calculate B coefficient of square function
B0=((y1-y0)*(x2*x2-x0*x0)-(y2-y0)*(x1*x1-x0*x0))/((x1-x0)*(x2*x2-x0*x0)-(x2-x0)*(x1*x1-x0*x0))
Return(B0)
}

function Find_A_2(float X0, float X1, float X2, float Y0, float Y1, float Y2, Float B0)
{
# Calculate A coefficient of square function
A0=(y2-y0-B0*(x2-x0))/(x2*x2-x0*x0)
Return(A0)
}

function Find_C_3(float X0, float X1, float X2, float Y0, float Y1, float Y2, Float A0, Float B0)
{
# Calculate C coefficient of square function
C0=y0-a0*x0*x0-b0*x0
Return(C0)
}

""")

Return(Last)
}