Add a simple colour MLS

applications/reconstruct/src/main.cpp

@@ -254,8 +254,8 @@ static void run(ftl::Configurable *root) {
 	pipeline1->append<ftl::operators::ColourChannels>("colour");  // Convert BGR to BGRA
 	pipeline1->append<ftl::operators::HFSmoother>("hfnoise");  // Remove high-frequency noise
 	pipeline1->append<ftl::operators::Normals>("normals");  // Estimate surface normals
-	pipeline1->append<ftl::operators::SmoothChannel>("smoothing");  // Generate a smoothing channel
-	pipeline1->append<ftl::operators::SimpleMLS>("mls");  // Perform MLS (using smoothing channel)
+	//pipeline1->append<ftl::operators::SmoothChannel>("smoothing");  // Generate a smoothing channel
+	pipeline1->append<ftl::operators::ColourMLS>("mls");  // Perform MLS (using smoothing channel)
 	// Alignment
 
 

components/operators/include/ftl/operators/smoothing.hpp

@@ -61,6 +61,22 @@ class SimpleMLS : public ftl::operators::Operator {
     private:
 };
 
+/**
+ * Perform Moving Least Squares smoothing with a smoothing amount determined
+ * by a simple colour similarity weighting.
+ */
+class ColourMLS : public ftl::operators::Operator {
+    public:
+    explicit ColourMLS(ftl::Configurable*);
+    ~ColourMLS();
+
+	inline Operator::Type type() const override { return Operator::Type::OneToOne; }
+
+    bool apply(ftl::rgbd::Frame &in, ftl::rgbd::Frame &out, ftl::rgbd::Source *src, cudaStream_t stream) override;
+
+    private:
+};
+
 /**
  * A version of Moving Least Squares where both the smoothing factor and
  * neighbourhood size are adapted over an extra large range using the colour

components/operators/src/smoothing.cpp

@@ -5,6 +5,7 @@
 
 using ftl::operators::HFSmoother;
 using ftl::operators::SimpleMLS;
+using ftl::operators::ColourMLS;
 using ftl::operators::SmoothChannel;
 using ftl::codecs::Channel;
 using cv::cuda::GpuMat;
@@ -120,3 +121,46 @@ bool SimpleMLS::apply(ftl::rgbd::Frame &in, ftl::rgbd::Frame &out, ftl::rgbd::So
 	return true;
 }
 
+
+
+ColourMLS::ColourMLS(ftl::Configurable *cfg) : ftl::operators::Operator(cfg) {
+
+}
+
+ColourMLS::~ColourMLS() {
+
+}
+
+bool ColourMLS::apply(ftl::rgbd::Frame &in, ftl::rgbd::Frame &out, ftl::rgbd::Source *s, cudaStream_t stream) {
+	float thresh = config()->value("mls_threshold", 0.04f);
+	float col_smooth = config()->value("mls_colour_smoothing", 30.0f);
+	int iters = config()->value("mls_iterations", 1);
+	int radius = config()->value("mls_radius",2);
+
+	if (!in.hasChannel(Channel::Normals)) {
+		LOG(ERROR) << "Required normals channel missing for MLS";
+		return false;
+	}
+
+	// FIXME: Assume in and out are the same frame.
+	for (int i=0; i<iters; ++i) {
+		ftl::cuda::colour_mls_smooth(
+			in.createTexture<float4>(Channel::Normals),
+			in.createTexture<float4>(Channel::Points, ftl::rgbd::Format<float4>(in.get<cv::cuda::GpuMat>(Channel::Depth).size())),
+			in.createTexture<float>(Channel::Depth),
+			in.createTexture<float>(Channel::Depth2, ftl::rgbd::Format<float>(in.get<cv::cuda::GpuMat>(Channel::Depth).size())),
+			in.createTexture<uchar4>(Channel::Colour),
+			thresh,
+			col_smooth,
+			radius,
+			s->parameters(),
+			0
+		);
+
+		in.swapChannels(Channel::Depth, Channel::Depth2);
+		in.swapChannels(Channel::Normals, Channel::Points);
+	}
+
+	return true;
+}
+

components/operators/src/smoothing.cu

@@ -99,6 +99,109 @@ void ftl::cuda::mls_smooth(
 	#endif
 }
 
+
+// ===== Colour MLS Smooth =====================================================
+
+/*
+ * Smooth depth map using Moving Least Squares
+ */
+ template <int SEARCH_RADIUS>
+ __global__ void colour_mls_smooth_kernel(
+		TextureObject<float4> normals_in,
+		TextureObject<float4> normals_out,
+        TextureObject<float> depth_in,        // Virtual depth map
+		TextureObject<float> depth_out,   // Accumulated output
+		TextureObject<uchar4> colour_in,
+		float smoothing,
+		float colour_smoothing,
+        ftl::rgbd::Camera camera) {
+        
+    const int x = blockIdx.x*blockDim.x + threadIdx.x;
+    const int y = blockIdx.y*blockDim.y + threadIdx.y;
+
+    if (x < 0 || y < 0 || x >= depth_in.width() || y >= depth_in.height()) return;
+
+	float3 aX = make_float3(0.0f,0.0f,0.0f);
+	float3 nX = make_float3(0.0f,0.0f,0.0f);
+    float contrib = 0.0f;
+
+	float d0 = depth_in.tex2D(x, y);
+	depth_out(x,y) = d0;
+	if (d0 < camera.minDepth || d0 > camera.maxDepth) return;
+	float3 X = camera.screenToCam((int)(x),(int)(y),d0);
+
+	uchar4 c0 = colour_in.tex2D(x, y);
+
+    // Neighbourhood
+    for (int v=-SEARCH_RADIUS; v<=SEARCH_RADIUS; ++v) {
+    for (int u=-SEARCH_RADIUS; u<=SEARCH_RADIUS; ++u) {
+		const float d = depth_in.tex2D(x+u, y+v);
+		if (d < camera.minDepth || d > camera.maxDepth) continue;
+
+		// Point and normal of neighbour
+		const float3 Xi = camera.screenToCam((int)(x)+u,(int)(y)+v,d);
+		const float3 Ni = make_float3(normals_in.tex2D((int)(x)+u, (int)(y)+v));
+
+		const uchar4 c = colour_in.tex2D(x+u, y+v);
+		const float cw = ftl::cuda::colourWeighting(c0,c,colour_smoothing);
+
+		// Gauss approx weighting function using point distance
+		const float w = ftl::cuda::spatialWeighting(X,Xi,smoothing*cw);
+
+		aX += Xi*w;
+		nX += Ni*w;
+		contrib += w;
+    }
+	}
+	
+	nX /= contrib;  // Weighted average normal
+	aX /= contrib;  // Weighted average point (centroid)
+
+	// Signed-Distance Field function
+	float fX = nX.x * (X.x - aX.x) + nX.y * (X.y - aX.y) + nX.z * (X.z - aX.z);
+
+	// Calculate new point using SDF function to adjust depth (and position)
+	X = X - nX * fX;
+	
+	//uint2 screen = camera.camToScreen<uint2>(X);
+
+    //if (screen.x < depth_out.width() && screen.y < depth_out.height()) {
+    //    depth_out(screen.x,screen.y) = X.z;
+	//}
+	depth_out(x,y) = X.z;
+	normals_out(x,y) = make_float4(nX / length(nX), 0.0f);
+}
+
+void ftl::cuda::colour_mls_smooth(
+		ftl::cuda::TextureObject<float4> &normals_in,
+		ftl::cuda::TextureObject<float4> &normals_out,
+		ftl::cuda::TextureObject<float> &depth_in,
+		ftl::cuda::TextureObject<float> &depth_out,
+		ftl::cuda::TextureObject<uchar4> &colour_in,
+		float smoothing,
+		float colour_smoothing,
+		int radius,
+		const ftl::rgbd::Camera &camera,
+		cudaStream_t stream) {
+
+	const dim3 gridSize((depth_out.width() + T_PER_BLOCK - 1)/T_PER_BLOCK, (depth_out.height() + T_PER_BLOCK - 1)/T_PER_BLOCK);
+	const dim3 blockSize(T_PER_BLOCK, T_PER_BLOCK);
+
+	switch (radius) {
+		case 5: colour_mls_smooth_kernel<5><<<gridSize, blockSize, 0, stream>>>(normals_in, normals_out, depth_in, depth_out, colour_in, smoothing, colour_smoothing, camera); break;
+		case 4: colour_mls_smooth_kernel<4><<<gridSize, blockSize, 0, stream>>>(normals_in, normals_out, depth_in, depth_out, colour_in, smoothing, colour_smoothing, camera); break;
+		case 3: colour_mls_smooth_kernel<3><<<gridSize, blockSize, 0, stream>>>(normals_in, normals_out, depth_in, depth_out, colour_in, smoothing, colour_smoothing, camera); break;
+		case 2: colour_mls_smooth_kernel<2><<<gridSize, blockSize, 0, stream>>>(normals_in, normals_out, depth_in, depth_out, colour_in, smoothing, colour_smoothing, camera); break;
+		case 1: colour_mls_smooth_kernel<1><<<gridSize, blockSize, 0, stream>>>(normals_in, normals_out, depth_in, depth_out, colour_in, smoothing, colour_smoothing, camera); break;
+	}
+	cudaSafeCall( cudaGetLastError() );
+
+
+	#ifdef _DEBUG
+	cudaSafeCall(cudaDeviceSynchronize());
+	#endif
+}
+
 // ===== Colour-based Smooth ===================================================
 
 template <int RADIUS>

components/operators/src/smoothing_cuda.hpp

@@ -17,6 +17,18 @@ void mls_smooth(
 		const ftl::rgbd::Camera &camera,
 		cudaStream_t stream);
 
+void colour_mls_smooth(
+		ftl::cuda::TextureObject<float4> &normals_in,
+		ftl::cuda::TextureObject<float4> &normals_out,
+		ftl::cuda::TextureObject<float> &depth_in,
+		ftl::cuda::TextureObject<float> &depth_out,
+		ftl::cuda::TextureObject<uchar4> &colour_in,
+		float smoothing,
+		float colour_smoothing,
+		int radius,
+		const ftl::rgbd::Camera &camera,
+		cudaStream_t stream);
+
 void depth_smooth(
 	ftl::cuda::TextureObject<float> &depth_in,
 	ftl::cuda::TextureObject<uchar4> &colour_in,