C++ implementation of SiamMask

.gitignore

@@ -1 +1,5 @@
-./build/*
+/build/*
+/models/*
+CMakeLists.txt.user
+
+

CMakeLists.txt

@@ -0,0 +1,18 @@
+cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
+project(SiamMaskCpp)
+
+find_package(Torch REQUIRED)
+find_package(OpenCV 3.4 REQUIRED)
+
+message(STATUS "CMAKE_PREFIX_PATH: " ${CMAKE_PREFIX_PATH})
+message(STATUS "TORCH_LIBRARIES: ${TORCH_LIBRARIES}")
+message(STATUS "OpenCV_LIBS: ${OpenCV_LIBS}")
+
+include_directories(
+    "."
+    "${CMAKE_PREFIX_PATH}/include"
+)
+
+add_executable(demo demo.cpp)
+target_link_libraries(demo "${TORCH_LIBRARIES}" "${OpenCV_LIBS}" dlib)
+set_property(TARGET demo PROPERTY CXX_STANDARD 11)

common.h

@@ -0,0 +1,38 @@
+#ifndef SIAMMASK_COMMON_Hh
+#define SIAMMASK_COMMON_Hh
+
+#include <torch/script.h>
+#include <opencv2/opencv.hpp>
+#include <vector>
+#include <string>
+#include <iostream>
+
+using Tensor = torch::Tensor;
+using IValue = torch::IValue;
+using Device = torch::Device;
+
+using GpuMat = cv::cuda::GpuMat;
+using Mat = cv::Mat;
+using Point = cv::Point;
+using Rect = cv::Rect;
+using Scalar = cv::Scalar;
+
+template<typename T, typename U>
+using pair = std::pair<T, U>;
+template<typename T>
+using vector = std::vector<T>;
+template<typename T, typename U>
+using map = std::map<T, U>;
+using string = std::string;
+using std::cout;
+using std::endl;
+
+#ifndef shapeof
+#define shapeof(m) vector<int>({m.rows, m.cols, m.channels()})
+#endif // shapeof
+
+#ifndef print
+#define print(x) cout << (#x) << ": " << x << endl
+#endif // print
+
+#endif // SIAMMASK_COMMON_Hh

config_davis.cfg

@@ -0,0 +1,15 @@
+hp {
+    instance_size = 255
+    base_size = 8
+    out_size = 127
+    seg_thr = 0.35
+    penalty_k = 0.04
+    window_influence = 0.4
+    lr = 1.0
+}
+anchors {
+    stride = 8
+    ratios = [0.33, 0.5, 1, 2, 3]
+    scales = [8]
+}
+

config_reader.h

@@ -0,0 +1,15 @@
+#include <dlib/config_reader.h>
+#include <argparse/argparse.hpp>
+
+namespace dlib {
+    // Add vector I/O compatible with argparse's
+    template <typename cr_type, typename T>
+    typename enable_if<is_config_reader<cr_type>, std::vector<T> >::type get_option (
+        const cr_type& cr,
+        const std::string& option_name,
+        const std::vector<T>& default_value
+    ) {
+        std::string value_str = get_option(cr, option_name, argparse::toString(default_value));
+        return argparse::castTo<std::vector<T> >(value_str);
+    }
+}

convert.h

@@ -0,0 +1,34 @@
+#ifndef SIAMMASK_CONVERT_Hh
+#define SIAMMASK_CONVERT_Hh
+
+#include "common.h"
+
+inline Tensor toTensor(const Mat& img) {
+    Mat imgf;
+    img.convertTo(imgf, CV_32FC3);
+    return torch::tensor(
+           torch::ArrayRef<float>(
+                imgf.ptr<float>(),
+                imgf.rows * imgf.cols * imgf.channels()
+           )
+    ).reshape({1, img.rows, img.cols, 3}).permute({0, 3, 1, 2});
+}
+
+inline Tensor toTensor(const GpuMat& gpuimg) {
+    Mat img;
+    gpuimg.download(img);
+    return toTensor(img);
+}
+
+inline void toMat(const Tensor& tensor, Mat& mat) {
+    CV_Assert(tensor.dim() == 2, tensor.scalar_type() == torch::ScalarType::Float);
+    Mat(tensor.size(0), tensor.size(1), CV_32FC1, tensor.cpu().data_ptr<float>()).copyTo(mat);
+}
+
+inline void toGpuMat(const Tensor& tensor, GpuMat& gmat) {
+    Mat mat;
+    toMat(tensor, mat);
+    gmat.upload(mat);
+}
+
+#endif // SIAMMASK_CONVERT_Hh

demo.cpp

@@ -0,0 +1,73 @@
+#include <argparse/argparse.hpp>
+#include <dlib/dir_nav.h>
+#include "siammask.h"
+
+void overlayMask(const Mat& src, const Mat& mask, Mat& dst) {
+    vector<Mat> chans;
+    cv::split(src, chans);
+    cv::max(chans[2], mask, chans[2]);
+    cv::merge(chans, dst);
+}
+
+int main(int argc, const char* argv[]) try {
+    argparse::ArgumentParser parser;
+    parser.addArgument("-m", "--modeldir", 1, false);
+    parser.addArgument("-c", "--config", 1, false);
+    parser.addFinalArgument("target");
+
+    parser.parse(argc, argv);
+
+    Device device(torch::kCUDA);
+
+    SiamMask siammask(parser.retrieve<string>("modeldir"), device);
+    State state;
+    state.load_config(parser.retrieve<string>("config"));
+
+    dlib::directory target_dir(parser.retrieve<string>("target"));
+    vector<dlib::file> image_files = dlib::get_files_in_directory_tree(
+        target_dir, dlib::match_endings("jpg png"), 0
+    );
+    std::sort(image_files.begin(), image_files.end());
+
+    cout << image_files.size() << " images found in " << target_dir << endl;
+
+    vector<Mat> images;
+    for(const auto& image_file : image_files) {
+        images.push_back(cv::imread(image_file.full_name()));
+    }
+
+    cv::namedWindow("SiamMask");
+    int64 toc = 0;
+    Rect roi = cv::selectROI("SiamMask", images.front(), false);
+
+    for(unsigned long i = 0; i < images.size(); ++i) {
+        int64 tic = cv::getTickCount();
+
+        Mat& src = images[i];
+        GpuMat gsrc;
+        gsrc.upload(src);
+
+        if (i == 0) {
+            cout << "Initializing..." << endl;
+            siameseInit(state, siammask, gsrc, roi, device);
+            print(state);
+            cv::rectangle(src, roi, Scalar(0, 255, 0));
+            cv::imshow("SiamMask", src);
+        } else {
+            siameseTrack(state, siammask, gsrc, device);
+            overlayMask(src, state.mask, src);
+            cv::rectangle(src, state.target, Scalar(0, 255, 0));
+            cv::imshow("SiamMask", src);
+        }
+
+        toc += cv::getTickCount() - tic;
+        cv::waitKey(1);
+    }
+
+    double total_time = toc / cv::getTickFrequency();
+    double fps = image_files.size() / total_time;
+    printf("SiamMask Time: %.1fs Speed: %.1ffps (with visulization!)\n", total_time, fps);
+} catch (std::exception& e) {
+    cout << "Exception thrown!\n" << e.what() << endl;
+}
+

geometry.h

@@ -0,0 +1,39 @@
+#ifndef SIAMMASK_GEOMETRY_Hh
+#define SIAMMASK_GEOMETRY_Hh
+
+#include "common.h"
+
+inline Point center(const Rect& rect) {
+    return (rect.tl() + rect.br()) / 2;
+}
+
+inline Rect centeredRect(const Point& center, int width, int height) {
+    return Rect(center.x - width / 2, center.y - height / 2, width, height);
+}
+
+template<typename MatType>
+Rect getRect(const MatType& img) {
+    return Rect(0, 0, img.cols, img.rows);
+}
+
+inline Rect uniteRects(const Rect& r1, const Rect& r2) {
+    const int left = std::min(r1.x, r2.x);
+    const int top = std::min(r1.y, r2.y);
+    const int right = std::max(r1.x + r1.width, r2.x + r2.width);
+    const int bottom = std::max(r1.y + r1.height, r2.y + r2.height);
+    return Rect(left, top, right - left, bottom - top);
+}
+
+inline Rect intersectRects(const Rect& r1, const Rect& r2) {
+    const int left = std::max(r1.x, r2.x);
+    const int top = std::max(r1.y, r2.y);
+    const int right = std::min(r1.x + r1.width, r2.x + r2.width);
+    const int bottom = std::min(r1.y + r1.height, r2.y + r2.height);
+    return Rect(left, top, right - left, bottom - top);
+}
+
+inline Rect translateRect(const Rect& rect, const Point& p) {
+    return Rect(rect.x + p.x, rect.y + p.y, rect.width, rect.height);
+}
+
+#endif // SIAMMASK_GEOMETRY_Hh

numpy.h

@@ -0,0 +1,108 @@
+#ifndef SIAMMASK_NUMPY_Hh
+#define SIAMMASK_NUMPY_Hh
+
+#include "geometry.h"
+
+namespace numpy {
+    inline Mat tile(const Mat& x, uint64_t reps, bool along_row = true) {
+        Mat concat;
+        if(along_row)
+            cv::hconcat(vector<Mat>(reps, x), concat);
+        else
+            cv::vconcat(vector<Mat>(reps, x), concat);
+        return concat;
+    }
+
+    inline GpuMat concat(const GpuMat& x, const GpuMat& y, bool along_row = true) {
+        if(x.empty())
+            return y.clone();
+        else if (y.empty())
+            return x.clone();
+        CV_Assert(x.channels() == y.channels(), x.type() == y.type());
+        if(along_row) {
+            CV_Assert(x.rows == y.rows);
+            GpuMat res = cv::cuda::createContinuous(x.rows, x.cols + y.cols, x.type());
+            x.copyTo(res(getRect(x)));
+            y.copyTo(res(translateRect(getRect(y), Point(x.cols, 0))));
+            return res;
+        } else {
+            CV_Assert(x.cols == y.cols);
+            GpuMat res = cv::cuda::createContinuous(x.rows + y.rows, x.cols, x.type());
+            x.copyTo(res(getRect(x)));
+            y.copyTo(res(translateRect(getRect(y), Point(0, x.rows))));
+            return res;
+        }
+    }
+
+    inline GpuMat tile(const GpuMat& x, uint64_t reps, bool along_row = true) {
+        uint8_t digits = log2(reps) + 1;
+        vector<GpuMat> concats;
+        concats.push_back(x);
+
+        for(uint8_t digit = 0; digit + 1 < digits; ++digit)
+            concats.push_back(concat(concats.back(), concats.back(), along_row));
+
+        GpuMat tiled;
+        for(uint8_t digit = 0; digit < digits; ++digit) {
+            if((reps >> digit) & 1)
+                tiled = concat(tiled, concats[digit], along_row);
+        }
+
+        return tiled;
+    }
+
+    inline pair<Mat, Mat> meshgrid(const Mat& x, const Mat& y) {
+        CV_Assert(x.rows == 1 && y.rows == 1);
+        return std::make_pair(
+            tile(x, y.cols, false),
+            tile(y.t(), x.cols)
+        );
+    }
+
+    inline pair<GpuMat, GpuMat> meshgrid(const GpuMat& x, const GpuMat& y) {
+        CV_Assert(x.rows == 1 && y.rows == 1);
+        GpuMat yt;
+        cv::cuda::transpose(y, yt);
+        return std::make_pair(
+            tile(x, y.cols, false),
+            tile(yt, x.cols)
+        );
+    }
+
+    inline Mat hanning(uint64_t M) {
+        Mat w(1, M, CV_32FC1);
+        for(uint64_t n = 0; n < M; ++n)
+            w.at<float>(0, n) = 0.5 - 0.5*cos(2*dlib::pi*n / (M - 1));
+        return w;
+    }
+
+    inline Mat outer(const Mat& x, const Mat& y) {
+        CV_Assert(x.rows == 1 && y.rows == 1);
+        return x.t() * y;
+    }
+
+    inline vector<long> unravel_index(long index, const vector<long>& dims) {
+        vector<long> indices;
+        for(unsigned long i = dims.size() - 1; i >= 1; --i) {
+            long ax_idx = index % dims[i];
+            indices.push_back(ax_idx);
+            index = (index - ax_idx) / dims[i];
+        }
+        indices.push_back(index);
+        std::reverse(indices.begin(), indices.end());
+        return indices;
+    }
+
+    inline Mat arange(long begins, long end, long step = 1) {
+        Mat range(1, int((end - begins) / step), CV_32FC1);
+        for(int r = 0; r < range.cols; ++r)
+            range.at<float>(0, r) = begins + r * step;
+        return range;
+    }
+
+    inline Mat arange(unsigned long n) {
+        return arange(0, n);
+    }
+}
+
+#endif // SIAMMASK_NUMPY_Hh