rgbd_camera: Added option to save stereo images to directory or side-by-side avi file

tools/CameraRGBD/main.cpp

@@ -31,6 +31,9 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include "rtabmap/core/util3d_transforms.h"
 #include "rtabmap/utilite/ULogger.h"
 #include "rtabmap/utilite/UMath.h"
+#include "rtabmap/utilite/UFile.h"
+#include "rtabmap/utilite/UDirectory.h"
+#include "rtabmap/utilite/UConversion.h"
 #include <opencv2/highgui/highgui.hpp>
 #include <opencv2/imgproc/imgproc.hpp>
 #include <pcl/visualization/cloud_viewer.h>
@@ -39,7 +42,7 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 void showUsage()
 {
 	printf("\nUsage:\n"
-			"rtabmap-rgbd_camera driver\n"
+			"rtabmap-rgbd_camera [options] driver\n"
 			"  driver       Driver number to use: 0=OpenNI-PCL (Kinect)\n"
 			"                                     1=OpenNI2    (Kinect and Xtion PRO Live)\n"
 			"                                     2=Freenect   (Kinect)\n"
@@ -47,10 +50,21 @@ void showUsage()
 			"                                     4=OpenNI-CV-ASUS (Xtion PRO Live)\n"
 			"                                     5=Freenect2  (Kinect v2)\n"
 			"                                     6=DC1394     (Bumblebee2)\n"
-			"                                     7=FlyCapture2 (Bumblebee2)\n");
+			"                                     7=FlyCapture2 (Bumblebee2)\n"
+			"  Options:\n"
+			"      -rate #.#                      Input rate Hz (default 0=inf)\n"
+			"      -save_stereo \"path\"            Save stereo images in a folder or a video file (side by side *.avi).\n");
 	exit(1);
 }
 
