handle rgb and postprocessing images

ZeoDarkflame · Sep 8, 2020 · ecd74fc · ecd74fc
1 parent b549ec1
commit ecd74fc
Show file tree

Hide file tree

Showing 6 changed files with 203 additions and 6 deletions.
diff --git a/flightlib/include/flightlib/bridges/unity_bridge.hpp b/flightlib/include/flightlib/bridges/unity_bridge.hpp
@@ -39,7 +39,7 @@ class UnityBridge {
 
   // public get functions
   bool getRender(const FrameID frame_id);
-  bool handleOutput(RenderMessage_t &output);
+  bool handleOutput();
   bool getPointCloud(PointCloudMessage_t &pointcloud_msg);
 
   // public set functions

diff --git a/flightlib/include/flightlib/sensors/rgb_camera.hpp b/flightlib/include/flightlib/sensors/rgb_camera.hpp
@@ -41,6 +41,10 @@ class RGBCamera : SensorBase {
   int getHeight(void) const;
   Scalar getFOV(void) const;
   Scalar getDepthScale(void) const;
+  bool getRGBImage(cv::Mat& rgb_img);
+  bool getDepthMap(cv::Mat& depth_map);
+  bool getSegmentation(cv::Mat& segmentation);
+  bool getOpticalFlow(cv::Mat& opticalflow);
 
   // auxiliary functions
   void enableDepth(const bool on);

diff --git a/flightlib/src/bridges/unity_bridge.cpp b/flightlib/src/bridges/unity_bridge.cpp
@@ -183,7 +183,7 @@ bool UnityBridge::addQuadrotor(std::shared_ptr<Quadrotor> quad) {
   return true;
 }
 
-bool UnityBridge::handleOutput(RenderMessage_t& output) {
+bool UnityBridge::handleOutput() {
   // create new message object
   zmqpp::message msg;
   sub_.receive(msg);
@@ -192,10 +192,41 @@ bool UnityBridge::handleOutput(RenderMessage_t& output) {
   // parse metadata
   SubMessage_t sub_msg = json::parse(json_sub_msg);
 
+  size_t image_i = 1;
   // ensureBufferIsAllocated(sub_msg);
   for (size_t idx = 0; idx < settings_.vehicles.size(); idx++) {
     // update vehicle collision flag
     unity_quadrotors_[idx]->setCollision(sub_msg.sub_vehicles[idx].collision);
+
+    // feed image data to RGB camera
+    for (const auto& cam : settings_.vehicles[idx].cameras) {
+      for (size_t layer_idx = 0; layer_idx <= cam.enabled_layers.size();
+           layer_idx++) {
+        if (!layer_idx == 0 && !cam.enabled_layers[layer_idx - 1]) continue;
+        uint32_t image_len = cam.width * cam.height * cam.channels;
+        // Get raw image bytes from ZMQ message.
+        // WARNING: This is a zero-copy operation that also casts the input to
+        // an array of unit8_t. when the message is deleted, this pointer is
+        // also dereferenced.
+        const uint8_t* image_data;
+        msg.get(image_data, image_i);
+        image_i = image_i + 1;
+        // Pack image into cv::Mat
+        cv::Mat new_image =
+          cv::Mat(cam.height, cam.width, CV_MAKETYPE(CV_8U, cam.channels));
+        memcpy(new_image.data, image_data, image_len);
+        // Flip image since OpenCV origin is upper left, but Unity's is lower
+        // left.
+        cv::flip(new_image, new_image, 0);
+
+        // Tell OpenCv that the input is RGB.
+        if (cam.channels == 3) {
+          cv::cvtColor(new_image, new_image, CV_RGB2BGR);
+        }
+        unity_quadrotors_[idx]->getCameras()[cam.output_index]->feedImageQueue(
+          layer_idx, new_image);
+      }
+    }
   }
   return true;
 }

diff --git a/flightlib/src/envs/vec_env.cpp b/flightlib/src/envs/vec_env.cpp
@@ -98,7 +98,7 @@ bool VecEnv<EnvBase>::step(Ref<MatrixRowMajor<>> act, Ref<MatrixRowMajor<>> obs,
 
   if (unity_render_ && unity_ready_) {
     unity_bridge_ptr_->getRender(0);
-    unity_bridge_ptr_->handleOutput(unity_output_);
+    unity_bridge_ptr_->handleOutput();
   }
   return true;
 }

diff --git a/flightlib/src/sensors/rgb_camera.cpp b/flightlib/src/sensors/rgb_camera.cpp
@@ -4,8 +4,8 @@ namespace flightlib {
 
 RGBCamera::RGBCamera()
   : channels_(3),
-    width_(480),
-    height_(720),
+    width_(720),
+    height_(480),
     fov_{70.0},
     depth_scale_{0.2},
     enabled_layers_({false, false, false}) {}
@@ -145,4 +145,40 @@ void RGBCamera::enableOpticalFlow(const bool on) {
   enabled_layers_[CameraLayer::OpticalFlow] = on;
 }
 
+bool RGBCamera::getRGBImage(cv::Mat& rgb_img) {
+  if (!rgb_queue_.empty()) {
+    rgb_img = rgb_queue_.front();
+    rgb_queue_.pop_front();
+    return true;
+  }
+  return false;
+}
+
+bool RGBCamera::getDepthMap(cv::Mat& depth_map) {
+  if (!depth_queue_.empty()) {
+    depth_map = depth_queue_.front();
+    depth_queue_.pop_front();
+    return true;
+  }
+  return false;
+}
+
+bool RGBCamera::getSegmentation(cv::Mat& segmentation) {
+  if (!segmentation_queue_.empty()) {
+    segmentation = segmentation_queue_.front();
+    segmentation_queue_.pop_front();
+    return true;
+  }
+  return false;
+}
+
+bool RGBCamera::getOpticalFlow(cv::Mat& opticalflow) {
+  if (!opticalflow_queue_.empty()) {
+    opticalflow = opticalflow_queue_.front();
+    opticalflow_queue_.pop_front();
+    return true;
+  }
+  return false;
+}
+
 }  // namespace flightlib
diff --git a/flightlib/tests/bridges/unity_bridge.cpp b/flightlib/tests/bridges/unity_bridge.cpp
@@ -19,6 +19,8 @@ TEST(UnityBridge, Constructor) {
 
   // if (unity_ready) logger.info("Unity Rendering is connected");
   // EXPECT_TRUE(unity_ready);
+  // //timeout flightmare
+  // usleep(5 * 1e6);
 }
 
 TEST(UnityBridge, PointCloud) {
@@ -37,5 +39,129 @@ TEST(UnityBridge, PointCloud) {
   // EXPECT_TRUE(unity_bridge.getPointCloud(pointcloud_msg));
   // std::experimental::filesystem::remove(pointcloud_msg.path +
   //                                       pointcloud_msg.file_name + ".ply");
-  // unity_bridge.disconnectUnity();
+  // //timeout flightmare
+  // usleep(5 * 1e6);
+}
+
+TEST(UnityBridge, HandleOutputRGB) {
+  Logger logger{"Test HandleOutputRGB"};
+  UnityBridge unity_bridge;
+  // QuadrotorDynamics dyn = QuadrotorDynamics(1.0, 0.2);
+  // std::shared_ptr<Quadrotor> quad = std::make_shared<Quadrotor>(dyn);
+  // std::shared_ptr<RGBCamera> rgb = std::make_shared<RGBCamera>();
+  // quad->addRGBCamera(rgb);
+  // unity_bridge.addQuadrotor(quad);
+
+  // bool unity_ready = unity_bridge.connectUnity(UnityScene::GARAGE);
+
+  // EXPECT_TRUE(unity_ready);
+
+  // FrameID frame_id = 1;
+  // unity_bridge.getRender(frame_id);
+  // bool handle_output = unity_bridge.handleOutput();
+
+  // EXPECT_TRUE(handle_output);
+
+  // cv::Mat test_img;
+
+  // EXPECT_TRUE(rgb->getRGBImage(test_img));
+  // EXPECT_FALSE(rgb->getDepthMap(test_img));
+  // EXPECT_FALSE(rgb->getSegmentation(test_img));
+  // EXPECT_FALSE(rgb->getOpticalFlow(test_img));
+
+  // // timeout flightmare
+  // usleep(5 * 1e6);
+}
+
+TEST(UnityBridge, HandleOutputDepth) {
+  Logger logger{"Test HandleOutputDepth"};
+  UnityBridge unity_bridge;
+  // QuadrotorDynamics dyn = QuadrotorDynamics(1.0, 0.2);
+  // std::shared_ptr<Quadrotor> quad = std::make_shared<Quadrotor>(dyn);
+  // std::shared_ptr<RGBCamera> rgb = std::make_shared<RGBCamera>();
+  // rgb->setPostPrecesscing(std::vector<bool>{true, false, false});
+  // quad->addRGBCamera(rgb);
+  // unity_bridge.addQuadrotor(quad);
+
+  // bool unity_ready = unity_bridge.connectUnity(UnityScene::GARAGE);
+
+  // EXPECT_TRUE(unity_ready);
+
+  // FrameID frame_id = 1;
+  // unity_bridge.getRender(frame_id);
+  // bool handle_output = unity_bridge.handleOutput();
+
+  // EXPECT_TRUE(handle_output);
+
+  // cv::Mat test_img;
+
+  // EXPECT_TRUE(rgb->getRGBImage(test_img));
+  // EXPECT_TRUE(rgb->getDepthMap(test_img));
+  // EXPECT_FALSE(rgb->getSegmentation(test_img));
+  // EXPECT_FALSE(rgb->getOpticalFlow(test_img));
+
+  // // timeout flightmare
+  // usleep(5 * 1e6);
+}
+
+TEST(UnityBridge, HandleOutputSegmentation) {
+  Logger logger{"Test HandleOutputSegmentation"};
+  UnityBridge unity_bridge;
+  // QuadrotorDynamics dyn = QuadrotorDynamics(1.0, 0.2);
+  // std::shared_ptr<Quadrotor> quad = std::make_shared<Quadrotor>(dyn);
+  // std::shared_ptr<RGBCamera> rgb = std::make_shared<RGBCamera>();
+  // rgb->setPostPrecesscing(std::vector<bool>{false, true, false});
+  // quad->addRGBCamera(rgb);
+  // unity_bridge.addQuadrotor(quad);
+
+  // bool unity_ready = unity_bridge.connectUnity(UnityScene::GARAGE);
+
+  // EXPECT_TRUE(unity_ready);
+
+  // FrameID frame_id = 1;
+  // unity_bridge.getRender(frame_id);
+  // bool handle_output = unity_bridge.handleOutput();
+
+  // EXPECT_TRUE(handle_output);
+
+  // cv::Mat test_img;
+
+  // EXPECT_TRUE(rgb->getRGBImage(test_img));
+  // EXPECT_FALSE(rgb->getDepthMap(test_img));
+  // EXPECT_TRUE(rgb->getSegmentation(test_img));
+  // EXPECT_FALSE(rgb->getOpticalFlow(test_img));
+
+  // // timeout flightmare
+  // usleep(5 * 1e6);
+}
+
+TEST(UnityBridge, HandleOutputOpticalFlow) {
+  Logger logger{"Test HandleOutputOpticalFlow"};
+  UnityBridge unity_bridge;
+  // QuadrotorDynamics dyn = QuadrotorDynamics(1.0, 0.2);
+  // std::shared_ptr<Quadrotor> quad = std::make_shared<Quadrotor>(dyn);
+  // std::shared_ptr<RGBCamera> rgb = std::make_shared<RGBCamera>();
+  // rgb->setPostPrecesscing(std::vector<bool>{false, false, true});
+  // quad->addRGBCamera(rgb);
+  // unity_bridge.addQuadrotor(quad);
+
+  // bool unity_ready = unity_bridge.connectUnity(UnityScene::GARAGE);
+
+  // EXPECT_TRUE(unity_ready);
+
+  // FrameID frame_id = 1;
+  // unity_bridge.getRender(frame_id);
+  // bool handle_output = unity_bridge.handleOutput();
+
+  // EXPECT_TRUE(handle_output);
+
+  // cv::Mat test_img;
+
+  // EXPECT_TRUE(rgb->getRGBImage(test_img));
+  // EXPECT_FALSE(rgb->getDepthMap(test_img));
+  // EXPECT_FALSE(rgb->getSegmentation(test_img));
+  // EXPECT_TRUE(rgb->getOpticalFlow(test_img));
+
+  // // timeout flightmare
+  // usleep(5 * 1e6);
 }