MCTemporalDenoise port attempt

mctd.py

import vapoursynth as vs
import havsfunc as haf
import math

def mvd(clip, c_super, fv_list, bv_list, chroma=False, thsad=300, thscd1=150, thscd2=150):
	core = vs.get_core()
	MDegrain1 = core.mv.Degrain1
	MDegrain2 = core.mv.Degrain2
	MDegrain3 = core.mv.Degrain3
	Merge = core.std.Merge
	
	radius = len(fv_list)
	common_params = dict(thsad=thsad, thscd1=thscd1, thscd2=thscd2, plane=4 if chroma else 0)
	if radius == 1:
		f1v, b1v = fv_list[0], bv_list[0]
		sm = MDegrain1(clip, c_super, b1v, f1v, **common_params)
	elif radius == 2:
		f1v, f2v = fv_list
		b1v, b2v = bv_list
		sm = MDegrain2(clip, c_super, b1v, f1v, b2v, f2v, **common_params)
	elif radius == 3:
		f1v, f2v, f3v = fv_list
		b1v, b2v, b3v = bv_list
		sm = MDegrain3(clip, c_super, b1v, f1v, b2v, f2v, b3v, f3v, **common_params)
	elif radius == 4:
		f1v, f2v, f3v, f4v = fv_list
		b1v, b2v, b3v, b4v = bv_list
		mv12 = MDegrain2(clip, c_super, b1v, f1v, b2v, f2v, **common_params)
		mv34 = MDegrain2(clip, c_super, b3v, f3v, b4v, f4v, **common_params)
		sm = Merge(mv12, mv34, 0.4444)
	elif radius == 5:
		f1v, f2v, f3v, f4v, f5v = fv_list
		b1v, b2v, b3v, b4v, b5v = bv_list
		mv123 = MDegrain3(clip, c_super, b1v, f1v, b2v, f2v, b3v, f3v, **common_params)
		mv45 = MDegrain2(clip, c_super, b4v, f4v, b5v, f5v, **common_params)
		sm = Merge(mv123, mv45, 0.4545)
	elif radius == 6:
		f1v, f2v, f3v, f4v, f5v, f6v = fv_list
		b1v, b2v, b3v, b4v, b5v, b6v = bv_list
		mv123 = MDegrain3(clip, c_super, b1v, f1v, b2v, f2v, b3v, f3v, **common_params)
		mv456 = MDegrain3(clip, c_super, b4v, f4v, b5v, f5v, b6v, f6v, **common_params)
		sm = Merge(mv123, mv456, 0.4615)
	else:
		raise RuntimeError("radius must be in [0, 6]")

	return sm

def defaults(preset):
	if preset == "very low":
		return dict(
			radius=1,
			sigma=2,
			limit=-1,
			chroma=False,
			bwbh=16,
			owoh=8,
			blksize=16,
			overlap=8,
			bt=1,
			thsad=200,
			thscd1=200,
			thscd2=90,
			truemotion=False,
			mv_global=True,
			pel=1,
			pelsearch=1,
			search=2,
			searchparam=2,
			mv_sharp=2
		)
	elif preset == "low":
		return dict(
			radius=2,
			sigma=4,
			limit=-1,
			chroma=False,
			bwbh=16,
			owoh=8,
			blksize=16,
			overlap=8,
			bt=3,
			thsad=300,
			thscd1=300,
			thscd2=100,
			truemotion=False,
			mv_global=True,
			pel=2,
			pelsearch=2,
			search=2,
			searchparam=2,
			mv_sharp=2
		)
	elif preset == "medium":
		return dict(
			radius=3,
			sigma=8,
			limit=-1,
			chroma=True,
			bwbh=16,
			owoh=8,
			blksize=16,
			overlap=8,
			bt=3,
			thsad=400,
			thscd1=400,
			thscd2=100,
			truemotion=False,
			mv_global=True,
			pel=2,
			pelsearch=2,
			search=2,
			searchparam=2,
			mv_sharp=2
		)
	elif preset == "high":
		return dict(
			radius=2,
			sigma=12,
			limit=-1,
			chroma=True,
			bwbh=16,
			owoh=8,
			blksize=16,
			overlap=8,
			bt=3,
			thsad=500,
			thscd1=500,
			thscd2=130,
			truemotion=False,
			mv_global=True,
			pel=2,
			pelsearch=2,
			search=2,
			searchparam=2,
			mv_sharp=1
		)
	elif preset == "very high":
		return dict(
			radius=3,
			sigma=8,
			limit=-1,
			chroma=True,
			bwbh=16,
			owoh=8,
			blksize=16,
			overlap=8,
			bt=4,
			thsad=500,
			thscd1=500,
			thscd2=130,
			truemotion=False,
			mv_global=True,
			pel=2,
			pelsearch=2,
			search=2,
			searchparam=2,
			mv_sharp=0
		)
	else:
		raise RuntimeError("incorrect preset {preset}".format(preset=preset))