+// catch ctrl-c
+bool running = true;
+void sighandler(int sig)
+{
+	printf("\nSignal %d caught...\n", sig);
+	running = false;
+}
+
 int main(int argc, char * argv[])
 {
 	ULogger::setType(ULogger::kTypeConsole);
@@ -59,74 +73,120 @@ int main(int argc, char * argv[])
 	//ULogger::setPrintWhere(false);
 
 	int driver = 0;
+	std::string stereoSavePath;
+	float rate = 0.0f;
 	if(argc < 2)
 	{
 		showUsage();
 	}
 	else
 	{
-		if(strcmp(argv[argc-1], "--help") == 0)
+		for(int i=1; i<argc; ++i)
 		{
-			showUsage();
+			if(strcmp(argv[i], "-rate") == 0)
-		}
+			{
-		driver = atoi(argv[argc-1]);
+				++i;
-		if(driver < 0 || driver > 7)
+				if(i < argc)
-		{
+				{
-			UERROR("driver should be between 0 and 6.");
+					rate = uStr2Float(argv[i]);
-			showUsage();
+					if(rate < 0.0f)
+					{
+						showUsage();
+					}
+				}
+				else
+				{
+					showUsage();
+				}
+				continue;
+			}
+			if(strcmp(argv[i], "-save_stereo") == 0)
+			{
+				++i;
+				if(i < argc)
+				{
+					stereoSavePath = argv[i];
+				}
+				else
+				{
+					showUsage();
+				}
+				continue;
+			}
+			if(strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-help") == 0)
+			{
+				showUsage();
+			}
+
+			// last
+			driver = atoi(argv[i]);
+			if(driver < 0 || driver > 7)
+			{
+				UERROR("driver should be between 0 and 6.");
+				showUsage();
+			}
 		}
 	}
 	UINFO("Using driver %d", driver);
 
 	rtabmap::Camera * camera = 0;
-	if(driver == 0)
+	if(driver < 6)
 	{
-		camera = new rtabmap::CameraOpenni();
+		if(!stereoSavePath.empty())
-	}
-	else if(driver == 1)
-	{
-		if(!rtabmap::CameraOpenNI2::available())
 		{
-			UERROR("Not built with OpenNI2 support...");
+			UWARN("-save_stereo option cannot be used with RGB-D drivers.");
-			exit(-1);
+			stereoSavePath.clear();
 		}
-		camera = new rtabmap::CameraOpenNI2();
+
-	}
+		if(driver == 0)
-	else if(driver == 2)
-	{
-		if(!rtabmap::CameraFreenect::available())
 		{
-			UERROR("Not built with Freenect support...");
+			camera = new rtabmap::CameraOpenni();
-			exit(-1);
 		}
-		camera = new rtabmap::CameraFreenect();
+		else if(driver == 1)
-	}
-	else if(driver == 3)
-	{
-		if(!rtabmap::CameraOpenNICV::available())
 		{
-			UERROR("Not built with OpenNI from OpenCV support...");
+			if(!rtabmap::CameraOpenNI2::available())
-			exit(-1);
+			{
+				UERROR("Not built with OpenNI2 support...");
+				exit(-1);
+			}
+			camera = new rtabmap::CameraOpenNI2();
 		}
-		camera = new rtabmap::CameraOpenNICV(false);
+		else if(driver == 2)
-	}
-	else if(driver == 4)
-	{
-		if(!rtabmap::CameraOpenNICV::available())
 		{
-			UERROR("Not built with OpenNI from OpenCV support...");
+			if(!rtabmap::CameraFreenect::available())
-			exit(-1);
+			{
+				UERROR("Not built with Freenect support...");
+				exit(-1);
+			}
+			camera = new rtabmap::CameraFreenect();
 		}
-		camera = new rtabmap::CameraOpenNICV(true);
+		else if(driver == 3)
-	}
-	else if(driver == 5)
-	{
-		if(!rtabmap::CameraFreenect2::available())
 		{
-			UERROR("Not built with Freenect2 support...");
+			if(!rtabmap::CameraOpenNICV::available())
-			exit(-1);
+			{
+				UERROR("Not built with OpenNI from OpenCV support...");
+				exit(-1);
+			}
+			camera = new rtabmap::CameraOpenNICV(false);
+		}
+		else if(driver == 4)
+		{
+			if(!rtabmap::CameraOpenNICV::available())
+			{
+				UERROR("Not built with OpenNI from OpenCV support...");
+				exit(-1);
+			}
+			camera = new rtabmap::CameraOpenNICV(true);
+		}
+		else if(driver == 5)
+		{
+			if(!rtabmap::CameraFreenect2::available())
+			{
+				UERROR("Not built with Freenect2 support...");
+				exit(-1);
+			}
+			camera = new rtabmap::CameraFreenect2(0, rtabmap::CameraFreenect2::kTypeRGBDepthSD);
 		}
-		camera = new rtabmap::CameraFreenect2(0, rtabmap::CameraFreenect2::kTypeRGBDepthSD);
 	}
 	else if(driver == 6)
 	{
@@ -157,21 +217,69 @@ int main(int argc, char * argv[])
 		delete camera;
 		exit(1);
 	}
+
 	rtabmap::SensorData data = camera->takeImage();
 	if(data.imageRaw().cols != data.depthOrRightRaw().cols || data.imageRaw().rows != data.depthOrRightRaw().rows)
 	{
 		UWARN("RGB (%d/%d) and depth (%d/%d) frames are not the same size! The registered cloud cannot be shown.",
 				data.imageRaw().cols, data.imageRaw().rows, data.depthOrRightRaw().cols, data.depthOrRightRaw().rows);
 	}
+	pcl::visualization::CloudViewer * viewer = 0;
 	if(!data.stereoCameraModel().isValid() && (data.cameraModels().size() == 0 || !data.cameraModels()[0].isValid()))
 	{
 		UWARN("Camera not calibrated! The registered cloud cannot be shown.");
 	}
-	pcl::visualization::CloudViewer viewer("cloud");
+	else
+	{
+		viewer = new pcl::visualization::CloudViewer("cloud");
+	}
 	rtabmap::Transform t(1, 0, 0, 0,
 						 0, -1, 0, 0,
 						 0, 0, -1, 0);
-	while(!data.imageRaw().empty() && !viewer.wasStopped())
+
+	cv::VideoWriter videoWriter;
+	UDirectory dir;
+	if(!stereoSavePath.empty() &&
+	   !data.imageRaw().empty() &&
+	   !data.rightRaw().empty())
+	{
+		if(UFile::getExtension(stereoSavePath).compare("avi") == 0)
+		{
+			if(data.imageRaw().size() == data.rightRaw().size())
+			{
+				if(rate <= 0)
+				{
+					UERROR("You should set the input rate when saving stereo images to a video file.");
+					showUsage();
+				}
+				cv::Size targetSize = data.imageRaw().size();
+				targetSize.width *= 2;
+				videoWriter.open(stereoSavePath, CV_FOURCC('M', 'J', 'P', 'G'), rate, targetSize, data.imageRaw().channels() == 3);
+			}
+			else
+			{
+				UERROR("Images not the same size, cannot save stereo images to the video file.");
+			}
+		}
+		else if(UDirectory::exists(stereoSavePath))
+		{
+			UDirectory::makeDir(stereoSavePath+"/"+"left");
+			UDirectory::makeDir(stereoSavePath+"/"+"right");
+		}
+		else
+		{
+			UERROR("Directory \"%s\" doesn't exist.", stereoSavePath.c_str());
+			stereoSavePath.clear();
+		}
+	}
+
+	// to catch the ctrl-c
+	signal(SIGABRT, &sighandler);
+	signal(SIGTERM, &sighandler);
+	signal(SIGINT, &sighandler);
+
+	int id=1;
+	while(!data.imageRaw().empty() && (viewer==0 || !viewer->wasStopped()) && running)
 	{
 		cv::Mat rgb = data.imageRaw();
 		if(!data.depthRaw().empty() && (data.depthRaw().type() == CV_16UC1 || data.depthRaw().type() == CV_32FC1))
@@ -194,7 +302,8 @@ int main(int argc, char * argv[])
 						data.cameraModels()[0].fx(),
 						data.cameraModels()[0].fy());
 				cloud = rtabmap::util3d::transformPointCloud(cloud, t);
-				viewer.showCloud(cloud, "cloud");
+				if(viewer)
+					viewer->showCloud(cloud, "cloud");
 			}
 			else if(!depth.empty() &&
 					data.cameraModels().size() &&
@@ -207,7 +316,7 @@ int main(int argc, char * argv[])
 						data.cameraModels()[0].fx(),
 						data.cameraModels()[0].fy());
 				cloud = rtabmap::util3d::transformPointCloud(cloud, t);
-				viewer.showCloud(cloud, "cloud");
+				viewer->showCloud(cloud, "cloud");
 			}
 
 			cv::Mat tmp;
