Delowpass with DFT

delowpass.py

# Full-frequency source (US)
full = depth(full, 32)
# Low-frequency source (JP)
blurred = depth(blurred, 32)

full = full.resize.Point(format=vs.YUV444PS, matrix_s='170m')
blurred = blurred.resize.Point(format=vs.YUV444PS, matrix_s='170m')

full = get_y(full)
blurred = get_y(blurred)

def lowpass(clip):
	clip = clip.fmtc.resample(clip.width, clip.height, kernel="lanczos", taps=4, fh=1/1.375)
	clip = clip.fmtc.resample(clip.width, clip.height, kernel="lanczos", taps=4, fh=1/1.25)
	return clip

full_lpf = lowpass(full)
blurred_lpf = lowpass(blurred)
# Impulse response from the low-pass filter, placed in the middle of the frame
impulse_lpf = lowpass(full.akarin.Expr(f'X {full.width // 2} ='))

def delowpass(n, f):
	dst = f[0].copy()

	for iplane in range(dst.format.num_planes):
		full, full_lpf, blurred, blurred_lpf, impulse_lpf = f

		full = np.asarray(full[iplane])
		full_lpf = np.asarray(full_lpf[iplane])
		blurred = np.asarray(blurred[iplane])
		blurred_lpf = np.asarray(blurred_lpf[iplane])
		impulse_lpf = np.asarray(impulse_lpf[iplane])

		height, width = full.shape

		# If noise amplitude ratio is known and fixed:
		# ratio_sqr = ratio ** 2

		a = abs(np.fft.rfft(np.fft.ifftshift(impulse_lpf[0]))) ** 2
		b = np.ones(width // 2 + 1)
		eqns = np.column_stack((a, b))
		vals = (abs(np.fft.rfft2(blurred - full_lpf)[1:]) ** 2).mean(0)
		# This can surely be optimized by inlining the lstsq
		full_s2, blurred_s2 = np.linalg.lstsq(eqns, vals, rcond=None)[0]
		ratio_sqr = full_s2 / blurred_s2

		# Mirror-pad the images to account for mirror padding during low-pass filtering
		full = np.fft.rfft(np.hstack((full, np.fliplr(full))))
		blurred_lpf = np.fft.rfft(np.hstack((blurred_lpf, np.fliplr(blurred_lpf))))
		impulse_lpf = np.fft.rfft(np.hstack([impulse_lpf[:, width//2:], np.zeros_like(impulse_lpf), impulse_lpf[:, :width//2]]))

		# Alternatively: zero-pad for some semblance of zero-extend border handling in the low-pass filtering.
		# `blurred` should actually have some nonzero data in it that's been cropped off during the original filtering,
		# so this is not fully accurate, but we don't know any better and this seems to improve the left edge of the image.
		# full = np.fft.rfft(full, width * 2)
		# blurred_lpf = np.fft.rfft(blurred_lpf, width * 2)

		# Without padding:
		# full = np.fft.rfft(full)
		# blurred_lpf = np.fft.rfft(blurred_lpf)
		# impulse_lpf = np.fft.rfft(np.fft.ifftshift(impulse_lpf, axes=1))

		# Least-squares/least-variance solution to:
		#     x * impulse_lpf + noise = blurred
		#     x + noise * ratio = full
		# The LPF's impulse response is symmetric, so its Fourier transform is real,
		# which makes the abs() redundant in a world of infinite-precision mathematics;
		# but finite computing precision may cause it to have a tiny imaginary part.
		dft = (ratio_sqr * blurred_lpf + full) / (ratio_sqr * abs(impulse_lpf) ** 2 + 1)

		out = np.fft.irfft(dft, width * 2)[:, :width]

		# Without padding:
		# out = np.fft.irfft(dft, width)

		np.copyto(np.asarray(dst[iplane]), out)

	return dst

set_output(full.std.ModifyFrame([full, full_lpf, blurred, blurred_lpf, impulse_lpf], delowpass))