def mctd(clip, radius, sigma, limit, chroma, bwbh, owoh, blksize, overlap, bt, thsad, thscd1, thscd2, truemotion, mv_global, pel, pelsearch, search, searchparam, mv_sharp):
	core = vs.get_core()
	dfttest = core.dfttest.DFTTest
	MAnalyse = core.mv.Analyse
	MSuper = core.mv.Super

	planes = (0, 1, 2) if chroma else (0,)	
	bitdepth = clip.format.bits_per_sample
	mod = bwbh if bwbh >= blksize else blksize
	w = clip.width
	h = clip.height
	wf = math.ceil(w/mod) * mod - w + mod
	hf = math.ceil(h/mod) * mod - h + mod
	wn = int(w + wf)
	hn = int(h + hf)
	
	i = haf.Resize(clip, wn, hn, -wf/2, -hf/2, wn, hn)

	# prefiltered clip
	p = dfttest(i, sigma=sigma,
				ftype=0,
				smode=1, sbsize=bwbh, sosize=owoh,
				tmode=0, tbsize=bt,
				planes=planes)
	d = i  # deblock (required?)

	### PREPARING
	p_super = MSuper(p, hpad=0, vpad=0, pel=pel, sharp=mv_sharp, chroma=chroma)
	
	# Common arguments
	m_analyse_kw = dict(
		truemotion=truemotion, blksize=blksize, overlap=overlap,
		pelsearch=pelsearch, search=search, searchparam=searchparam, chroma=chroma
	)
	m_analyse_kw["global"] = mv_global  # hack due to Python's reserved kw
	fv_list = [
		MAnalyse(p_super, delta=i+1, isb=False, **m_analyse_kw)
		for i in range(radius)
	]
	bv_list = [
		MAnalyse(p_super, delta=i+1, isb=True, **m_analyse_kw)
		for i in range(radius)
	]
	
	d_super = MSuper(d, hpad=0, vpad=0, pel=pel, sharp=mv_sharp, chroma=chroma)
	sm = mvd(d, d_super, fv_list=fv_list, bv_list=bv_list, chroma=chroma,
				  thsad=thsad, thscd1=thscd1, thscd2=thscd2)
				  
	p_diff = core.std.MakeDiff(i, p, planes=planes)
	sm_diff = core.std.MakeDiff(i, sm, planes=planes)
	if limit == -1:
		expr = "x {thr} - abs y {thr} - abs < x y ?".format(thr=haf.scale(128, bitdepth))
		d_diff = core.std.Expr((p_diff, sm_diff), expr=expr)
	elif limit > 0:
		expr = "x y - abs {thr} <= x x y - 0 < y {thr} - y {thr} + ? ?".format(thr=limit)
		d_diff = core.std.Expr((sm, i), expr=expr)
	else:
		d_diff = sm
	
	sm_limited = core.std.MakeDiff(i, d_diff, planes=planes) if limit == -1 else d_diff
	
	# TODO: Second pass
	# TODO: Postprocessing?
	# TODO: interlaced?
	
	out = core.std.CropRel(sm_limited, wf//2, hf//2, wf//2, hf//2)

	return out