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