Add toml config file to customize: labels, model-file, objects to tra…

…ck, etc

.gitignore

@@ -28,3 +28,4 @@ deps/densecrf/
 deps/freetype-gl-cpp/
 deps/gSLICr/
 deps/opencv/
+deps/toml11/

Core/CMakeLists.txt

@@ -16,11 +16,6 @@ find_package(OpenCV 3.1 REQUIRED)
 find_package(OpenMP)
 find_package(PythonLibs)
 
-set(MASKFUSION_MASK_RCNN_DIR ${CMAKE_SOURCE_DIR}/deps/Mask_RCNN CACHE STRING "Location of MaskRCNN")
-set(MASKFUSION_PYTHON_VE_PATH "" CACHE STRING "Python virtual environment to use")
-configure_file("Segmentation/MaskRCNN/MaskRCNN.py.in" "${CMAKE_CURRENT_BINARY_DIR}/../GUI/MaskRCNN.py")
-configure_file("Segmentation/MaskRCNN/helpers.py" "${CMAKE_CURRENT_BINARY_DIR}/../GUI/helpers.py")
-
 if (OPENMP_FOUND)
  set (CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}")
  set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}")

Core/MaskFusion.cpp

@@ -244,7 +244,6 @@ bool MaskFusion::processFrame(FrameDataPointer frame, const Eigen::Matrix4f* inP
  Model::generateCUDATextures(textureDepthMetricFiltered.get(), textureMask.get(), cudaIntrinsics, depthCutoff);
 
  TICK("odom");
-
  ModelListIterator itr = models.begin();
  (*itr)->performTracking(frameToFrameRGB,
  rgbOnly, icpWeight,
@@ -259,19 +258,20 @@ bool MaskFusion::processFrame(FrameDataPointer frame, const Eigen::Matrix4f* inP
 
  for(++itr;itr!=models.end();itr++){
  ModelPointer& mp = (*itr);
- if(!(*itr)->isNonstatic() && !trackAllModels){
- mp->updateStaticPose(globalModel->getPose()); // cam->cam_0=object_0 (cam_0->object_0 = identity)
+ bool trackable = trackableClassIds.empty() || trackableClassIds.count(mp->getClassID());
 
- } else {
+ if(((*itr)->isNonstatic() || trackAllModels) && trackable){
  Eigen::Matrix4f t = (*itr)->performTracking(frameToFrameRGB, rgbOnly, icpWeight, pyramid, fastOdom, so3, maxDepthProcessed, textureRGB.get(),
  frame->timestamp, requiresFillIn(*itr));
 
  // Don't allow big jumps (remove jumping models)
  float d = t.topRightCorner(3, 1).norm(); // Hack, do something reasonable
  if(d > 0.2){
  std::cout << "Removing model due to movement." << std::endl;
- //itr = inactivateModel(itr);
+ itr = inactivateModel(itr);
  }
+ } else {
+ mp->updateStaticPose(globalModel->getPose()); // cam->cam_0=object_0 (cam_0->object_0 = identity)
  }
  }
 
@@ -937,6 +937,7 @@ void MaskFusion::setMfMorphMaskIterations(int val) { labelGenerator.getMfSegment
 void MaskFusion::setMfThreshold(float val) { labelGenerator.getMfSegmentationPerformer()->threshold = val; }
 void MaskFusion::setMfWeightConvexity(float val) { labelGenerator.getMfSegmentationPerformer()->weightConvexity = val; }
 void MaskFusion::setMfWeightDistance(float val) { labelGenerator.getMfSegmentationPerformer()->weightDistance = val; }
+void MaskFusion::setTrackableClassIds(const std::set<int>& ids) { trackableClassIds = ids; }
 
 void MaskFusion::setFernThresh(const float& val) { fernThresh = val; }
 

Core/MaskFusion.h

@@ -243,6 +243,7 @@ class MaskFusion {
  void setMfWeightDistance(float val);
  void setMfWeightConvexity(float val);
  void setMfNonstaticThreshold(float val);
+ void setTrackableClassIds(const std::set<int>& ids);
 
  void setModelSpawnOffset(const unsigned& val);
  void setModelDeactivateCount(const unsigned& val);
@@ -415,7 +416,7 @@ class MaskFusion {
  unsigned modelSpawnOffset; // setting
  unsigned spawnOffset = 0; // current value
 
- const std::vector<int> nonstaticClasses = { 40 };
+ std::set<int> trackableClassIds;
 
  bool exportSegmentation;
  std::string exportDir;

Core/Model/Model.h

@@ -304,7 +304,6 @@ class Model {
  unsigned int id;
  int classID = -1;
 
- //bool isStatic = true;
  unsigned unseenCount = 0;
  unsigned age = 0;
 

Core/Segmentation/MaskRCNN/MaskRCNN.py.in

@@ -41,6 +41,7 @@ import numpy as np
 from PIL import Image
 import time
 import scipy.misc
+import pytoml as toml
 
 import tensorflow as tf
 config = tf.ConfigProto()
@@ -65,34 +66,22 @@ current_bounding_boxes = None
 # Root directory of the project
 DATA_DIR = os.path.join(MASK_RCNN_DIR, "data")
 MODEL_DIR = os.path.join(DATA_DIR, "logs")
-MODEL_PATH = os.path.join(DATA_DIR, "mask_rcnn_coco.h5")
 
 # PARAMETERS
-# Threshold for score:
-SCORE_T = 0.55
+
 # This only has an effect if the list contains at least 1 class name:
-FILTER_CLASSES = [] # Example: ['teddy bear']
 SPECIAL_ASSIGNMENTS = {} # Example: {'person': 255}
 SINGLE_INSTANCES = False
 OUTPUT_FRAMES = True
 STORE_CLASS_IDS = True
 START_INDEX = 0
 
-class_names = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
- 'bus', 'train', 'truck', 'boat', 'traffic light',
- 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird',
- 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',
- 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
- 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
- 'kite', 'baseball bat', 'baseball glove', 'skateboard',
- 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
- 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
- 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
- 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',
- 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
- 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster',
- 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors',
- 'teddy bear', 'hair drier', 'toothbrush']
+with open("config.toml", 'rb') as toml_file:
+ toml_config = toml.load(toml_file)
+ class_names = toml_config["MaskRCNN"]["class_names"]
+ model_path = toml_config["MaskRCNN"]["model_path"]
+ filter_classes = toml_config["MaskRCNN"]["filter_classes"]
+ score_threshold = toml_config["MaskRCNN"]["score_threshold"]
 
 class InferenceConfig(coco.CocoConfig):
  # Set batch size to 1 since we'll be running inference on
@@ -104,9 +93,9 @@ class InferenceConfig(coco.CocoConfig):
 config = InferenceConfig()
 config.display()
 model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)
-model.load_weights(MODEL_PATH, by_name=True)
+model.load_weights(model_path, by_name=True)
 
-FILTER_CLASSES = [class_names.index(x) for x in FILTER_CLASSES]
+filter_classes = [class_names.index(x) for x in filter_classes]
 SPECIAL_ASSIGNMENTS = {class_names.index(x): SPECIAL_ASSIGNMENTS[x] for x in SPECIAL_ASSIGNMENTS}
 
 def execute(rgb_image):
@@ -119,4 +108,4 @@ def execute(rgb_image):
  #if SINGLE_INSTANCES:
  # merge_instances(r)
 
- current_segmentation, current_class_ids, current_bounding_boxes = generate_id_image(r, SCORE_T, FILTER_CLASSES, SPECIAL_ASSIGNMENTS)
+ current_segmentation, current_class_ids, current_bounding_boxes = generate_id_image(r, score_threshold, filter_classes, SPECIAL_ASSIGNMENTS)

Core/Segmentation/MaskRCNN/offline_runner.py

@@ -37,10 +37,12 @@
 from PIL import Image
 from helpers import *
 import time
+import pytoml as toml
 import scipy.misc
 
 parser = argparse.ArgumentParser()
 parser.add_argument("-i", required=True, help="Input directory (all files are being processed)")
+parser.add_argument("-c", required=False, help="Optional config file, otherwise MsCoco is assumed")
 parser.add_argument("-o", required=True, help="Output directory")
 parser.add_argument("--filter", nargs='+', required=False,
  help="Specify which labels you would like to use (empty means all), example: --filter teddy_bear pizza baseball_bat")
@@ -49,7 +51,7 @@
 # FURTHER PARAMETERS
 EXTENSIONS = ['jpg', 'png']
 FILTER_IMAGE_NAME = "" # only use images, whose name contains this string (eg "Color")
-SCORE_T = 0.85
+score_threshold = 0.85
 SPECIAL_ASSIGNMENTS = {} #{'person': 255}
 SINGLE_INSTANCES = False
 OUTPUT_FRAMES = True
@@ -60,12 +62,13 @@
 OUTPUT_DIR = args.o
 DATA_DIR = os.path.join(mask_rcnn_path, "data")
 MODEL_DIR = os.path.join(DATA_DIR, "logs")
-COCO_MODEL_PATH = os.path.join(DATA_DIR, "mask_rcnn_coco.h5")
-FILTER_CLASSES = []
+model_path = os.path.join(DATA_DIR, "mask_rcnn_coco.h5")
+
+filter_classes = []
 if args.filter:
- FILTER_CLASSES = args.filter
- FILTER_CLASSES = [f.replace("_", " ") for f in FILTER_CLASSES]
-CLASS_NAMES = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
+ filter_classes = args.filter
+ filter_classes = [f.replace("_", " ") for f in filter_classes]
+class_names = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
  'bus', 'train', 'truck', 'boat', 'traffic light',
  'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird',
  'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',
@@ -80,14 +83,24 @@
  'keyboard', 'cell phone', 'microwave', 'oven', 'toaster',
  'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors',
  'teddy bear', 'hair drier', 'toothbrush']
-FILTER_CLASSES = [CLASS_NAMES.index(x) for x in FILTER_CLASSES]
-SPECIAL_ASSIGNMENTS = {CLASS_NAMES.index(x): SPECIAL_ASSIGNMENTS[x] for x in SPECIAL_ASSIGNMENTS}
+
+if args.c:
+ with open(args.c, 'rb') as toml_file:
+ toml_config = toml.load(toml_file)
+ class_names = toml_config["MaskRCNN"]["class_names"]
+ model_path = toml_config["MaskRCNN"]["model_path"]
+ filter_classes = toml_config["MaskRCNN"]["filter_classes"]
+ score_threshold = toml_config["MaskRCNN"]["score_threshold"]
+
+filter_classes = [class_names.index(x) for x in filter_classes]
+SPECIAL_ASSIGNMENTS = {class_names.index(x): SPECIAL_ASSIGNMENTS[x] for x in SPECIAL_ASSIGNMENTS}
 
 class InferenceConfig(coco.CocoConfig):
  # Set batch size to 1 since we'll be running inference on
  # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
  GPU_COUNT = 1
  IMAGES_PER_GPU = 1
+ NUM_CLASSES = len(class_names)
 
 config = InferenceConfig()
 config.display()
@@ -96,7 +109,7 @@ class InferenceConfig(coco.CocoConfig):
 model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)
 
 # Load weights trained on MS-COCO
-model.load_weights(COCO_MODEL_PATH, by_name=True)
+model.load_weights(model_path, by_name=True)
 
 file_names = [fn for fn in os.listdir(IMAGE_DIR) if any(fn.endswith(ext) for ext in EXTENSIONS)]
 file_names.sort()
@@ -150,17 +163,17 @@ class InferenceConfig(coco.CocoConfig):
  merge_instances(r)
 
  #out_path = os.path.join(OUTPUT_DIR, "{}.png".format(idx))
- id_image, exported_class_ids, exported_rois = generate_id_image(r, SCORE_T, FILTER_CLASSES, SPECIAL_ASSIGNMENTS)
+ id_image, exported_class_ids, exported_rois = generate_id_image(r, score_threshold, filter_classes, SPECIAL_ASSIGNMENTS)
  save_id_image(id_image, OUTPUT_DIR, base_name, exported_class_ids, STORE_CLASS_IDS, exported_rois)
 
 
  # Visualise
  ax.clear()
- filter_result(r, FILTER_CLASSES)
+ filter_result(r, filter_classes)
  #visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'],
- # CLASS_NAMES, r['scores'], SCORE_T, ax=ax) # requires patched version
+ # class_names, r['scores'], score_threshold, ax=ax) # requires patched version
  visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'],
- CLASS_NAMES, r['scores'], ax=ax)
+ class_names, r['scores'], ax=ax)
  fig.canvas.draw()
  if OUTPUT_FRAMES:
  plt.savefig(os.path.join(OUTPUT_DIR, base_name+".jpg"))

GUI/CMakeLists.txt

@@ -21,17 +21,19 @@ endif()
 
 option(WITH_FREETYPE_GL_CPP "Enable drawing of labels" OFF)
 if(WITH_FREETYPE_GL_CPP)
- find_package(freetype-gl-cpp REQUIRED ${CMAKE_SOURCE_DIR}/deps/freetype-gl-cpp/build)
+# find_package(freetype-gl-cpp REQUIRED ${CMAKE_SOURCE_DIR}/deps/freetype-gl-cpp/build)
  link_directories(${CMAKE_SOURCE_DIR}/deps/freetype-gl-cpp/install/lib/freetype-gl-cpp)
  include_directories(${CMAKE_SOURCE_DIR}/deps/freetype-gl-cpp/install/include)
  add_definitions(-DWITH_FREETYPE_GL_CPP)
  list(APPEND LIBRARIES "freetype-gl-cpp")
+ message("WITH_FREETYPE_GL_CPP")
 endif()
 
+include_directories(${CMAKE_SOURCE_DIR}/Core)
+include_directories(${CMAKE_SOURCE_DIR}/deps/toml11)
 include_directories(${ZLIB_INCLUDE_DIR})
 include_directories(${EIGEN_INCLUDE_DIRS})
 include_directories(${CUDA_INCLUDE_DIRS})
-include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../Core)
 include_directories(${OPENNI2_INCLUDE_DIR})
 include_directories(${Pangolin_INCLUDE_DIRS})
 include_directories(${BOOST_INCLUDE_DIRS})
@@ -44,6 +46,13 @@ message("CUDA_INCLUDE_DIRS: ${CUDA_INCLUDE_DIRS}")
 message("OPENNI2_INCLUDE_DIR: ${OPENNI2_INCLUDE_DIR}")
 message("CMAKE_CURRENT_SOURCE_DIR/../Core: ${${CMAKE_CURRENT_SOURCE_DIR}/../Core}")
 
+set(MASKFUSION_MASK_RCNN_DIR ${CMAKE_SOURCE_DIR}/deps/Mask_RCNN CACHE STRING "Location of MaskRCNN")
+set(MASKFUSION_PYTHON_VE_PATH "" CACHE STRING "Python virtual environment to use")
+set(MASKFUSION_MODEL_FILE "${MASKFUSION_MASK_RCNN_DIR}/data/mask_rcnn_coco.h5" CACHE STRING "H5 MaskRCNN model weights")
+configure_file("../Core/Segmentation/MaskRCNN/MaskRCNN.py.in" "${CMAKE_CURRENT_BINARY_DIR}/MaskRCNN.py")
+configure_file("../Core/Segmentation/MaskRCNN/helpers.py" "${CMAKE_CURRENT_BINARY_DIR}/helpers.py")
+configure_file("../config.toml.in" "${CMAKE_CURRENT_BINARY_DIR}/config.toml")
+
 file(GLOB srcs *.cpp)
 file(GLOB tools_srcs Tools/*.cpp)
 

GUI/MainController.cpp

@@ -29,6 +29,7 @@
 #include <boost/algorithm/string.hpp>
 #include <GUI/Tools/PangolinReader.h>
 #include <opencv2/core/ocl.hpp>
+#include <toml.hpp>
 
 /*
  * Parameters:
@@ -264,8 +265,26 @@ MainController::MainController(int argc, char* argv[])
  gui->addBifoldParameters();
  }
 
+ // Load configuration from files
  if(pangolin::FileExists("parameters.cfg"))
  pangolin::ParseVarsFile("parameters.cfg");
+ if(!pangolin::FileExists("config.toml"))
+ throw std::runtime_error("Could not read 'config.toml', which specifies class-names and weights.");
+
+ std::vector<std::string> classNames;
+ std::vector<std::string> trackableClasses;
+ try {
+ toml::table tomlConfig = toml::parse("config.toml");
+ const auto tomlMaskRCNN = toml::get<toml::Table>(tomlConfig.at("MaskRCNN"));
+ classNames = toml::get<std::vector<std::string>>(tomlMaskRCNN.at("class_names"));
+ trackableClasses = toml::get<std::vector<std::string>>(tomlMaskRCNN.at("trackable_classes"));
+
+ for(std::string& c : trackableClasses){
+ trackableClassIds.insert(std::distance(classNames.begin(), std::find(classNames.begin(), classNames.end(), c)));
+ }
+ } catch(...){
+ throw std::invalid_argument("Unable to parse input toml configuration file.");
+ }
 
  if (Parse::get().arg(argc, argv, "-d", tmpFloat) > -1) gui->depthCutoff->Ref().Set(tmpFloat);
  if (Parse::get().arg(argc, argv, "-i", tmpFloat) > -1) gui->icpWeight->Ref().Set(tmpFloat);
@@ -318,23 +337,9 @@ MainController::MainController(int argc, char* argv[])
  "Exporting results to: "
  << exportDir << std::endl;
 
- //  // Setup rendering for labels
+ // Setup rendering for labels
 #ifdef WITH_FREETYPE_GL_CPP
- gui->addLabelTexts({"BG", "person", "bicycle", "car", "motorcycle", "airplane",
- "bus", "train", "truck", "boat", "traffic light",
- "fire hydrant", "stop sign", "parking meter", "bench", "bird",
- "cat", "dog", "horse", "sheep", "cow", "elephant", "bear",
- "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie",
- "suitcase", "frisbee", "skis", "snowboard", "sports ball",
- "kite", "baseball bat", "baseball glove", "skateboard",
- "surfboard", "tennis racket", "bottle", "wine glass", "cup",
- "fork", "knife", "spoon", "bowl", "banana", "apple",
- "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza",
- "donut", "cake", "chair", "couch", "potted plant", "bed",
- "dining table", "toilet", "tv", "laptop", "mouse", "remote",
- "keyboard", "cell phone", "microwave", "oven", "toaster",
- "sink", "refrigerator", "book", "clock", "vase", "scissors",
- "teddy bear", "hair drier", "toothbrush"}, {{0,0,0,1}});
+ gui->addLabelTexts(classNames, {{0,0,0,1}});
 #endif
 }
 
@@ -395,6 +400,7 @@ void MainController::launch() {
  !openLoop, iclnuim, reloc, photoThresh, confGlobalInit, confObjectInit, gui->depthCutoff->Get(),
  gui->icpWeight->Get(), fastOdom, fernThresh, so3, frameToFrameRGB, gui->modelSpawnOffset->Get(),
  Model::MatchingType::Drost, segmentationMethod, exportDir, exportSegmentation, usePrecomputedMasksOnly, frameQueueSize);
+ maskFusion->setTrackableClassIds(trackableClassIds);
 
  gui->addTextureColumn(maskFusion->getDrawableTextures());
 

GUI/MainController.h

@@ -72,4 +72,6 @@ class MainController {
  Segmentation::Method segmentationMethod;
 
  GPUResize* resizeStream;
+
+ std::set<int> trackableClassIds;
 };

build.sh

@@ -138,6 +138,7 @@ pip3 install h5py
 pip3 install cython
 pip3 install imgaug
 pip3 install opencv-python
+pip3 install pytoml
 ln -s `python -c "import numpy as np; print(np.__path__[0])"`/core/include/numpy Core/Segmentation/MaskRCNN || true # Provide numpy headers to C++
 
 
@@ -286,6 +287,10 @@ if [[ $* == *--build-dependencies* ]] ; then
  wget --no-clobber https://github.com/matterport/Mask_RCNN/releases/download/v1.0/mask_rcnn_coco.h5
  cd ../..
 
+ # c++ toml
+ highlight "Building toml11..."
+ git_clone "git clone --depth=1 https://github.com/ToruNiina/toml11.git"
+
  cd ..
 fi # --build-dependencies