@@ -238,7 +347,8 @@ int main(int argc, char * argv[])
 						data.stereoCameraModel().left().fx(),
 						data.stereoCameraModel().baseline());
 				cloud = rtabmap::util3d::transformPointCloud(cloud, t);
-				viewer.showCloud(cloud, "cloud");
+				if(viewer)
+					viewer->showCloud(cloud, "cloud");
 			}
 		}
 
@@ -246,8 +356,50 @@ int main(int argc, char * argv[])
 		if(c == 27)
 			break; // if ESC, break and quit
 
+		if(videoWriter.isOpened())
+		{
+			cv::Mat left = data.imageRaw();
+			cv::Mat right = data.rightRaw();
+			if(left.size() == right.size())
+			{
+				cv::Size targetSize = left.size();
+				targetSize.width *= 2;
+				cv::Mat targetImage(targetSize, left.type());
+				if(right.type() != left.type())
+				{
+					cv::Mat tmp;
+					cv::cvtColor(right, tmp, left.channels()==3?CV_GRAY2BGR:CV_BGR2GRAY);
+					right = tmp;
+				}
+				UASSERT(left.type() == right.type());
+
+				cv::Mat roiA(targetImage, cv::Rect( 0, 0, left.size().width, left.size().height ));
+				left.copyTo(roiA);
+				cv::Mat roiB( targetImage, cvRect( left.size().width, 0, left.size().width, left.size().height ) );
+				right.copyTo(roiB);
+
+				videoWriter.write(targetImage);
+				printf("Saved frame %d to \"%s\"\n", id, stereoSavePath.c_str());
+			}
+			else
+			{
+				UERROR("Left and right images are not the same size!?");
+			}
+		}
+		else if(!stereoSavePath.empty())
+		{
+			cv::imwrite(stereoSavePath+"/"+"left/"+uNumber2Str(id) + ".jpg", data.imageRaw());
+			cv::imwrite(stereoSavePath+"/"+"right/"+uNumber2Str(id) + ".jpg", data.rightRaw());
+			printf("Saved frames %d to \"%s/left\" and \"%s/right\" directories\n", id, stereoSavePath.c_str(), stereoSavePath.c_str());
+		}
+		++id;
 		data = camera->takeImage();
 	}
+	printf("Closing...\n");
+	if(viewer)
+	{
+		delete viewer;
+	}
 	cv::destroyWindow("Video");
 	cv::destroyWindow("Depth");
 	delete camera;