HDCAM dehalo experiments

dehdcam.py

from vskernels import Bilinear, Box, Catrom, Hermite
from vsscale import ArtCNN
from vstools import vs, core, depth, limiter, set_output, get_y, get_u, get_v, join

def disc(n):
	return (core.dgdecodenv.DGSource(f'YESPRECURE5_DISC{n}.dgi')
		.std.SetFrameProps(_Matrix=vs.MATRIX_BT709, _Transfer=vs.TRANSFER_BT709, _Primaries=vs.PRIMARIES_BT709)
		.std.SetFieldBased(vs.FIELD_PROGRESSIVE)
		.resize.Point(format=vs.YUV420PS)
	)

ep14 = disc(3)[43520:86978][3388-2788:]

def upsample_chroma(clip):
	return ArtCNN.R8F64_Chroma(chroma_scaler=Bilinear).scale(clip)

def light(clip):
	clip = clip.resize.Point(format=clip.format.replace(color_family=vs.RGB))
	return clip.ort.Model('net_g_22000_fp32.onnx', provider='CUDA', fp16=True)

def descale(clip, taps):
	scaling_factor = 0.75 if clip.height == 1080 else 0.5
	return clip.descale.Delanczos(width=clip.width*scaling_factor, height=clip.height, src_left=(1-scaling_factor)/2, taps=taps, border_handling=1)

def scale(clip, taps):
	scaling_factor = 0.75 if clip.height == 1080 else 0.5
	return clip.descale.Lanczos(width=clip.width/scaling_factor, height=clip.height, src_left=(1-scaling_factor)/2, taps=taps, border_handling=1)

def horizontal_artcnn(clip):
	return ArtCNN.R16F96(scaler=Box).scale(clip, width=clip.width*2, height=clip.height)

def horizontal_nnedi3(clip):
	return clip.std.Transpose().sneedif.NNEDI3(field=1, dh=True, nsize=3, nns=4, qual=1).std.Transpose()

def upscale_artcnn_nnedi3(clip):
	descaled_y = descale(get_y(clip), taps=7)
	descaled_u = descale(get_u(clip), taps=5)
	descaled_v = descale(get_v(clip), taps=5)

	rescaled_y = scale(descaled_y, taps=7)
	rescaled_u = scale(descaled_u, taps=5)
	rescaled_v = scale(descaled_v, taps=5)

	descaled_u = descaled_u.akarin.Expr('x 0.5 +')
	descaled_v = descaled_v.akarin.Expr('x 0.5 +')

	artcnn_y = horizontal_artcnn(descaled_y)
	artcnn_uv = horizontal_artcnn(core.std.Interleave([descaled_u, descaled_v])).akarin.Expr('x 0.5 -')
	artcnn_u = artcnn_uv[0::2]
	artcnn_v = artcnn_uv[1::2]

	nnedi3_y = horizontal_nnedi3(descaled_y)
	nnedi3_uv = horizontal_nnedi3(core.std.Interleave([descaled_u, descaled_v])).akarin.Expr('x 0.5 -')
	nnedi3_u = nnedi3_uv[0::2]
	nnedi3_v = nnedi3_uv[1::2]

	artcnn_final_y = Hermite.scale(artcnn_y, width=clip.width, height=clip.height, src_left=0.25)
	nnedi3_final_y = Hermite.scale(nnedi3_y, width=clip.width, height=clip.height, src_left=-0.25)

	artcnn_final_u = Catrom.scale(artcnn_u, src_left=0.5)
	artcnn_final_v = Catrom.scale(artcnn_v, src_left=0.5)

	rescaled = join(rescaled_y, rescaled_u, rescaled_v)
	artcnn = join(artcnn_final_y, artcnn_final_u, artcnn_final_v)
	nnedi3 = join(nnedi3_final_y, nnedi3_u, nnedi3_v)
	mixed = join(artcnn_final_y, nnedi3_u, nnedi3_v)
	mixed2 = join(artcnn_final_y, nnedi3_u.std.Merge(artcnn_final_u), nnedi3_v.std.Merge(artcnn_final_v))

	return rescaled, artcnn, nnedi3, mixed, mixed2

rescaled, artcnn, nnedi3, mixed, mixed2 = upscale_artcnn_nnedi3(ep14)

yuv420 = core.std.Interleave([
	ep14,
	rescaled,
	artcnn,
	nnedi3,
	mixed,
	mixed2,
])
yuv444 = upsample_chroma(yuv420)
ep14     = yuv444[0::6]
rescaled = yuv444[1::6]
artcnn   = yuv444[2::6]
nnedi3   = yuv444[3::6]
mixed    = yuv444[4::6]
mixed2   = yuv444[5::6]
set_output(ep14)
set_output(rescaled, "ep14 (same-kernel rescale)")
set_output(light(ep14), "ep14 (net_g_22000_fp32)")
set_output(artcnn, "ep14 (ArtCNN/Hermite, ArtCNN/Catrom)")
set_output(nnedi3, "ep14 (nnedi3/Hermite, nnedi3)")
set_output(mixed, "ep14 (ArtCNN/Hermite, nnedi3)")
set_output(mixed2, "ep14 (ArtCNN/Hermite, ArtCNN/Catrom+nnedi3)")