Merge pull request #21 from uzh-rpg/dev/unity_bridge

Dev/unity bridge

LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2018 Robotics and Perception Group, 
+Copyright (c) 2020 Robotics and Perception Group, 
 University of Zurich, Switzerland
 
 Permission is hereby granted, free of charge, to any person obtaining a copy

flightlib/CMakeLists.txt

@@ -8,6 +8,7 @@ message(STATUS "======================  !Flightmare!  ======================")
 # Options
 ################################################################################
 option(BUILD_TESTS "Building the tests" ON)
+option(BUILD_UNITY_BRIDGE_TESTS "Building the Unity Bridge tests" ON)
 option(BUILD_BENCH "Building the benchmark." OFF)
 option(ENABLE_FAST "Build with optimizations for speed" ON)
 option(ENABLE_BLAS "Build using BLAS and LAPACK libraries" OFF)
@@ -136,7 +137,7 @@ message(STATUS "The activated CXX DEBUG configuration is:\n ${CMAKE_CXX_FLAGS_DE
 ################################################################################
 message(STATUS "======>  Setup Build  ")
 
-# Create file lists for flightlib
+# Create file lists for flightlib source
 file(GLOB_RECURSE FLIGHTLIB_SOURCES
   src/bridges/*.cpp
   src/dynamics/*.cpp
@@ -148,25 +149,28 @@ file(GLOB_RECURSE FLIGHTLIB_SOURCES
 
 # Create file lists for flightlib tests
 file(GLOB_RECURSE FLIGHTLIB_TEST_SOURCES
-  tests/bridges/*.cpp
   tests/dynamics/*.cpp
   tests/objects/*.cpp
   tests/sensors/*.cpp  
   tests/envs/*.cpp
   tests/common/*.cpp
 )
 
-# Create file lists for flightlib_gym tests
+# Create file lists for flightlib_gym source 
 file(GLOB_RECURSE FLIGHTLIB_GYM_SOURCES
   src/wrapper/*.cpp 
 )
 
-message(${FLIGHTLIB_GYM_SOURCES})
-
+# Create file lists for flightlib_gym tests
 file(GLOB_RECURSE FLIGHTLIB_GYM_TEST_SOURCES
   tests/wrapper/*.cpp
 )
 
+# Create file lists for flightlib_unity_bridge tests
+file(GLOB_RECURSE FLIGHTLIB_UNITY_BRIDGE_TEST_SOURCES
+  tests/bridges/*.cpp
+)
+
 ################################################################################
 # Optional Catkin Build
 ################################################################################
@@ -199,7 +203,8 @@ else()
     ${OpenCV_LIBRARIES}
     yaml-cpp
     zmq
-    zmqpp)
+    zmqpp
+    stdc++fs)
   set(LIBRARY_NAME ${PROJECT_NAME})
   set_target_properties(${PROJECT_NAME} PROPERTIES POSITION_INDEPENDENT_CODE TRUE)
 endif()
@@ -233,27 +238,34 @@ if(ENABLE_BLAS AND BLAS_FOUND AND LAPACK_FOUND)
   )
 endif()
 
-# Build tests
+# Build tests for flightlib gym wrapper
 if(BUILD_TESTS AND FLIGHTLIB_GYM_TEST_SOURCES)
   add_executable(test_gym ${FLIGHTLIB_GYM_TEST_SOURCES})
   target_link_libraries(test_gym
     ${LIBRARY_NAME}
     gtest
-    gtest_main
-)
+    gtest_main)
 add_test(test_gym test_gym)
 endif()
 
-# Build tests
+# Build tests for flightlib
 if(BUILD_TESTS AND FLIGHTLIB_TEST_SOURCES)
   add_executable(test_lib ${FLIGHTLIB_TEST_SOURCES})
   target_link_libraries(test_lib PUBLIC
     ${LIBRARY_NAME}
     gtest
-    gtest_main
-)
+    gtest_main)
 add_test(test_lib test_lib)
 endif()
 
+# Build tests for flightlib unity bridge
+if(BUILD_UNITY_BRIDGE_TESTS AND FLIGHTLIB_UNITY_BRIDGE_TEST_SOURCES)
+  add_executable(test_unity_bridge ${FLIGHTLIB_UNITY_BRIDGE_TEST_SOURCES})
+  target_link_libraries(test_unity_bridge PUBLIC
+    ${LIBRARY_NAME}
+    gtest
+    gtest_main)
+add_test(test_unity_bridge test_unity_bridge)
+endif()
 
 message(STATUS "================  !Done. No more nightmare!  ================")

flightlib/include/flightlib/bridges/unity_bridge.hpp

@@ -1,13 +1,19 @@
+//
+// This bridge was originally from FlightGoggles.
+// We made several changes on top of it.
+//
 #pragma once
 
 // std libs
 #include <unistd.h>
 #include <chrono>
+#include <experimental/filesystem>
 #include <fstream>
 #include <map>
 #include <string>
 #include <unordered_map>
 
+
 // Include ZMQ bindings for communications with Unity.
 #include <zmqpp/zmqpp.hpp>
 
@@ -18,6 +24,7 @@
 #include "flightlib/common/quad_state.hpp"
 #include "flightlib/common/types.hpp"
 #include "flightlib/objects/quadrotor.hpp"
+#include "flightlib/objects/static_object.hpp"
 #include "flightlib/objects/unity_camera.hpp"
 #include "flightlib/sensors/rgb_camera.hpp"
 
@@ -37,14 +44,17 @@ class UnityBridge {
 
   // public get functions
   bool getRender(const FrameID frame_id);
-  bool handleOutput(RenderMessage_t &output);
+  bool handleOutput();
+  bool getPointCloud(PointCloudMessage_t &pointcloud_msg,
+                     Scalar time_out = 600.0);
 
   // public set functions
   bool setScene(const SceneID &scene_id);
 
   // add object
   bool addQuadrotor(std::shared_ptr<Quadrotor> quad);
   bool addCamera(std::shared_ptr<UnityCamera> unity_camera);
+  bool addStaticObject(std::shared_ptr<StaticObject> static_object);
 
   // public auxiliary functions
   inline void setPubPort(const std::string &pub_port) { pub_port_ = pub_port; };
@@ -66,6 +76,7 @@ class UnityBridge {
 
   std::vector<std::shared_ptr<Quadrotor>> unity_quadrotors_;
   std::vector<std::shared_ptr<RGBCamera>> rgb_cameras_;
+  std::vector<std::shared_ptr<StaticObject>> static_objects_;
 
   // ZMQ variables and functions
   std::string client_address_;

flightlib/include/flightlib/bridges/unity_message_types.hpp

@@ -1,3 +1,7 @@
+//
+// This bridge message types was originally from FlightGoggles.
+// We made several changes on top of it.
+//
 #pragma once
 
 // std
@@ -112,6 +116,15 @@ struct SubMessage_t {
   std::vector<Sub_Vehicle_t> sub_vehicles;
 };
 
+struct PointCloudMessage_t {
+  // define point cloud box range [x, y, z] / meter
+  std::vector<Scalar> range{20.0, 20.0, 20.0};
+  std::vector<Scalar> origin{0.0, 0.0, 0.0};
+  Scalar resolution{0.15};
+  std::string path{"point_clouds_data/"};
+  std::string file_name{"default"};
+};
+
 /*********************
  * JSON constructors *
  *********************/
@@ -184,6 +197,14 @@ inline void from_json(const json &j, SubMessage_t &o) {
   o.sub_vehicles = j.at("pub_vehicles").get<std::vector<Sub_Vehicle_t>>();
 }
 
+inline void to_json(json &j, const PointCloudMessage_t &o) {
+  j = json{{"range", o.range},
+           {"origin", o.origin},
+           {"resolution", o.resolution},
+           {"path", o.path},
+           {"file_name", o.file_name}};
+}
+
 // Struct for outputting parsed received messages to handler functions
 struct RenderMessage_t {
   SubMessage_t sub_msg;

flightlib/include/flightlib/envs/env_base.hpp

@@ -1,3 +1,6 @@
+//
+// This is inspired by RaiGym, thanks.
+//
 #pragma once
 
 // standard library

flightlib/include/flightlib/envs/vec_env.hpp

@@ -1,3 +1,6 @@
+//
+// This is inspired by RaiGym, thanks.
+//
 #pragma once
 
 // std

flightlib/include/flightlib/objects/static_gate.hpp

@@ -0,0 +1,30 @@
+#pragma once
+
+#include "flightlib/objects/static_object.hpp"
+
+namespace flightlib {
+class StaticGate : public StaticObject {
+ public:
+  StaticGate(std::string id, std::string prefab_id);
+  ~StaticGate(){};
+
+  // publich set functions
+  void setPosition(const Vector<3>& position) { position_ = position; };
+  void setRotation(const Quaternion& quaternion) { quat_ = quaternion; };
+  void setSize(const Vector<3>& size) { size_ = size; };
+
+  // publich get functions
+  Vector<3> getPos(void) { return position_; };
+  Quaternion getQuat(void) { return quat_; };
+  Vector<3> getSize(void) { return size_; };
+
+ private:
+  std::string id_;
+  std::string prefab_id_;
+
+  Vector<3> position_{0.0, 0.0, 0.0};
+  Quaternion quat_{1.0, 0.0, 0.0, 0.0};
+  Vector<3> size_{1.0, 1.0, 1.0};
+};
+
+}  // namespace flightlib

flightlib/include/flightlib/objects/static_object.hpp

@@ -0,0 +1,27 @@
+#pragma once
+
+#include "flightlib/common/types.hpp"
+
+namespace flightlib {
+class StaticObject {
+ public:
+  StaticObject(std::string id, std::string prefab_id)
+    : id_(id), prefab_id_(prefab_id){};
+  virtual ~StaticObject(){};
+
+  // publich get functions
+  virtual Vector<3> getPos(void) = 0;
+  virtual Quaternion getQuat(void) = 0;
+  virtual Vector<3> getSize(void) = 0;
+
+  // public get functions
+  const std::string getID(void) { return id_; };
+  const std::string getPrefabID(void) { return prefab_id_; };
+
+
+ protected:
+  const std::string id_;
+  const std::string prefab_id_;
+};
+
+}  // namespace flightlib

flightlib/include/flightlib/sensors/rgb_camera.hpp

@@ -30,7 +30,7 @@ class RGBCamera : SensorBase {
   bool setHeight(const int height);
   bool setFOV(const Scalar fov);
   bool setDepthScale(const Scalar depth_scale);
-  bool setPostPrecesscing(const std::vector<bool>& enabled_layers);
+  bool setPostProcesscing(const std::vector<bool>& enabled_layers);
   bool feedImageQueue(const int image_layer, const cv::Mat& image_mat);
 
   // public get functions
@@ -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);

flightlib/src/bridges/unity_bridge.cpp

@@ -183,7 +183,20 @@ bool UnityBridge::addQuadrotor(std::shared_ptr<Quadrotor> quad) {
   return true;
 }
 
-bool UnityBridge::handleOutput(RenderMessage_t& output) {
+bool UnityBridge::addStaticObject(std::shared_ptr<StaticObject> static_object) {
+  Object_t object_t;
+  object_t.ID = static_object->getID();
+  object_t.prefab_ID = static_object->getPrefabID();
+  object_t.position = positionRos2Unity(static_object->getPos());
+  object_t.rotation = quaternionRos2Unity(static_object->getQuat());
+  object_t.size = scalarRos2Unity(static_object->getSize());
+
+  static_objects_.push_back(static_object);
+  settings_.objects.push_back(object_t);
+  pub_msg_.objects.push_back(object_t);
+}
+
+bool UnityBridge::handleOutput() {
   // create new message object
   zmqpp::message msg;
   sub_.receive(msg);
@@ -192,10 +205,71 @@ 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;
+}
+
+bool UnityBridge::getPointCloud(PointCloudMessage_t& pointcloud_msg,
+                                Scalar time_out) {
+  // create new message object
+  zmqpp::message msg;
+  // add topic header
+  msg << "PointCloud";
+  // create JSON object for initial settings
+  json json_msg = pointcloud_msg;
+  msg << json_msg.dump();
+  // send message without blocking
+  pub_.send(msg, true);
+
+  std::cout << "Generate PointCloud: Timeout=" << (int)time_out << " seconds."
+            << std::endl;
+
+  Scalar run_time = 0.0;
+  while (!std::experimental::filesystem::exists(
+    pointcloud_msg.path + pointcloud_msg.file_name + ".ply")) {
+    if (run_time >= time_out) {
+      logger_.warn("Timeout... PointCloud was not saved within expected time.");
+      return false;
+    }
+    std::cout << "Waiting for Pointcloud: Current Runtime=" << (int)run_time
+              << " seconds." << std::endl;
+    usleep((time_out / 10.0) * 1e6);
+    run_time += time_out / 10.0;
   }
   return true;
 }