I suggest looking into Halide as it will make trying different paths much easier (https://halide-lang.org/).

I haven't looked at your code closely so can't say with certainty it would be the right fit but worth a look.