<commit_before>/* Copyright (c) 2005 - 2011 Vertica, an HP company -*- C++ -*- */
/*
 *
 * Description: Functions that help with full-text search over documents
 *
 * Create Date: Nov 25, 2011
 */
#include <sstream>
#include <set>
#include "Vertica.h"
#include "libstemmer.h"
#include <curl/curl.h>
#include <json/json.h>
#include "Words.h"

using namespace Vertica;
using namespace std;

/*
 * This is a simple function that stems an input word
 */
class Stem : public ScalarFunction
{
    sb_stemmer *stemmer;

public:
    // Set up the stemmer
    virtual void setup(ServerInterface &srvInterface, const SizedColumnTypes &argTypes)
    {
        stemmer = sb_stemmer_new("english", NULL);
    }

    // Destroy the stemmer
    virtual void destroy(ServerInterface &srvInterface, const SizedColumnTypes &argTypes)
    {
        sb_stemmer_delete(stemmer);
    }

    /*
     * This method processes a block of rows in a single invocation.
     *
     * The inputs are retrieved via arg_reader
     * The outputs are returned via arg_writer
     */
    virtual void processBlock(ServerInterface &srvInterface,
                              BlockReader &arg_reader,
                              BlockWriter &res_writer)
    {
        // Basic error checking
        if (arg_reader.getNumCols() != 1)
            vt_report_error(0, "Function only accept 1 arguments, but %zu provided",
                            arg_reader.getNumCols());

        // While we have inputs to process
        do {
            // Get a copy of the input string
            std::string  inStr = arg_reader.getStringRef(0).str();

            const unsigned char* stemword = sb_stemmer_stem(stemmer, reinterpret_cast<const sb_symbol*>(inStr.c_str()), inStr.size());

            // Copy string into results
            res_writer.getStringRef().copy(reinterpret_cast<const char*>(stemword));
            res_writer.next();
        } while (arg_reader.next());
    }
};

class StemFactory : public ScalarFunctionFactory
{
    // return an instance of RemoveSpace to perform the actual addition.
    virtual ScalarFunction *createScalarFunction(ServerInterface &interface)
    { return vt_createFuncObj(interface.allocator, Stem); }

    virtual void getPrototype(ServerInterface &interface,
                              ColumnTypes &argTypes,
                              ColumnTypes &returnType)
    {
        argTypes.addVarchar();
        returnType.addVarchar();
    }

    virtual void getReturnType(ServerInterface &srvInterface,
                               const SizedColumnTypes &argTypes,
                               SizedColumnTypes &returnType)
    {
        const VerticaType &t = argTypes.getColumnType(0);
        returnType.addVarchar(t.getStringLength());
    }
};

RegisterFactory(StemFactory);

static size_t
WriteMemoryCallback(void *contents, size_t size, size_t nmemb, void *userp)
{
    size_t realsize = size * nmemb;
    stringstream *mem = (stringstream *)userp;

    mem->write(static_cast<const char*>(contents), realsize);

    return realsize;
}

static string fetch_url(CURL *curl_handle, const char *url)
{
    stringstream chunk;

    /* specify URL to get */
    curl_easy_setopt(curl_handle, CURLOPT_URL, url);

    /* send all data to this function  */
    curl_easy_setopt(curl_handle, CURLOPT_WRITEFUNCTION, WriteMemoryCallback);

    /* we pass our 'chunk' struct to the callback function */
    curl_easy_setopt(curl_handle, CURLOPT_WRITEDATA, (void *)&chunk);

    /* some servers don't like requests that are made without a user-agent
       field, so we provide one */
    curl_easy_setopt(curl_handle, CURLOPT_USERAGENT, "libcurl-agent/1.0");

    /* get it! */
    curl_easy_perform(curl_handle);

    return chunk.str();
}

/*
 * Fetch the given URL as a string
 */
class FetchURL : public ScalarFunction
{
    CURL *curl_handle;

public:
    virtual void setup(ServerInterface &srvInterface, const SizedColumnTypes &argTypes)
    {
        curl_global_init(CURL_GLOBAL_ALL);

        /* init the curl session */
        curl_handle = curl_easy_init();
    }

    // Destroy the stemmer
    virtual void destroy(ServerInterface &srvInterface, const SizedColumnTypes &argTypes)
    {
        /* cleanup curl stuff */
        curl_easy_cleanup(curl_handle);
        /* we're done with libcurl, so clean it up */
        curl_global_cleanup();
    }

    /*
     * This method processes a block of rows in a single invocation.
     *
     * The inputs are retrieved via arg_reader
     * The outputs are returned via arg_writer
     */
    virtual void processBlock(ServerInterface &srvInterface,
                              BlockReader &arg_reader,
                              BlockWriter &res_writer)
    {
        // Basic error checking
        if (arg_reader.getNumCols() != 1)
            vt_report_error(0, "Function only accept 1 arguments, but %zu provided",
                            arg_reader.getNumCols());

        try {
            // While we have inputs to process
            do {
                // Get a copy of the input string
                std::string  inStr = arg_reader.getStringRef(0).str();

                /* fetch contents of URL */
                string contents = fetch_url(curl_handle, inStr.c_str()).substr(0, 65000);

                // Copy string into results
                res_writer.getStringRef().copy(contents);
                res_writer.next();
            } while (arg_reader.next());
        } catch (std::exception &e) {
            vt_report_error(0, e.what());
        }
    }
};

class FetchURLFactory : public ScalarFunctionFactory
{
    // return an instance of RemoveSpace to perform the actual addition.
    virtual ScalarFunction *createScalarFunction(ServerInterface &interface)
    { return vt_createFuncObj(interface.allocator, FetchURL); }

    virtual void getPrototype(ServerInterface &interface,
                              ColumnTypes &argTypes,
                              ColumnTypes &returnType)
    {
        argTypes.addVarchar();
        returnType.addVarchar();
    }

    virtual void getReturnType(ServerInterface &srvInterface,
                               const SizedColumnTypes &argTypes,
                               SizedColumnTypes &returnType)
    {
        returnType.addVarchar(65000); // max varchar size
    }
};

RegisterFactory(FetchURLFactory);

/*
 * Parse the Twitter JSON results
 */
class TwitterJSONParser : public TransformFunction
{
    Json::Value root;   // will contains the root value after parsing.
    Json::Reader reader;

public:

    virtual void processPartition(ServerInterface &srvInterface,
                                  PartitionReader &input,
                                  PartitionWriter &output)
    {
        // Basic error checking
        if (input.getNumCols() != 1)
            vt_report_error(0, "Function only accept 1 arguments, but %zu provided",
                            input.getNumCols());

        // While we have inputs to process
        do {
            // Get a copy of the input string
            string  inStr = input.getStringRef(0).str();

            bool success = reader.parse(inStr, root, false);
            if (!success)
                vt_report_error(0, "Malformed JSON text");

            // Get the query string and results
            string query = root.get("query", "-").asString();
            Json::Value results = root.get("results",  Json::Value::null);

            for (uint i=0; i<results.size(); i++) {
                Json::Value result = results[i];
                string text = result.get("text", "").asString();

                // Copy string into results
                output.getStringRef(0).copy(query);
                output.getStringRef(1).copy(text);
                output.next();
            }
        } while (input.next());
    }
};

class TwitterJSONFactory : public TransformFunctionFactory
{
    // return an instance of the transform function
    virtual TransformFunction *createTransformFunction(ServerInterface &interface)
    { return vt_createFuncObj(interface.allocator, TwitterJSONParser); }

    virtual void getPrototype(ServerInterface &interface,
                              ColumnTypes &argTypes,
                              ColumnTypes &returnType)
    {
        argTypes.addVarchar();
        returnType.addVarchar();
        returnType.addVarchar();
    }

    virtual void getReturnType(ServerInterface &srvInterface,
                               const SizedColumnTypes &argTypes,
                               SizedColumnTypes &returnType)
    {
        returnType.addVarchar(100, "query"); // query string size <= 100
        returnType.addVarchar(300, "text"); // tweet size <= 300
    }
};

RegisterFactory(TwitterJSONFactory);

/*
 * Search Twitter and return results
 */
class TwitterSearch : public TransformFunction
{
    Json::Value root;   // will contains the root value after parsing.
    Json::Reader reader;
    CURL *curl_handle;

public:
    virtual void setup(ServerInterface &srvInterface, const SizedColumnTypes &argTypes)
    {
        curl_global_init(CURL_GLOBAL_ALL);

        /* init the curl session */
        curl_handle = curl_easy_init();
    }

    virtual void destroy(ServerInterface &srvInterface, const SizedColumnTypes &argTypes)
    {
        /* cleanup curl stuff */
        curl_easy_cleanup(curl_handle);
        /* we're done with libcurl, so clean it up */
        curl_global_cleanup();
    }

    virtual void processPartition(ServerInterface &srvInterface,
                                  PartitionReader &input,
                                  PartitionWriter &output)
    {
        // Basic error checking
        if (input.getNumCols() != 1)
            vt_report_error(0, "Function only accept 1 arguments, but %zu provided",
                            input.getNumCols());

        try {
            // While we have inputs to process
            do {
                // Get a copy of the input string
                string  inStr = input.getStringRef(0).str();

                // URL encode the search string
                char *sstr = curl_easy_escape(curl_handle, inStr.c_str(), 0);

                string url = "http://search.twitter.com/search.json?q=";
                url += sstr;

                while(true) {
                    string contents = fetch_url(curl_handle, url.c_str());

                    bool success = reader.parse(contents, root, false);
                    if (!success) {
                        srvInterface.log("Malformed JSON text: %s", contents.c_str());
                        break;
                    }
                    // Get the query string, next page and results
                    string query = root.get("query", "-").asString();
                    url = "http://search.twitter.com/search.json";
                    url += root.get("next_page", "").asString();

                    Json::Value results = root.get("results",  Json::Value::null);

                    if (results.size() == 0)
                        break;

                    for (uint i=0; i<results.size(); i++) {
                        Json::Value result = results[i];
                        string text = result.get("text", "").asString();

                        // Copy string into results
                        output.getStringRef(0).copy(query);
                        output.getStringRef(1).copy(text);
                        output.next();
                    }
                }

                curl_free(sstr);
            } while (input.next());
        } catch (std::exception &e) {
            vt_report_error(0, e.what());
        }
    }
};

class TwitterSearchFactory : public TransformFunctionFactory
{
    // return an instance of the transform function
    virtual TransformFunction *createTransformFunction(ServerInterface &interface)
    { return vt_createFuncObj(interface.allocator, TwitterSearch); }

    virtual void getPrototype(ServerInterface &interface,
                              ColumnTypes &argTypes,
                              ColumnTypes &returnType)
    {
        argTypes.addVarchar();
        returnType.addVarchar();
        returnType.addVarchar();
    }

    virtual void getReturnType(ServerInterface &srvInterface,
                               const SizedColumnTypes &argTypes,
                               SizedColumnTypes &returnType)
    {
        const VerticaType &t = argTypes.getColumnType(0);
        returnType.addVarchar(t.getStringLength() + 50, "query"); // to account for URL encoding
        returnType.addVarchar(300, "text"); // tweet size <= 300
    }
};

RegisterFactory(TwitterSearchFactory);

/*
 * This is a simple function that gives the sentiment score of a piece of
 * text. The score is simply the number of positive words minus the number of
 * negative words in the text.
 */
class Sentiment : public ScalarFunction
{
    set<string> positive;
    set<string> negative;

    sb_stemmer *stemmer;

    static inline void tokenize(const string &text, vector<string> &tokens)
    {
        istringstream ss(text);

        do {
            std::string buffer;
            ss >> buffer;

            if (!buffer.empty())
                tokens.push_back(buffer);
        } while (ss);

    }
public:
    // Set up the positive and negative word sets
    virtual void setup(ServerInterface &srvInterface, const SizedColumnTypes &argTypes)
    {
        stemmer = sb_stemmer_new("english", NULL);

        int psize = sizeof(pwords)/sizeof(char*);
        int nsize = sizeof(nwords)/sizeof(char*);

        for (int i=0; i<psize; i++) {
            string stemword = reinterpret_cast<const char*>(sb_stemmer_stem(stemmer, reinterpret_cast<const sb_symbol*>(pwords[i]), strlen(pwords[i])));
            positive.insert(stemword);
        }

        for (int i=0; i<nsize; i++) {
            string stemword = reinterpret_cast<const char*>(sb_stemmer_stem(stemmer, reinterpret_cast<const sb_symbol*>(nwords[i]), strlen(nwords[i])));
            negative.insert(stemword);
        }
    }

    // Destroy the stemmer
    virtual void destroy(ServerInterface &srvInterface, const SizedColumnTypes &argTypes)
    {
        sb_stemmer_delete(stemmer);
    }

    /*
     * This method processes a block of rows in a single invocation.
     *
     * The inputs are retrieved via arg_reader
     * The outputs are returned via arg_writer
     */
    virtual void processBlock(ServerInterface &srvInterface,
                              BlockReader &arg_reader,
                              BlockWriter &res_writer)
    {
        // Basic error checking
        if (arg_reader.getNumCols() != 1)
            vt_report_error(0, "Function only accept 1 arguments, but %zu provided",
                            arg_reader.getNumCols());

        // While we have inputs to process
        do {
            // Get a copy of the input string
            std::string  inStr = arg_reader.getStringRef(0).str();

            vector<string> words;
            tokenize(inStr, words);

            int score = 0;

            for (uint i=0; i<words.size(); i++) {
                string stemword = reinterpret_cast<const char*>(sb_stemmer_stem(stemmer, reinterpret_cast<const sb_symbol*>(words[i].c_str()), words[i].size()));

                if (positive.count(stemword))
                    score++;
                if (negative.count(stemword))
                    score--;
            }

            // Copy score into results
            res_writer.setInt(score);
            res_writer.next();
        } while (arg_reader.next());
    }
};

class SentimentFactory : public ScalarFunctionFactory
{
    // return an instance of RemoveSpace to perform the actual addition.
    virtual ScalarFunction *createScalarFunction(ServerInterface &interface)
    { return vt_createFuncObj(interface.allocator, Sentiment); }

    virtual void getPrototype(ServerInterface &interface,
                              ColumnTypes &argTypes,
                              ColumnTypes &returnType)
    {
        argTypes.addVarchar();
        returnType.addInt();
    }
};

RegisterFactory(SentimentFactory);
<commit_msg>Rearrange includes to try and fix build problems<commit_after>/* Copyright (c) 2005 - 2011 Vertica, an HP company -*- C++ -*- */
/*
 *
 * Description: Functions that help with full-text search over documents
 *
 * Create Date: Nov 25, 2011
 */
#include <curl/curl.h>

#include <sstream>
#include <set>
#include "Vertica.h"
#include "libstemmer.h"
#include <json/json.h>
#include "Words.h"

using namespace Vertica;
using namespace std;

/*
 * This is a simple function that stems an input word
 */
class Stem : public ScalarFunction
{
    sb_stemmer *stemmer;

public:
    // Set up the stemmer
    virtual void setup(ServerInterface &srvInterface, const SizedColumnTypes &argTypes)
    {
        stemmer = sb_stemmer_new("english", NULL);
    }

    // Destroy the stemmer
    virtual void destroy(ServerInterface &srvInterface, const SizedColumnTypes &argTypes)
    {
        sb_stemmer_delete(stemmer);
    }

    /*
     * This method processes a block of rows in a single invocation.
     *
     * The inputs are retrieved via arg_reader
     * The outputs are returned via arg_writer
     */
    virtual void processBlock(ServerInterface &srvInterface,
                              BlockReader &arg_reader,
                              BlockWriter &res_writer)
    {
        // Basic error checking
        if (arg_reader.getNumCols() != 1)
            vt_report_error(0, "Function only accept 1 arguments, but %zu provided",
                            arg_reader.getNumCols());

        // While we have inputs to process
        do {
            // Get a copy of the input string
            std::string  inStr = arg_reader.getStringRef(0).str();

            const unsigned char* stemword = sb_stemmer_stem(stemmer, reinterpret_cast<const sb_symbol*>(inStr.c_str()), inStr.size());

            // Copy string into results
            res_writer.getStringRef().copy(reinterpret_cast<const char*>(stemword));
            res_writer.next();
        } while (arg_reader.next());
    }
};

class StemFactory : public ScalarFunctionFactory
{
    // return an instance of RemoveSpace to perform the actual addition.
    virtual ScalarFunction *createScalarFunction(ServerInterface &interface)
    { return vt_createFuncObj(interface.allocator, Stem); }

    virtual void getPrototype(ServerInterface &interface,
                              ColumnTypes &argTypes,
                              ColumnTypes &returnType)
    {
        argTypes.addVarchar();
        returnType.addVarchar();
    }

    virtual void getReturnType(ServerInterface &srvInterface,
                               const SizedColumnTypes &argTypes,
                               SizedColumnTypes &returnType)
    {
        const VerticaType &t = argTypes.getColumnType(0);
        returnType.addVarchar(t.getStringLength());
    }
};

RegisterFactory(StemFactory);

static size_t
WriteMemoryCallback(void *contents, size_t size, size_t nmemb, void *userp)
{
    size_t realsize = size * nmemb;
    stringstream *mem = (stringstream *)userp;

    mem->write(static_cast<const char*>(contents), realsize);

    return realsize;
}

static string fetch_url(CURL *curl_handle, const char *url)
{
    stringstream chunk;

    /* specify URL to get */
    curl_easy_setopt(curl_handle, CURLOPT_URL, url);

    /* send all data to this function  */
    curl_easy_setopt(curl_handle, CURLOPT_WRITEFUNCTION, WriteMemoryCallback);

    /* we pass our 'chunk' struct to the callback function */
    curl_easy_setopt(curl_handle, CURLOPT_WRITEDATA, (void *)&chunk);

    /* some servers don't like requests that are made without a user-agent
       field, so we provide one */
    curl_easy_setopt(curl_handle, CURLOPT_USERAGENT, "libcurl-agent/1.0");

    /* get it! */
    curl_easy_perform(curl_handle);

    return chunk.str();
}

/*
 * Fetch the given URL as a string
 */
class FetchURL : public ScalarFunction
{
    CURL *curl_handle;

public:
    virtual void setup(ServerInterface &srvInterface, const SizedColumnTypes &argTypes)
    {
        curl_global_init(CURL_GLOBAL_ALL);

        /* init the curl session */
        curl_handle = curl_easy_init();
    }

    // Destroy the stemmer
    virtual void destroy(ServerInterface &srvInterface, const SizedColumnTypes &argTypes)
    {
        /* cleanup curl stuff */
        curl_easy_cleanup(curl_handle);
        /* we're done with libcurl, so clean it up */
        curl_global_cleanup();
    }

    /*
     * This method processes a block of rows in a single invocation.
     *
     * The inputs are retrieved via arg_reader
     * The outputs are returned via arg_writer
     */
    virtual void processBlock(ServerInterface &srvInterface,
                              BlockReader &arg_reader,
                              BlockWriter &res_writer)
    {
        // Basic error checking
        if (arg_reader.getNumCols() != 1)
            vt_report_error(0, "Function only accept 1 arguments, but %zu provided",
                            arg_reader.getNumCols());

        try {
            // While we have inputs to process
            do {
                // Get a copy of the input string
                std::string  inStr = arg_reader.getStringRef(0).str();

                /* fetch contents of URL */
                string contents = fetch_url(curl_handle, inStr.c_str()).substr(0, 65000);

                // Copy string into results
                res_writer.getStringRef().copy(contents);
                res_writer.next();
            } while (arg_reader.next());
        } catch (std::exception &e) {
            vt_report_error(0, e.what());
        }
    }
};

class FetchURLFactory : public ScalarFunctionFactory
{
    // return an instance of RemoveSpace to perform the actual addition.
    virtual ScalarFunction *createScalarFunction(ServerInterface &interface)
    { return vt_createFuncObj(interface.allocator, FetchURL); }

    virtual void getPrototype(ServerInterface &interface,
                              ColumnTypes &argTypes,
                              ColumnTypes &returnType)
    {
        argTypes.addVarchar();
        returnType.addVarchar();
    }

    virtual void getReturnType(ServerInterface &srvInterface,
                               const SizedColumnTypes &argTypes,
                               SizedColumnTypes &returnType)
    {
        returnType.addVarchar(65000); // max varchar size
    }
};

RegisterFactory(FetchURLFactory);

/*
 * Parse the Twitter JSON results
 */
class TwitterJSONParser : public TransformFunction
{
    Json::Value root;   // will contains the root value after parsing.
    Json::Reader reader;

public:

    virtual void processPartition(ServerInterface &srvInterface,
                                  PartitionReader &input,
                                  PartitionWriter &output)
    {
        // Basic error checking
        if (input.getNumCols() != 1)
            vt_report_error(0, "Function only accept 1 arguments, but %zu provided",
                            input.getNumCols());

        // While we have inputs to process
        do {
            // Get a copy of the input string
            string  inStr = input.getStringRef(0).str();

            bool success = reader.parse(inStr, root, false);
            if (!success)
                vt_report_error(0, "Malformed JSON text");

            // Get the query string and results
            string query = root.get("query", "-").asString();
            Json::Value results = root.get("results",  Json::Value::null);

            for (uint i=0; i<results.size(); i++) {
                Json::Value result = results[i];
                string text = result.get("text", "").asString();

                // Copy string into results
                output.getStringRef(0).copy(query);
                output.getStringRef(1).copy(text);
                output.next();
            }
        } while (input.next());
    }
};

class TwitterJSONFactory : public TransformFunctionFactory
{
    // return an instance of the transform function
    virtual TransformFunction *createTransformFunction(ServerInterface &interface)
    { return vt_createFuncObj(interface.allocator, TwitterJSONParser); }

    virtual void getPrototype(ServerInterface &interface,
                              ColumnTypes &argTypes,
                              ColumnTypes &returnType)
    {
        argTypes.addVarchar();
        returnType.addVarchar();
        returnType.addVarchar();
    }

    virtual void getReturnType(ServerInterface &srvInterface,
                               const SizedColumnTypes &argTypes,
                               SizedColumnTypes &returnType)
    {
        returnType.addVarchar(100, "query"); // query string size <= 100
        returnType.addVarchar(300, "text"); // tweet size <= 300
    }
};

RegisterFactory(TwitterJSONFactory);

/*
 * Search Twitter and return results
 */
class TwitterSearch : public TransformFunction
{
    Json::Value root;   // will contains the root value after parsing.
    Json::Reader reader;
    CURL *curl_handle;

public:
    virtual void setup(ServerInterface &srvInterface, const SizedColumnTypes &argTypes)
    {
        curl_global_init(CURL_GLOBAL_ALL);

        /* init the curl session */
        curl_handle = curl_easy_init();
    }

    virtual void destroy(ServerInterface &srvInterface, const SizedColumnTypes &argTypes)
    {
        /* cleanup curl stuff */
        curl_easy_cleanup(curl_handle);
        /* we're done with libcurl, so clean it up */
        curl_global_cleanup();
    }

    virtual void processPartition(ServerInterface &srvInterface,
                                  PartitionReader &input,
                                  PartitionWriter &output)
    {
        // Basic error checking
        if (input.getNumCols() != 1)
            vt_report_error(0, "Function only accept 1 arguments, but %zu provided",
                            input.getNumCols());

        try {
            // While we have inputs to process
            do {
                // Get a copy of the input string
                string  inStr = input.getStringRef(0).str();

                // URL encode the search string
                char *sstr = curl_easy_escape(curl_handle, inStr.c_str(), 0);

                string url = "http://search.twitter.com/search.json?q=";
                url += sstr;

                while(true) {
                    string contents = fetch_url(curl_handle, url.c_str());

                    bool success = reader.parse(contents, root, false);
                    if (!success) {
                        srvInterface.log("Malformed JSON text: %s", contents.c_str());
                        break;
                    }
                    // Get the query string, next page and results
                    string query = root.get("query", "-").asString();
                    url = "http://search.twitter.com/search.json";
                    url += root.get("next_page", "").asString();

                    Json::Value results = root.get("results",  Json::Value::null);

                    if (results.size() == 0)
                        break;

                    for (uint i=0; i<results.size(); i++) {
                        Json::Value result = results[i];
                        string text = result.get("text", "").asString();

                        // Copy string into results
                        output.getStringRef(0).copy(query);
                        output.getStringRef(1).copy(text);
                        output.next();
                    }
                }

                curl_free(sstr);
            } while (input.next());
        } catch (std::exception &e) {
            vt_report_error(0, e.what());
        }
    }
};

class TwitterSearchFactory : public TransformFunctionFactory
{
    // return an instance of the transform function
    virtual TransformFunction *createTransformFunction(ServerInterface &interface)
    { return vt_createFuncObj(interface.allocator, TwitterSearch); }

    virtual void getPrototype(ServerInterface &interface,
                              ColumnTypes &argTypes,
                              ColumnTypes &returnType)
    {
        argTypes.addVarchar();
        returnType.addVarchar();
        returnType.addVarchar();
    }

    virtual void getReturnType(ServerInterface &srvInterface,
                               const SizedColumnTypes &argTypes,
                               SizedColumnTypes &returnType)
    {
        const VerticaType &t = argTypes.getColumnType(0);
        returnType.addVarchar(t.getStringLength() + 50, "query"); // to account for URL encoding
        returnType.addVarchar(300, "text"); // tweet size <= 300
    }
};

RegisterFactory(TwitterSearchFactory);

/*
 * This is a simple function that gives the sentiment score of a piece of
 * text. The score is simply the number of positive words minus the number of
 * negative words in the text.
 */
class Sentiment : public ScalarFunction
{
    set<string> positive;
    set<string> negative;

    sb_stemmer *stemmer;

    static inline void tokenize(const string &text, vector<string> &tokens)
    {
        istringstream ss(text);

        do {
            std::string buffer;
            ss >> buffer;

            if (!buffer.empty())
                tokens.push_back(buffer);
        } while (ss);

    }
public:
    // Set up the positive and negative word sets
    virtual void setup(ServerInterface &srvInterface, const SizedColumnTypes &argTypes)
    {
        stemmer = sb_stemmer_new("english", NULL);

        int psize = sizeof(pwords)/sizeof(char*);
        int nsize = sizeof(nwords)/sizeof(char*);

        for (int i=0; i<psize; i++) {
            string stemword = reinterpret_cast<const char*>(sb_stemmer_stem(stemmer, reinterpret_cast<const sb_symbol*>(pwords[i]), strlen(pwords[i])));
            positive.insert(stemword);
        }

        for (int i=0; i<nsize; i++) {
            string stemword = reinterpret_cast<const char*>(sb_stemmer_stem(stemmer, reinterpret_cast<const sb_symbol*>(nwords[i]), strlen(nwords[i])));
            negative.insert(stemword);
        }
    }

    // Destroy the stemmer
    virtual void destroy(ServerInterface &srvInterface, const SizedColumnTypes &argTypes)
    {
        sb_stemmer_delete(stemmer);
    }

    /*
     * This method processes a block of rows in a single invocation.
     *
     * The inputs are retrieved via arg_reader
     * The outputs are returned via arg_writer
     */
    virtual void processBlock(ServerInterface &srvInterface,
                              BlockReader &arg_reader,
                              BlockWriter &res_writer)
    {
        // Basic error checking
        if (arg_reader.getNumCols() != 1)
            vt_report_error(0, "Function only accept 1 arguments, but %zu provided",
                            arg_reader.getNumCols());

        // While we have inputs to process
        do {
            // Get a copy of the input string
            std::string  inStr = arg_reader.getStringRef(0).str();

            vector<string> words;
            tokenize(inStr, words);

            int score = 0;

            for (uint i=0; i<words.size(); i++) {
                string stemword = reinterpret_cast<const char*>(sb_stemmer_stem(stemmer, reinterpret_cast<const sb_symbol*>(words[i].c_str()), words[i].size()));

                if (positive.count(stemword))
                    score++;
                if (negative.count(stemword))
                    score--;
            }

            // Copy score into results
            res_writer.setInt(score);
            res_writer.next();
        } while (arg_reader.next());
    }
};

class SentimentFactory : public ScalarFunctionFactory
{
    // return an instance of RemoveSpace to perform the actual addition.
    virtual ScalarFunction *createScalarFunction(ServerInterface &interface)
    { return vt_createFuncObj(interface.allocator, Sentiment); }

    virtual void getPrototype(ServerInterface &interface,
                              ColumnTypes &argTypes,
                              ColumnTypes &returnType)
    {
        argTypes.addVarchar();
        returnType.addInt();
    }
};

RegisterFactory(SentimentFactory);
<|endoftext|>


<commit_before><commit_msg>INTEGRATION: CWS vclcleanup01 (1.12.6); FILE MERGED 2003/11/28 07:32:45 mt 1.12.6.1: #i22952# Removed App Server code<commit_after><|endoftext|>


<commit_before>/*
*   Copyright (c) 2010-2016, MIT Probabilistic Computing Project
*
*   Lead Developers: Dan Lovell and Jay Baxter
*   Authors: Dan Lovell, Baxter Eaves, Jay Baxter, Vikash Mansinghka
*   Research Leads: Vikash Mansinghka, Patrick Shafto
*
*   Licensed under the Apache License, Version 2.0 (the "License");
*   you may not use this file except in compliance with the License.
*   You may obtain a copy of the License at
*
*       http://www.apache.org/licenses/LICENSE-2.0
*
*   Unless required by applicable law or agreed to in writing, software
*   distributed under the License is distributed on an "AS IS" BASIS,
*   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
*   See the License for the specific language governing permissions and
*   limitations under the License.
*/
#include "MultinomialComponentModel.h"
#include <boost/random/uniform_01.hpp>

using namespace std;

MultinomialComponentModel::MultinomialComponentModel(const CM_Hypers& in_hypers) {
    count = 0;
    score = 0;
    p_hypers = &in_hypers;
    hyper_K = get(*p_hypers, (string) "K");
    hyper_dirichlet_alpha = get(*p_hypers, (string) "dirichlet_alpha");
    init_suffstats();
    set_log_Z_0();
}

MultinomialComponentModel::MultinomialComponentModel(const CM_Hypers& in_hypers,
        int count_in,
        const map<string, double>& counts) {
    count = 0;
    score = 0;
    p_hypers = &in_hypers;
    hyper_K = get(*p_hypers, (string) "K");
    hyper_dirichlet_alpha = get(*p_hypers, (string) "dirichlet_alpha");
    set_log_Z_0();
    // set suffstats
    count = count_in;
    init_suffstats();
    for (map<string, double>::const_iterator it = counts.begin();
         it != counts.end();
         ++it) {
        int i = intify(it->first);
        assert(0 <= i);
        assert(i < hyper_K);
        assert(0 <= it->second);
        assert(it->second == trunc(it->second));
        suffstats[i] = static_cast<int>(it->second);
    }
    score = calc_marginal_logp();
}

void MultinomialComponentModel::init_suffstats() {
    suffstats.resize(hyper_K);
}

double MultinomialComponentModel::calc_marginal_logp() const {
    const vector<int>& counts = suffstats;
    int K = hyper_K;
    double dirichlet_alpha = hyper_dirichlet_alpha;
    return numerics::calc_multinomial_marginal_logp(count, counts, K,
            dirichlet_alpha);
}

double MultinomialComponentModel::calc_element_predictive_logp(
    double element) const {
    if (isnan(element)) {
        return 0;
    }
    int K = hyper_K;
    double dirichlet_alpha = hyper_dirichlet_alpha;
    double logp = numerics::calc_multinomial_predictive_logp(element,
                  suffstats, count,
                  K, dirichlet_alpha);
    return logp;
}

double MultinomialComponentModel::calc_element_predictive_logp_constrained(
    double element, const vector<double>& constraints) const {
    if (isnan(element)) {
        return 0;
    }
    int K = hyper_K;
    double dirichlet_alpha = hyper_dirichlet_alpha;
    //
    vector<int> suffstats_copy = suffstats;
    int count_copy = count;
    int num_constraints = (int) constraints.size();
    for (int constraint_idx = 0; constraint_idx < num_constraints;
            constraint_idx++) {
        double constraint = constraints[constraint_idx];
        assert(0 <= constraint);
        assert(constraint < K);
        assert(constraint == trunc(constraint));
        int i = static_cast<int>(constraint);
        count_copy++;
        suffstats_copy[i]++;
    }
    double predictive = \
                        numerics::calc_multinomial_predictive_logp(element,
                                suffstats_copy,
                                count_copy,
                                K, dirichlet_alpha);
    return predictive;
}

vector<double> MultinomialComponentModel::calc_hyper_conditionals(
    const string& which_hyper, const vector<double>& hyper_grid) const {
    const vector<int>& counts = suffstats;
    int K = hyper_K;
    if (which_hyper == "dirichlet_alpha") {
        return numerics::calc_multinomial_dirichlet_alpha_conditional(hyper_grid,
                count,
                counts,
                K);
    } else {
        // error condition
        cout << "MultinomialComponentModel::calc_hyper_conditional: bad value for which_hyper="
             << which_hyper << endl;
        assert(0);
        vector<double> vd;
        return vd;
    }
}

double MultinomialComponentModel::insert_element(double element) {
    if (isnan(element)) {
        return 0;
    }
    assert(element == trunc(element));
    int i = static_cast<int>(element);
    double delta_score = calc_element_predictive_logp(element);
    suffstats[i] += 1;
    count += 1;
    score += delta_score;
    return delta_score;
}

double MultinomialComponentModel::remove_element(double element) {
    if (isnan(element)) {
        return 0;
    }
    assert(element == trunc(element));
    int i = static_cast<int>(element);
    assert(0 < suffstats[i]);
    suffstats[i] -= 1;
    double delta_score = calc_element_predictive_logp(element);
    count -= 1;
    score -= delta_score;
    return delta_score;
}

double MultinomialComponentModel::incorporate_hyper_update() {
    hyper_K = get(*p_hypers, (string) "K");
    hyper_dirichlet_alpha = get(*p_hypers, (string) "dirichlet_alpha");
    double score_0 = score;
    // hypers[which_hyper] = value; // set by owner of hypers object
    score = calc_marginal_logp();
    double score_delta = score - score_0;
    return score_delta;
}

void MultinomialComponentModel::set_log_Z_0() {
    log_Z_0 = calc_marginal_logp();
}

void MultinomialComponentModel::get_keys_counts_for_draw(vector<int>& keys,
        vector<double>& log_counts_for_draw,
        const vector<int>& counts) const {
    double dirichlet_alpha = hyper_dirichlet_alpha;
    for (int key = 0; key < hyper_K; key++) {
        int count_for_draw = counts[key];
        // "update" counts by adding dirichlet alpha to each value
        count_for_draw += dirichlet_alpha;
        keys.push_back(key);
        log_counts_for_draw.push_back(log(count_for_draw));
    }
    return;
}

double MultinomialComponentModel::get_draw(int random_seed) const {
    // get modified suffstats
    const vector<int>& counts = suffstats;
    // get a random draw
    boost::mt19937  _engine(random_seed);
    boost::uniform_01<boost::mt19937> _dist(_engine);
    double uniform_draw = _dist();
    //
    vector<int> keys;
    vector<double> log_counts_for_draw;
    get_keys_counts_for_draw(keys, log_counts_for_draw, counts);
    //
    int key_idx = numerics::draw_sample_unnormalized(log_counts_for_draw,
                  uniform_draw);
    double draw = static_cast<double>(keys[key_idx]);
    return draw;
}

double MultinomialComponentModel::get_draw_constrained(int random_seed,
        const vector<double>& constraints) const {
    // get modified suffstats
    const vector<int>& counts = suffstats;
    // get a random draw
    boost::mt19937  _engine(random_seed);
    boost::uniform_01<boost::mt19937> _dist(_engine);
    double uniform_draw = _dist();
    //
    vector<int> keys;
    vector<double> log_counts_for_draw;
    get_keys_counts_for_draw(keys, log_counts_for_draw, counts);
    map<int, int> index_lookup = construct_lookup_map(keys);
    int num_constraints = (int) constraints.size();
    for (int constraint_idx = 0; constraint_idx < num_constraints;
            constraint_idx++) {
        double constraint = constraints[constraint_idx];
        assert(0 <= constraint);
        assert(constraint == trunc(constraint));
        int count_idx = index_lookup[static_cast<int>(constraint)];
        double log_count = log_counts_for_draw[count_idx];
        log_counts_for_draw[count_idx] = log(exp(log_count) + 1);
    }
    //
    int key_idx = numerics::draw_sample_unnormalized(log_counts_for_draw,
                  uniform_draw);
    double draw = static_cast<double>(keys[key_idx]);
    return draw;
}

map<string, double> MultinomialComponentModel::_get_suffstats() const {
    map<string, double> counts;
    for (int key = 0; key < hyper_K; key++) {
        counts[stringify(key)] = suffstats[key];
    }
    return counts;
}

map<string, double> MultinomialComponentModel::get_hypers() const {
    map<string, double> hypers;
    hypers["K"] = hyper_K;
    hypers["dirichlet_alpha"] = hyper_dirichlet_alpha;
    return hypers;
}

void MultinomialComponentModel::get_suffstats(int& count_out,
        map<string, double>& counts) const {
    count_out = count;
    counts = _get_suffstats();
}
<commit_msg>Use RandomNumberGenerator in MultinomialComponentModel.<commit_after>/*
*   Copyright (c) 2010-2016, MIT Probabilistic Computing Project
*
*   Lead Developers: Dan Lovell and Jay Baxter
*   Authors: Dan Lovell, Baxter Eaves, Jay Baxter, Vikash Mansinghka
*   Research Leads: Vikash Mansinghka, Patrick Shafto
*
*   Licensed under the Apache License, Version 2.0 (the "License");
*   you may not use this file except in compliance with the License.
*   You may obtain a copy of the License at
*
*       http://www.apache.org/licenses/LICENSE-2.0
*
*   Unless required by applicable law or agreed to in writing, software
*   distributed under the License is distributed on an "AS IS" BASIS,
*   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
*   See the License for the specific language governing permissions and
*   limitations under the License.
*/

#include "RandomNumberGenerator.h"

#include "MultinomialComponentModel.h"

using namespace std;

MultinomialComponentModel::MultinomialComponentModel(const CM_Hypers& in_hypers) {
    count = 0;
    score = 0;
    p_hypers = &in_hypers;
    hyper_K = get(*p_hypers, (string) "K");
    hyper_dirichlet_alpha = get(*p_hypers, (string) "dirichlet_alpha");
    init_suffstats();
    set_log_Z_0();
}

MultinomialComponentModel::MultinomialComponentModel(const CM_Hypers& in_hypers,
        int count_in,
        const map<string, double>& counts) {
    count = 0;
    score = 0;
    p_hypers = &in_hypers;
    hyper_K = get(*p_hypers, (string) "K");
    hyper_dirichlet_alpha = get(*p_hypers, (string) "dirichlet_alpha");
    set_log_Z_0();
    // set suffstats
    count = count_in;
    init_suffstats();
    for (map<string, double>::const_iterator it = counts.begin();
         it != counts.end();
         ++it) {
        int i = intify(it->first);
        assert(0 <= i);
        assert(i < hyper_K);
        assert(0 <= it->second);
        assert(it->second == trunc(it->second));
        suffstats[i] = static_cast<int>(it->second);
    }
    score = calc_marginal_logp();
}

void MultinomialComponentModel::init_suffstats() {
    suffstats.resize(hyper_K);
}

double MultinomialComponentModel::calc_marginal_logp() const {
    const vector<int>& counts = suffstats;
    int K = hyper_K;
    double dirichlet_alpha = hyper_dirichlet_alpha;
    return numerics::calc_multinomial_marginal_logp(count, counts, K,
            dirichlet_alpha);
}

double MultinomialComponentModel::calc_element_predictive_logp(
    double element) const {
    if (isnan(element)) {
        return 0;
    }
    int K = hyper_K;
    double dirichlet_alpha = hyper_dirichlet_alpha;
    double logp = numerics::calc_multinomial_predictive_logp(element,
                  suffstats, count,
                  K, dirichlet_alpha);
    return logp;
}

double MultinomialComponentModel::calc_element_predictive_logp_constrained(
    double element, const vector<double>& constraints) const {
    if (isnan(element)) {
        return 0;
    }
    int K = hyper_K;
    double dirichlet_alpha = hyper_dirichlet_alpha;
    //
    vector<int> suffstats_copy = suffstats;
    int count_copy = count;
    int num_constraints = (int) constraints.size();
    for (int constraint_idx = 0; constraint_idx < num_constraints;
            constraint_idx++) {
        double constraint = constraints[constraint_idx];
        assert(0 <= constraint);
        assert(constraint < K);
        assert(constraint == trunc(constraint));
        int i = static_cast<int>(constraint);
        count_copy++;
        suffstats_copy[i]++;
    }
    double predictive = \
                        numerics::calc_multinomial_predictive_logp(element,
                                suffstats_copy,
                                count_copy,
                                K, dirichlet_alpha);
    return predictive;
}

vector<double> MultinomialComponentModel::calc_hyper_conditionals(
    const string& which_hyper, const vector<double>& hyper_grid) const {
    const vector<int>& counts = suffstats;
    int K = hyper_K;
    if (which_hyper == "dirichlet_alpha") {
        return numerics::calc_multinomial_dirichlet_alpha_conditional(hyper_grid,
                count,
                counts,
                K);
    } else {
        // error condition
        cout << "MultinomialComponentModel::calc_hyper_conditional: bad value for which_hyper="
             << which_hyper << endl;
        assert(0);
        vector<double> vd;
        return vd;
    }
}

double MultinomialComponentModel::insert_element(double element) {
    if (isnan(element)) {
        return 0;
    }
    assert(element == trunc(element));
    int i = static_cast<int>(element);
    double delta_score = calc_element_predictive_logp(element);
    suffstats[i] += 1;
    count += 1;
    score += delta_score;
    return delta_score;
}

double MultinomialComponentModel::remove_element(double element) {
    if (isnan(element)) {
        return 0;
    }
    assert(element == trunc(element));
    int i = static_cast<int>(element);
    assert(0 < suffstats[i]);
    suffstats[i] -= 1;
    double delta_score = calc_element_predictive_logp(element);
    count -= 1;
    score -= delta_score;
    return delta_score;
}

double MultinomialComponentModel::incorporate_hyper_update() {
    hyper_K = get(*p_hypers, (string) "K");
    hyper_dirichlet_alpha = get(*p_hypers, (string) "dirichlet_alpha");
    double score_0 = score;
    // hypers[which_hyper] = value; // set by owner of hypers object
    score = calc_marginal_logp();
    double score_delta = score - score_0;
    return score_delta;
}

void MultinomialComponentModel::set_log_Z_0() {
    log_Z_0 = calc_marginal_logp();
}

void MultinomialComponentModel::get_keys_counts_for_draw(vector<int>& keys,
        vector<double>& log_counts_for_draw,
        const vector<int>& counts) const {
    double dirichlet_alpha = hyper_dirichlet_alpha;
    for (int key = 0; key < hyper_K; key++) {
        int count_for_draw = counts[key];
        // "update" counts by adding dirichlet alpha to each value
        count_for_draw += dirichlet_alpha;
        keys.push_back(key);
        log_counts_for_draw.push_back(log(count_for_draw));
    }
    return;
}

double MultinomialComponentModel::get_draw(int random_seed) const {
    // get modified suffstats
    const vector<int>& counts = suffstats;
    // get a random draw
    double uniform_draw = RandomNumberGenerator(random_seed).next();
    //
    vector<int> keys;
    vector<double> log_counts_for_draw;
    get_keys_counts_for_draw(keys, log_counts_for_draw, counts);
    //
    int key_idx = numerics::draw_sample_unnormalized(log_counts_for_draw,
                  uniform_draw);
    double draw = static_cast<double>(keys[key_idx]);
    return draw;
}

double MultinomialComponentModel::get_draw_constrained(int random_seed,
        const vector<double>& constraints) const {
    // get modified suffstats
    const vector<int>& counts = suffstats;
    // get a random draw
    double uniform_draw = RandomNumberGenerator(random_seed).next();
    //
    vector<int> keys;
    vector<double> log_counts_for_draw;
    get_keys_counts_for_draw(keys, log_counts_for_draw, counts);
    map<int, int> index_lookup = construct_lookup_map(keys);
    int num_constraints = (int) constraints.size();
    for (int constraint_idx = 0; constraint_idx < num_constraints;
            constraint_idx++) {
        double constraint = constraints[constraint_idx];
        assert(0 <= constraint);
        assert(constraint == trunc(constraint));
        int count_idx = index_lookup[static_cast<int>(constraint)];
        double log_count = log_counts_for_draw[count_idx];
        log_counts_for_draw[count_idx] = log(exp(log_count) + 1);
    }
    //
    int key_idx = numerics::draw_sample_unnormalized(log_counts_for_draw,
                  uniform_draw);
    double draw = static_cast<double>(keys[key_idx]);
    return draw;
}

map<string, double> MultinomialComponentModel::_get_suffstats() const {
    map<string, double> counts;
    for (int key = 0; key < hyper_K; key++) {
        counts[stringify(key)] = suffstats[key];
    }
    return counts;
}

map<string, double> MultinomialComponentModel::get_hypers() const {
    map<string, double> hypers;
    hypers["K"] = hyper_K;
    hypers["dirichlet_alpha"] = hyper_dirichlet_alpha;
    return hypers;
}

void MultinomialComponentModel::get_suffstats(int& count_out,
        map<string, double>& counts) const {
    count_out = count;
    counts = _get_suffstats();
}
<|endoftext|>


<commit_before>/// HEADER
#include "transform_cloud.h"

/// PROJECT
#include <csapex/msg/io.h>
#include <csapex_transform/transform_message.h>
#include <csapex/model/node_modifier.h>
#include <csapex/utility/register_apex_plugin.h>

/// SYSTEM
#include <boost/mpl/for_each.hpp>
#define BOOST_SIGNALS_NO_DEPRECATION_WARNING
#include <pcl_ros/transforms.h>

CSAPEX_REGISTER_CLASS(csapex::TransformCloud, csapex::Node)

using namespace csapex;
using namespace csapex::connection_types;

TransformCloud::TransformCloud()
{
}

void TransformCloud::setup(NodeModifier& node_modifier)
{
    input_cloud_ = node_modifier.addInput<PointCloudMessage>("PointCloud");
    input_transform_ = node_modifier.addInput<TransformMessage>("Transformation");

    output_ = node_modifier.addOutput<PointCloudMessage>("PointCloud");
}

void TransformCloud::process()
{
    PointCloudMessage::ConstPtr msg = msg::getMessage<PointCloudMessage>(input_cloud_);

    boost::apply_visitor (PointCloudMessage::Dispatch<TransformCloud>(this, msg), msg->value);
}

template <class PointT>
void TransformCloud::inputCloud(typename pcl::PointCloud<PointT>::ConstPtr cloud)
{
    TransformMessage::ConstPtr transform = msg::getMessage<TransformMessage>(input_transform_);
    const tf::Transform& t = transform->value;

    typename pcl::PointCloud<PointT>::Ptr out(new pcl::PointCloud<PointT>);
    pcl_ros::transformPointCloud(*cloud, *out, t);

    std::string frame = cloud->header.frame_id;

    if(frame == transform->child_frame) {
        frame = transform->frame_id;
    }

    PointCloudMessage::Ptr msg(new PointCloudMessage(frame, cloud->header.stamp));
    msg->value = out;

    msg::publish(output_, msg);
}
<commit_msg>manual cloud transform to avoid pcl_ros<commit_after>/// HEADER
#include "transform_cloud.h"

/// PROJECT
#include <csapex/msg/io.h>
#include <csapex_transform/transform_message.h>
#include <csapex/model/node_modifier.h>
#include <csapex/utility/register_apex_plugin.h>

/// SYSTEM
#include <boost/mpl/for_each.hpp>

CSAPEX_REGISTER_CLASS(csapex::TransformCloud, csapex::Node)

using namespace csapex;
using namespace csapex::connection_types;

TransformCloud::TransformCloud()
{
}

void TransformCloud::setup(NodeModifier& node_modifier)
{
    input_cloud_ = node_modifier.addInput<PointCloudMessage>("PointCloud");
    input_transform_ = node_modifier.addInput<TransformMessage>("Transformation");

    output_ = node_modifier.addOutput<PointCloudMessage>("PointCloud");
}

void TransformCloud::process()
{
    PointCloudMessage::ConstPtr msg = msg::getMessage<PointCloudMessage>(input_cloud_);

    boost::apply_visitor (PointCloudMessage::Dispatch<TransformCloud>(this, msg), msg->value);
}

template <class PointT>
void TransformCloud::inputCloud(typename pcl::PointCloud<PointT>::ConstPtr cloud)
{
    TransformMessage::ConstPtr transform = msg::getMessage<TransformMessage>(input_transform_);
    const tf::Transform& t = transform->value;

    typename pcl::PointCloud<PointT>::Ptr out(new pcl::PointCloud<PointT>);
    out->header = cloud->header;
    out->width = cloud->width;
    out->height = cloud->height;
    out->is_dense = cloud->is_dense;

    std::size_t N = cloud->points.size();
    out->points.resize(N);

    for(std::size_t i = 0; i < N; ++i) {
        PointT pt = cloud->points[i];
        tf::Quaternion q = t.getRotation();
        tf::Vector3 tr = t.getOrigin();

        tf::Vector3 p(pt.x, pt.y, pt.z);

//        tf::Vector3 tfd =  tf::quatRotate(q, p) + tr;
        tf::Vector3 tfd = t * p;
        pt.x = tfd.x();
        pt.y = tfd.y();
        pt.z = tfd.z();
        out->points[i] = pt;
    }

    std::string frame = cloud->header.frame_id;

    if(frame == transform->child_frame) {
        frame = transform->frame_id;
    }

    PointCloudMessage::Ptr msg(new PointCloudMessage(frame, cloud->header.stamp));
    msg->value = out;

    msg::publish(output_, msg);
}
<|endoftext|>


<commit_before>#define BOOST_TEST_MODULE "test_unlimited_cell_list"

#ifdef BOOST_TEST_DYN_LINK
#include <boost/test/unit_test.hpp>
#else
#include <boost/test/included/unit_test.hpp>
#endif

#include <mjolnir/util/empty.hpp>
#include <mjolnir/core/SimulatorTraits.hpp>
#include <mjolnir/core/BoundaryCondition.hpp>
#include <mjolnir/core/UnlimitedGridCellList.hpp>
#include <mjolnir/core/System.hpp>
#include <mjolnir/core/Topology.hpp>
#include <numeric>
#include <random>

template<typename T>
struct dummy_potential
{
    using real_type      = T;
    using parameter_type = mjolnir::empty_t;

    using topology_type        = mjolnir::Topology;
    using molecule_id_type     = typename topology_type::molecule_id_type;
    using connection_kind_type = typename topology_type::connection_kind_type;

    explicit dummy_potential(const real_type cutoff,
                             const std::vector<std::size_t>& participants)
        : cutoff_(cutoff), participants_(participants)
    {}

    real_type max_cutoff_length() const noexcept {return this->cutoff_;}

    parameter_type prepare_params(std::size_t, std::size_t) const noexcept
    {
        return parameter_type{};
    }

    bool is_ignored_molecule(std::size_t, std::size_t) const {return false;}

    std::vector<std::pair<connection_kind_type, std::size_t>> ignore_within() const
    {
        return std::vector<std::pair<connection_kind_type, std::size_t>>{};
    }

    std::vector<std::size_t> const& participants() const noexcept
    {
        return this->participants_;
    }

    std::string name() const {return "dummy potential";}

    real_type cutoff_;
    std::vector<std::size_t> participants_;
};

BOOST_AUTO_TEST_CASE(test_CellList_UnlimitedBoundary)
{
    mjolnir::LoggerManager::set_default_logger("test_cell_list.log");
    using traits_type     = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using real_type       = typename traits_type::real_type;
    using boundary_type   = typename traits_type::boundary_type;
    using coordinate_type = typename traits_type::coordinate_type;
    using parameter_type  = typename dummy_potential<real_type>::parameter_type;

    constexpr std::size_t N = 1000;
    constexpr double      L = 10.0;
    constexpr double cutoff = 2.0;
    constexpr double margin = 0.25; // threshold = 2.0 * (1+0.25) = 2.5
    constexpr double threshold = cutoff * (1.0 + margin);

    const auto distribute_particle = [](std::mt19937& mt, double l) -> coordinate_type
    {
        return coordinate_type(
            l * std::generate_canonical<real_type, std::numeric_limits<real_type>::digits>(mt),
            l * std::generate_canonical<real_type, std::numeric_limits<real_type>::digits>(mt),
            l * std::generate_canonical<real_type, std::numeric_limits<real_type>::digits>(mt)
        );
    };

    std::vector<std::size_t> participants(N);
    std::iota(participants.begin(), participants.end(), 0u);

    dummy_potential<real_type> pot(cutoff, participants);

    mjolnir::System<traits_type> sys(N, boundary_type{});

    std::mt19937 mt(123456789);
    for(std::size_t i=0; i < N; ++i)
    {
        sys.at(i).mass     = 1.0;
        sys.at(i).position = distribute_particle(mt, L);
    }
    sys.topology().construct_molecules();

    mjolnir::UnlimitedGridCellList<traits_type, parameter_type> vlist(margin);

    using neighbor_type = typename decltype(vlist)::neighbor_type;

    BOOST_TEST(!vlist.valid());

    vlist.initialize(sys, pot);
    vlist.make(sys, pot);
    BOOST_TEST(vlist.valid());

    for(std::size_t i=0; i<N; ++i)
    {
        for(std::size_t j=i+1; j<N; ++j)
        {
            const auto partners = vlist.partners(i);
            if(std::find_if(partners.begin(), partners.end(),
                [=](const neighbor_type& elem) -> bool {return elem.index == j;}
                        ) == partners.end())
            {
                // should be enough distant (>= threshold)
                const auto dist = mjolnir::math::length(sys.adjust_direction(
                            sys.position(j) - sys.position(i)));
                BOOST_TEST(dist >= threshold);
            }
            else
            {
                // should be enough close (< threshold)
                const auto dist = mjolnir::math::length(sys.adjust_direction(
                            sys.position(j) - sys.position(i)));
                BOOST_TEST(dist < threshold);
            }
        }
    }
}
<commit_msg>test: add test to unlimited_cell_list<commit_after>#define BOOST_TEST_MODULE "test_unlimited_cell_list"

#ifdef BOOST_TEST_DYN_LINK
#include <boost/test/unit_test.hpp>
#else
#include <boost/test/included/unit_test.hpp>
#endif

#include <mjolnir/util/empty.hpp>
#include <mjolnir/core/SimulatorTraits.hpp>
#include <mjolnir/core/BoundaryCondition.hpp>
#include <mjolnir/core/UnlimitedGridCellList.hpp>
#include <mjolnir/core/System.hpp>
#include <mjolnir/core/Topology.hpp>
#include <numeric>
#include <random>

template<typename T>
struct dummy_potential
{
    using real_type      = T;
    using parameter_type = mjolnir::empty_t;

    using topology_type        = mjolnir::Topology;
    using molecule_id_type     = typename topology_type::molecule_id_type;
    using connection_kind_type = typename topology_type::connection_kind_type;

    explicit dummy_potential(const real_type cutoff,
                             const std::vector<std::size_t>& participants)
        : cutoff_(cutoff), participants_(participants)
    {}

    real_type max_cutoff_length() const noexcept {return this->cutoff_;}

    parameter_type prepare_params(std::size_t, std::size_t) const noexcept
    {
        return parameter_type{};
    }

    bool is_ignored_molecule(std::size_t, std::size_t) const {return false;}

    std::vector<std::pair<connection_kind_type, std::size_t>> ignore_within() const
    {
        return std::vector<std::pair<connection_kind_type, std::size_t>>{};
    }

    std::vector<std::size_t> const& participants() const noexcept
    {
        return this->participants_;
    }

    std::string name() const {return "dummy potential";}

    real_type cutoff_;
    std::vector<std::size_t> participants_;
};

BOOST_AUTO_TEST_CASE(test_CellList_UnlimitedBoundary)
{
    mjolnir::LoggerManager::set_default_logger("test_cell_list.log");
    using traits_type     = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using real_type       = typename traits_type::real_type;
    using boundary_type   = typename traits_type::boundary_type;
    using coordinate_type = typename traits_type::coordinate_type;
    using parameter_type  = typename dummy_potential<real_type>::parameter_type;

    constexpr std::size_t N = 1000;
    constexpr double      L = 10.0;
    constexpr double cutoff = 2.0;
    constexpr double margin = 0.25; // threshold = 2.0 * (1+0.25) = 2.5
    constexpr double threshold = cutoff * (1.0 + margin);

    const auto distribute_particle = [](std::mt19937& mt, double l) -> coordinate_type
    {
        return coordinate_type(
            l * std::generate_canonical<real_type, std::numeric_limits<real_type>::digits>(mt),
            l * std::generate_canonical<real_type, std::numeric_limits<real_type>::digits>(mt),
            l * std::generate_canonical<real_type, std::numeric_limits<real_type>::digits>(mt)
        );
    };

    std::vector<std::size_t> participants(N);
    std::iota(participants.begin(), participants.end(), 0u);

    dummy_potential<real_type> pot(cutoff, participants);

    mjolnir::System<traits_type> sys(N, boundary_type{});

    std::mt19937 mt(123456789);
    for(std::size_t i=0; i < N; ++i)
    {
        sys.at(i).mass     = 1.0;
        sys.at(i).position = distribute_particle(mt, L);
    }
    sys.topology().construct_molecules();

    mjolnir::UnlimitedGridCellList<traits_type, parameter_type> vlist(margin);

    using neighbor_type = typename decltype(vlist)::neighbor_type;

    BOOST_TEST(!vlist.valid());

    vlist.initialize(sys, pot);
    vlist.make(sys, pot);
    BOOST_TEST(vlist.valid());

    for(std::size_t i=0; i<N; ++i)
    {
        for(std::size_t j=i+1; j<N; ++j)
        {
            const auto partners = vlist.partners(i);
            if(std::find_if(partners.begin(), partners.end(),
                [=](const neighbor_type& elem) -> bool {return elem.index == j;}
                        ) == partners.end())
            {
                // should be enough distant (>= threshold)
                const auto dist = mjolnir::math::length(sys.adjust_direction(
                            sys.position(j) - sys.position(i)));
                BOOST_TEST(dist >= threshold);
            }
            else
            {
                // should be enough close (< threshold)
                const auto dist = mjolnir::math::length(sys.adjust_direction(
                            sys.position(j) - sys.position(i)));
                BOOST_TEST(dist < threshold);
            }
        }
    }
}

BOOST_AUTO_TEST_CASE(test_CellList_UnlimitedBoundary_partial)
{
    mjolnir::LoggerManager::set_default_logger("test_cell_list.log");

    using traits_type     = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using real_type       = typename traits_type::real_type;
    using boundary_type   = typename traits_type::boundary_type;
    using coordinate_type = typename traits_type::coordinate_type;
    using parameter_type  = typename dummy_potential<real_type>::parameter_type;

    constexpr std::size_t N = 1000;
    constexpr double      L = 10.0;
    constexpr double cutoff = 2.0;
    constexpr double margin = 0.25; // threshold = 2.0 * (1+0.25) = 2.5
    constexpr double threshold = cutoff * (1.0 + margin);

    const auto distribute_particle = [](std::mt19937& mt, double l) -> coordinate_type
    {
        return coordinate_type(
            l * std::generate_canonical<real_type, std::numeric_limits<real_type>::digits>(mt),
            l * std::generate_canonical<real_type, std::numeric_limits<real_type>::digits>(mt),
            l * std::generate_canonical<real_type, std::numeric_limits<real_type>::digits>(mt)
        );
    };

    std::vector<std::size_t> participants(500);
    // [200 ... 699]
    std::iota(participants.begin(), participants.end(), 200u);

    dummy_potential<real_type> pot(cutoff, participants);

    mjolnir::System<traits_type> sys(N, boundary_type{});

    std::mt19937 mt(123456789);
    for(std::size_t i=0; i < N; ++i)
    {
        sys.at(i).mass     = 1.0;
        sys.at(i).position = distribute_particle(mt, L);
    }
    sys.topology().construct_molecules();

    mjolnir::UnlimitedGridCellList<traits_type, parameter_type> vlist(margin);

    using neighbor_type = typename decltype(vlist)::neighbor_type;

    BOOST_TEST(!vlist.valid());

    vlist.initialize(sys, pot);
    vlist.make(sys, pot);
    BOOST_TEST(vlist.valid());

    for(std::size_t i=0; i<N; ++i)
    {
        const auto partners = vlist.partners(i);

        // if particle i is not related to the potential, it should not have
        // any interacting partners.
        if(pot.participants().end() == std::find(
            pot.participants().begin(), pot.participants().end(), i))
        {
            BOOST_TEST(partners.size() == 0u);
            continue;
        }

        for(std::size_t j=i+1; j<N; ++j)
        {
            if(std::find_if(partners.begin(), partners.end(),
                [=](const neighbor_type& elem) -> bool {return elem.index == j;}
                        ) == partners.end())
            {
                // should be enough distant (>= threshold)
                const auto dist = mjolnir::math::length(sys.adjust_direction(
                            sys.position(j) - sys.position(i)));
                const bool enough_distant = dist >= threshold;

                // or not a participant
                const bool is_participant = pot.participants().end() != std::find(
                    pot.participants().begin(), pot.participants().end(), j);

                const bool is_ok = enough_distant || (!is_participant);
                BOOST_TEST(is_ok);
            }
            else
            {
                // should be a participant
                const bool found = pot.participants().end() != std::find(
                    pot.participants().begin(), pot.participants().end(), j);
                BOOST_TEST(found);

                // should be enough close (< threshold)
                const auto dist = mjolnir::math::length(sys.adjust_direction(
                            sys.position(j) - sys.position(i)));
                BOOST_TEST(dist < threshold);
            }
        }
    }
}

BOOST_AUTO_TEST_CASE(test_CellList_UnlimitedBoundary_partial_2)
{
    mjolnir::LoggerManager::set_default_logger("test_cell_list.log");

    using traits_type     = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using real_type       = typename traits_type::real_type;
    using boundary_type   = typename traits_type::boundary_type;
    using coordinate_type = typename traits_type::coordinate_type;
    using parameter_type  = typename dummy_potential<real_type>::parameter_type;

    constexpr std::size_t N = 1000;
    constexpr double      L = 10.0;
    constexpr double cutoff = 2.0;
    constexpr double margin = 0.25; // threshold = 2.0 * (1+0.25) = 2.5
    constexpr double threshold = cutoff * (1.0 + margin);

    const auto distribute_particle = [](std::mt19937& mt, double l) -> coordinate_type
    {
        return coordinate_type(
            l * std::generate_canonical<real_type, std::numeric_limits<real_type>::digits>(mt),
            l * std::generate_canonical<real_type, std::numeric_limits<real_type>::digits>(mt),
            l * std::generate_canonical<real_type, std::numeric_limits<real_type>::digits>(mt)
        );
    };

    std::vector<std::size_t> participants; participants.reserve(500);
    for(std::size_t i=0; i<500; ++i)
    {
        participants.push_back(i*2);
    }

    dummy_potential<real_type> pot(cutoff, participants);

    mjolnir::System<traits_type> sys(N, boundary_type{});

    std::mt19937 mt(123456789);
    for(std::size_t i=0; i < N; ++i)
    {
        sys.at(i).mass     = 1.0;
        sys.at(i).position = distribute_particle(mt, L);
    }
    sys.topology().construct_molecules();

    mjolnir::UnlimitedGridCellList<traits_type, parameter_type> vlist(margin);

    using neighbor_type = typename decltype(vlist)::neighbor_type;

    BOOST_TEST(!vlist.valid());

    vlist.initialize(sys, pot);
    vlist.make(sys, pot);
    BOOST_TEST(vlist.valid());

    for(std::size_t i=0; i<N; ++i)
    {
        const auto partners = vlist.partners(i);

        // if particle i is not related to the potential, it should not have
        // any interacting partners.
        if(pot.participants().end() == std::find(
            pot.participants().begin(), pot.participants().end(), i))
        {
            BOOST_TEST(partners.size() == 0u);
            continue;
        }

        for(std::size_t j=i+1; j<N; ++j)
        {
            if(std::find_if(partners.begin(), partners.end(),
                [=](const neighbor_type& elem) -> bool {return elem.index == j;}
                        ) == partners.end())
            {
                // should be enough distant (>= threshold)
                const auto dist = mjolnir::math::length(sys.adjust_direction(
                            sys.position(j) - sys.position(i)));
                const bool enough_distant = dist >= threshold;

                // or not a participant
                const bool is_participant = pot.participants().end() != std::find(
                    pot.participants().begin(), pot.participants().end(), j);

                const bool is_ok = enough_distant || (!is_participant);
                BOOST_TEST(is_ok);
            }
            else
            {
                // should be a participant
                const bool found = pot.participants().end() != std::find(
                    pot.participants().begin(), pot.participants().end(), j);
                BOOST_TEST(found);

                // should be enough close (< threshold)
                const auto dist = mjolnir::math::length(sys.adjust_direction(
                            sys.position(j) - sys.position(i)));
                BOOST_TEST(dist < threshold);
            }
        }
    }
}

<|endoftext|>


<commit_before><commit_msg>fixed Uint conversion<commit_after><|endoftext|>


<commit_before>/** @file
    @brief Implementation

    @date 2015

    @author
    Sensics, Inc.
    <http://sensics.com/osvr>
*/

// Copyright 2015 Sensics, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Internal Includes
#include "SBDBlobExtractor.h"

// Library/third-party includes
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/features2d/features2d.hpp>

// Standard includes
// - none

#include <iostream>

namespace osvr {
namespace vbtracker {

    /// This class used to be the "keypoint enhancer" - it now is used to
    /// after-the-fact extract additional data per keypoint.
    class KeypointDetailer {
      public:
        using ContourType = std::vector<cv::Point2i>;
        std::vector<LedMeasurement>
        augmentKeypoints(cv::Mat const &grayImage,
                         std::vector<cv::KeyPoint> const &foundKeyPoints) {

            std::vector<LedMeasurement> ret;
            cv::Mat greyCopy = grayImage.clone();

            /// Reset the flood fill mask to just have a one-pixel border on the
            /// edges.
            m_floodFillMask =
                cv::Mat::zeros(grayImage.rows + 2, grayImage.cols + 2, CV_8UC1);
            m_floodFillMask.row(0) = 1;
            m_floodFillMask.row(m_floodFillMask.rows - 1) = 1;
            m_floodFillMask.col(0) = 1;
            m_floodFillMask.col(m_floodFillMask.cols - 1) = 1;
            m_origBounds = cv::Rect(1, 1, m_floodFillMask.cols - 2,
                                    m_floodFillMask.rows - 2);
            for (auto const &keypoint : foundKeyPoints) {
                ret.push_back(augmentKeypoint(greyCopy, keypoint));
            }

            return ret;
        }

        cv::Mat getDebugImage() { return m_floodFillMask(m_origBounds); }

      private:
        LedMeasurement augmentKeypoint(cv::Mat grayImage,
                                       cv::KeyPoint origKeypoint) {
            // Saving this now before we monkey with m_floodFillMask
            cv::bitwise_not(m_floodFillMask, m_scratchNotMask);

            const int loDiff = 2;
            const int upDiff = 2;
            auto m_area = cv::floodFill(
                grayImage, m_floodFillMask, origKeypoint.pt, 255,
                &m_filledBounds, loDiff, upDiff,
                CV_FLOODFILL_MASK_ONLY | (/* connectivity 4 or 8 */ 4) |
                    (/* value to write in to mask */ 255 << 8));
            // Now m_floodFillMask contains the mask with both our point
            // and all other points so far. We need to split them by ANDing with
            // the NOT of the old flood-fill mask we saved earlier.
            cv::bitwise_and(m_scratchNotMask, m_floodFillMask,
                            m_perPointResults);
            // OK, now we have the results for just this point in per-point
            // results

            /// Initialize return value
            auto ret = static_cast<LedMeasurement>(origKeypoint);
            if (m_area <= 0) {
                // strange...
                return ret;
            }
            m_binarySubImage = m_perPointResults(m_filledBounds);
            computeContour();
            m_moments = haveContour() ? cv::moments(m_contour, false)
                                      : cv::moments(m_binarySubImage, true);

// the first two we are using from the keypoint description, the latter we can't
// always get (contour-based)
#if 0
            /// Estimate diameter
            ret.diameter = diameter();
            ret.area = area();
            if (haveContour()) {
                ret.circularity = circularity();
            }
#endif
            ret.knowBoundingBox = true;
            ret.boundingBox =
                cv::Size2f(m_filledBounds.width, m_filledBounds.height);
            return ret;
        }

        /// Gets the contour of the binary sub image.
        void computeContour() {

            m_haveContour = false;
            return;
            // Nudge out one pixel each direction if we can.
            auto expandedBounds = m_filledBounds;
            if (expandedBounds.x > 0) {
                expandedBounds.x--;
                expandedBounds.width++;
            }
            if (expandedBounds.y > 0) {
                expandedBounds.y--;
                expandedBounds.height++;
            }
            if (expandedBounds.br().x < m_perPointResults.cols) {
                expandedBounds.width++;
            }
            if (expandedBounds.br().y < m_perPointResults.rows) {
                expandedBounds.height++;
            }

            cv::Mat input = m_perPointResults(expandedBounds).clone();
            std::vector<ContourType> contours;
            cv::findContours(input, contours, CV_RETR_EXTERNAL,
                             CV_CHAIN_APPROX_NONE,
                             cv::Point(-1, -1) + expandedBounds.tl());
#if 0
            if (contours.size() != 1) {
                std::cout << "Weird, we have " << contours.size()
                          << " contours!" << std::endl;
            }
#endif
            if (!contours.empty()) {
                m_contour = cv::Mat(contours[0], true);
                m_haveContour = true;
            } else {
                m_haveContour = false;
            }
        }

        bool haveContour() { return m_haveContour; }

        double area() const { return static_cast<double>(m_area); }

        /// Approximation of a diameter based on assumption of circularity.
        double diameter() const { return 2 * std::sqrt(area() / CV_PI); }

        /// Gets perimeter - requires contour
        double perimeter() const { return cv::arcLength(m_contour, true); }

        /// As used by OpenCV, return value in [0, 1] - requires contour
        double circularity() const {
            auto perim = perimeter();
            return 4 * CV_PI * area() / (perim * perim);
        }
        cv::Rect m_origBounds;
        cv::Mat m_scratchNotMask;
        cv::Mat m_floodFillMask;
        /// @name Per-keypoint data
        /// @{
        int m_area;
        cv::Mat m_perPointResults;
        cv::Mat m_binarySubImage;
        cv::Rect m_filledBounds;
        bool m_haveContour = false;
        cv::Mat m_contour;
        cv::Moments m_moments;
        /// @}
    };

    SBDBlobExtractor::SBDBlobExtractor(ConfigParams const &params)
        : m_params(params), m_keypointDetailer(new KeypointDetailer) {
        auto &p = m_params.blobParams;
        /// Set up blob params
        m_sbdParams.minDistBetweenBlobs = p.minDistBetweenBlobs;

        m_sbdParams.minArea = p.minArea; // How small can the blobs be?

        // Look for bright blobs: there is a bug in this code
        m_sbdParams.filterByColor = false;
        // m_sbdParams.blobColor = static_cast<uchar>(255);

        m_sbdParams.filterByInertia =
            false; // Do we test for non-elongated blobs?
                   // m_sbdParams.minInertiaRatio = 0.5;
                   // m_sbdParams.maxInertiaRatio = 1.0;

        m_sbdParams.filterByCircularity =
            p.filterByCircularity; // Test for circularity?
        m_sbdParams.minCircularity =
            p.minCircularity; // default is 0.8, but the edge of the
                              // case can make the blobs "weird-shaped"

        m_sbdParams.filterByConvexity =
            p.filterByConvexity; // Test for convexity?
        m_sbdParams.minConvexity = p.minConvexity;
    }

    SBDBlobExtractor::~SBDBlobExtractor() {
        /// Needed here where KeypointDetailer is defined.
    }

    std::vector<LedMeasurement> const &
    SBDBlobExtractor::extractBlobs(cv::Mat const &grayImage) {
        m_latestMeasurements.clear();
        m_lastGrayImage = grayImage.clone();
        m_debugThresholdImageDirty = true;
        m_debugBlobImageDirty = true;

        getKeypoints(grayImage);

        /// We'll apply the threshold here first, instead of trusting the flood
        /// fill to repeatably find the same points.
        cv::Mat thresholded;
        cv::threshold(grayImage, thresholded, m_sbdParams.minThreshold, 255,
                      cv::THRESH_BINARY);

        m_latestMeasurements =
            m_keypointDetailer->augmentKeypoints(thresholded, m_keyPoints);

        return m_latestMeasurements;
    }

    void SBDBlobExtractor::getKeypoints(cv::Mat const &grayImage) {
        m_keyPoints.clear();
        //================================================================
        // Tracking the points

        // Construct a blob detector and find the blobs in the image.
        double minVal, maxVal;
        cv::minMaxIdx(grayImage, &minVal, &maxVal);
        auto &p = m_params.blobParams;
        if (maxVal < p.absoluteMinThreshold) {
            /// empty image, early out!
            return;
        }

        auto imageRangeLerp = [=](double alpha) {
            return minVal + (maxVal - minVal) * alpha;
        };
        // 0.3 LERP between min and max as the min threshold, but
        // don't let really dim frames confuse us.
        m_sbdParams.minThreshold = std::max(imageRangeLerp(p.minThresholdAlpha),
                                            p.absoluteMinThreshold);
        m_sbdParams.maxThreshold =
            std::max(imageRangeLerp(0.8), p.absoluteMinThreshold);
        m_sbdParams.thresholdStep =
            (m_sbdParams.maxThreshold - m_sbdParams.minThreshold) /
            p.thresholdSteps;
/// @todo: Make a different set of parameters optimized for the
/// Oculus Dk2.
/// @todo: Determine the maximum size of a trackable blob by seeing
/// when we're so close that we can't view at least four in the
/// camera.
#if CV_MAJOR_VERSION == 2
        cv::Ptr<cv::SimpleBlobDetector> detector =
            new cv::SimpleBlobDetector(m_sbdParams);
#elif CV_MAJOR_VERSION == 3
        auto detector = cv::SimpleBlobDetector::create(m_sbdParams);
#else
#error "Unrecognized OpenCV version!"
#endif
        detector->detect(grayImage, m_keyPoints);

        // @todo: Consider computing the center of mass of a dilated bounding
        // rectangle around each keypoint to produce a more precise subpixel
        // localization of each LED.  The moments() function may be helpful
        // with this.

        // @todo: Estimate the summed brightness of each blob so that we can
        // detect when they are getting brighter and dimmer.  Pass this as
        // the brightness parameter to the Led class when adding a new one
        // or augmenting with a new frame.
    }

    cv::Mat SBDBlobExtractor::generateDebugThresholdImage() const {

        // Fake the thresholded image to give an idea of what the
        // blob detector is doing.
        auto getCurrentThresh = [&](int i) {
            return i * m_sbdParams.thresholdStep + m_sbdParams.minThreshold;
        };
        cv::Mat ret;
        cv::Mat temp;
        cv::threshold(m_lastGrayImage, ret, m_sbdParams.minThreshold, 255,
                      CV_THRESH_BINARY);
        cv::Mat tempOut;
        for (int i = 1; getCurrentThresh(i) < m_sbdParams.maxThreshold; ++i) {
            auto currentThresh = getCurrentThresh(i);
            cv::threshold(m_lastGrayImage, temp, currentThresh, currentThresh,
                          CV_THRESH_BINARY);
            cv::addWeighted(ret, 0.5, temp, 0.5, 0, tempOut);
            ret = tempOut;
        }
        return ret;
    }
    cv::Mat const &SBDBlobExtractor::getDebugThresholdImage() {
        if (m_debugThresholdImageDirty) {
            m_debugThresholdImage = generateDebugThresholdImage();
            m_debugThresholdImageDirty = false;
        }
        return m_debugThresholdImage;
    }

    cv::Mat SBDBlobExtractor::generateDebugBlobImage() const {
        cv::Mat ret;
        cv::Mat tempColor;
        cv::cvtColor(m_lastGrayImage, tempColor, CV_GRAY2BGR);
        // Draw detected blobs as blue circles.
        cv::drawKeypoints(tempColor, m_keyPoints, ret, cv::Scalar(255, 0, 0),
                          cv::DrawMatchesFlags::DRAW_RICH_KEYPOINTS);

        return ret;
    }
    cv::Mat const &SBDBlobExtractor::getDebugBlobImage() {
        if (m_debugBlobImageDirty) {
            m_debugBlobImage = generateDebugBlobImage();
            m_debugBlobImageDirty = false;
        }
        return m_debugBlobImage;
    }

    cv::Mat const &SBDBlobExtractor::getDebugExtraImage() {
        m_extraImage = m_keypointDetailer->getDebugImage().clone();
        return m_extraImage;
    }

} // namespace vbtracker
} // namespace osvr
<commit_msg>Modify the blob extractor to skip the extra postprocessing for now, since we're not using bounding boxes.<commit_after>/** @file
    @brief Implementation

    @date 2015

    @author
    Sensics, Inc.
    <http://sensics.com/osvr>
*/

// Copyright 2015 Sensics, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Internal Includes
#include "SBDBlobExtractor.h"

// Library/third-party includes
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/features2d/features2d.hpp>

// Standard includes
// - none

#include <iostream>

namespace osvr {
namespace vbtracker {

    /// This class used to be the "keypoint enhancer" - it now is used to
    /// after-the-fact extract additional data per keypoint.
    class KeypointDetailer {
      public:
        using ContourType = std::vector<cv::Point2i>;
        std::vector<LedMeasurement>
        augmentKeypoints(cv::Mat const &grayImage,
                         std::vector<cv::KeyPoint> const &foundKeyPoints) {

            std::vector<LedMeasurement> ret;
            cv::Mat greyCopy = grayImage.clone();

            /// Reset the flood fill mask to just have a one-pixel border on the
            /// edges.
            m_floodFillMask =
                cv::Mat::zeros(grayImage.rows + 2, grayImage.cols + 2, CV_8UC1);
            m_floodFillMask.row(0) = 1;
            m_floodFillMask.row(m_floodFillMask.rows - 1) = 1;
            m_floodFillMask.col(0) = 1;
            m_floodFillMask.col(m_floodFillMask.cols - 1) = 1;
            m_origBounds = cv::Rect(1, 1, m_floodFillMask.cols - 2,
                                    m_floodFillMask.rows - 2);
            for (auto const &keypoint : foundKeyPoints) {
                ret.push_back(augmentKeypoint(greyCopy, keypoint));
            }

            return ret;
        }

        cv::Mat getDebugImage() { return m_floodFillMask(m_origBounds); }

      private:
        LedMeasurement augmentKeypoint(cv::Mat grayImage,
                                       cv::KeyPoint origKeypoint) {
            // Saving this now before we monkey with m_floodFillMask
            cv::bitwise_not(m_floodFillMask, m_scratchNotMask);

            const int loDiff = 2;
            const int upDiff = 2;
            auto m_area = cv::floodFill(
                grayImage, m_floodFillMask, origKeypoint.pt, 255,
                &m_filledBounds, loDiff, upDiff,
                CV_FLOODFILL_MASK_ONLY | (/* connectivity 4 or 8 */ 4) |
                    (/* value to write in to mask */ 255 << 8));
            // Now m_floodFillMask contains the mask with both our point
            // and all other points so far. We need to split them by ANDing with
            // the NOT of the old flood-fill mask we saved earlier.
            cv::bitwise_and(m_scratchNotMask, m_floodFillMask,
                            m_perPointResults);
            // OK, now we have the results for just this point in per-point
            // results

            /// Initialize return value
            auto ret = static_cast<LedMeasurement>(origKeypoint);
            if (m_area <= 0) {
                // strange...
                return ret;
            }
            m_binarySubImage = m_perPointResults(m_filledBounds);
            computeContour();
            m_moments = haveContour() ? cv::moments(m_contour, false)
                                      : cv::moments(m_binarySubImage, true);

// the first two we are using from the keypoint description, the latter we can't
// always get (contour-based)
#if 0
            /// Estimate diameter
            ret.diameter = diameter();
            ret.area = area();
            if (haveContour()) {
                ret.circularity = circularity();
            }
#endif
            ret.knowBoundingBox = true;
            ret.boundingBox =
                cv::Size2f(m_filledBounds.width, m_filledBounds.height);
            return ret;
        }

        /// Gets the contour of the binary sub image.
        void computeContour() {

            m_haveContour = false;
            return;
            // Nudge out one pixel each direction if we can.
            auto expandedBounds = m_filledBounds;
            if (expandedBounds.x > 0) {
                expandedBounds.x--;
                expandedBounds.width++;
            }
            if (expandedBounds.y > 0) {
                expandedBounds.y--;
                expandedBounds.height++;
            }
            if (expandedBounds.br().x < m_perPointResults.cols) {
                expandedBounds.width++;
            }
            if (expandedBounds.br().y < m_perPointResults.rows) {
                expandedBounds.height++;
            }

            cv::Mat input = m_perPointResults(expandedBounds).clone();
            std::vector<ContourType> contours;
            cv::findContours(input, contours, CV_RETR_EXTERNAL,
                             CV_CHAIN_APPROX_NONE,
                             cv::Point(-1, -1) + expandedBounds.tl());
#if 0
            if (contours.size() != 1) {
                std::cout << "Weird, we have " << contours.size()
                          << " contours!" << std::endl;
            }
#endif
            if (!contours.empty()) {
                m_contour = cv::Mat(contours[0], true);
                m_haveContour = true;
            } else {
                m_haveContour = false;
            }
        }

        bool haveContour() { return m_haveContour; }

        double area() const { return static_cast<double>(m_area); }

        /// Approximation of a diameter based on assumption of circularity.
        double diameter() const { return 2 * std::sqrt(area() / CV_PI); }

        /// Gets perimeter - requires contour
        double perimeter() const { return cv::arcLength(m_contour, true); }

        /// As used by OpenCV, return value in [0, 1] - requires contour
        double circularity() const {
            auto perim = perimeter();
            return 4 * CV_PI * area() / (perim * perim);
        }
        cv::Rect m_origBounds;
        cv::Mat m_scratchNotMask;
        cv::Mat m_floodFillMask;
        /// @name Per-keypoint data
        /// @{
        int m_area;
        cv::Mat m_perPointResults;
        cv::Mat m_binarySubImage;
        cv::Rect m_filledBounds;
        bool m_haveContour = false;
        cv::Mat m_contour;
        cv::Moments m_moments;
        /// @}
    };

    SBDBlobExtractor::SBDBlobExtractor(ConfigParams const &params)
        : m_params(params), m_keypointDetailer(new KeypointDetailer) {
        auto &p = m_params.blobParams;
        /// Set up blob params
        m_sbdParams.minDistBetweenBlobs = p.minDistBetweenBlobs;

        m_sbdParams.minArea = p.minArea; // How small can the blobs be?

        // Look for bright blobs: there is a bug in this code
        m_sbdParams.filterByColor = false;
        // m_sbdParams.blobColor = static_cast<uchar>(255);

        m_sbdParams.filterByInertia =
            false; // Do we test for non-elongated blobs?
                   // m_sbdParams.minInertiaRatio = 0.5;
                   // m_sbdParams.maxInertiaRatio = 1.0;

        m_sbdParams.filterByCircularity =
            p.filterByCircularity; // Test for circularity?
        m_sbdParams.minCircularity =
            p.minCircularity; // default is 0.8, but the edge of the
                              // case can make the blobs "weird-shaped"

        m_sbdParams.filterByConvexity =
            p.filterByConvexity; // Test for convexity?
        m_sbdParams.minConvexity = p.minConvexity;
    }

    SBDBlobExtractor::~SBDBlobExtractor() {
        /// Needed here where KeypointDetailer is defined.
    }

    std::vector<LedMeasurement> const &
    SBDBlobExtractor::extractBlobs(cv::Mat const &grayImage) {
        m_latestMeasurements.clear();
        m_lastGrayImage = grayImage.clone();
        m_debugThresholdImageDirty = true;
        m_debugBlobImageDirty = true;

        getKeypoints(grayImage);

#if 0
        // This code uses the Keypoint Detailer, but is slightly unreliable.

        /// We'll apply the threshold here first, instead of trusting the flood
        /// fill to repeatably find the same points.
        cv::Mat thresholded;
        cv::threshold(grayImage, thresholded, m_sbdParams.minThreshold, 255,
                      cv::THRESH_BINARY);

        m_latestMeasurements =
            m_keypointDetailer->augmentKeypoints(thresholded, m_keyPoints);
#endif

        /// Use the LedMeasurement constructor to do the conversion from
        /// keypoint to measurement right now.
        m_latestMeasurements.resize(m_keyPoints.size());
        std::transform(
            begin(m_keyPoints), end(m_keyPoints), begin(m_latestMeasurements),
            [](cv::KeyPoint const &kp) { return LedMeasurement{kp}; });

        return m_latestMeasurements;
    }

    void SBDBlobExtractor::getKeypoints(cv::Mat const &grayImage) {
        m_keyPoints.clear();
        //================================================================
        // Tracking the points

        // Construct a blob detector and find the blobs in the image.
        double minVal, maxVal;
        cv::minMaxIdx(grayImage, &minVal, &maxVal);
        auto &p = m_params.blobParams;
        if (maxVal < p.absoluteMinThreshold) {
            /// empty image, early out!
            return;
        }

        auto imageRangeLerp = [=](double alpha) {
            return minVal + (maxVal - minVal) * alpha;
        };
        // 0.3 LERP between min and max as the min threshold, but
        // don't let really dim frames confuse us.
        m_sbdParams.minThreshold = std::max(imageRangeLerp(p.minThresholdAlpha),
                                            p.absoluteMinThreshold);
        m_sbdParams.maxThreshold =
            std::max(imageRangeLerp(0.8), p.absoluteMinThreshold);
        m_sbdParams.thresholdStep =
            (m_sbdParams.maxThreshold - m_sbdParams.minThreshold) /
            p.thresholdSteps;
/// @todo: Make a different set of parameters optimized for the
/// Oculus Dk2.
/// @todo: Determine the maximum size of a trackable blob by seeing
/// when we're so close that we can't view at least four in the
/// camera.
#if CV_MAJOR_VERSION == 2
        cv::Ptr<cv::SimpleBlobDetector> detector =
            new cv::SimpleBlobDetector(m_sbdParams);
#elif CV_MAJOR_VERSION == 3
        auto detector = cv::SimpleBlobDetector::create(m_sbdParams);
#else
#error "Unrecognized OpenCV version!"
#endif
        detector->detect(grayImage, m_keyPoints);

        // @todo: Consider computing the center of mass of a dilated bounding
        // rectangle around each keypoint to produce a more precise subpixel
        // localization of each LED.  The moments() function may be helpful
        // with this.

        // @todo: Estimate the summed brightness of each blob so that we can
        // detect when they are getting brighter and dimmer.  Pass this as
        // the brightness parameter to the Led class when adding a new one
        // or augmenting with a new frame.
    }

    cv::Mat SBDBlobExtractor::generateDebugThresholdImage() const {

        // Fake the thresholded image to give an idea of what the
        // blob detector is doing.
        auto getCurrentThresh = [&](int i) {
            return i * m_sbdParams.thresholdStep + m_sbdParams.minThreshold;
        };
        cv::Mat ret;
        cv::Mat temp;
        cv::threshold(m_lastGrayImage, ret, m_sbdParams.minThreshold, 255,
                      CV_THRESH_BINARY);
        cv::Mat tempOut;
        for (int i = 1; getCurrentThresh(i) < m_sbdParams.maxThreshold; ++i) {
            auto currentThresh = getCurrentThresh(i);
            cv::threshold(m_lastGrayImage, temp, currentThresh, currentThresh,
                          CV_THRESH_BINARY);
            cv::addWeighted(ret, 0.5, temp, 0.5, 0, tempOut);
            ret = tempOut;
        }
        return ret;
    }
    cv::Mat const &SBDBlobExtractor::getDebugThresholdImage() {
        if (m_debugThresholdImageDirty) {
            m_debugThresholdImage = generateDebugThresholdImage();
            m_debugThresholdImageDirty = false;
        }
        return m_debugThresholdImage;
    }

    cv::Mat SBDBlobExtractor::generateDebugBlobImage() const {
        cv::Mat ret;
        cv::Mat tempColor;
        cv::cvtColor(m_lastGrayImage, tempColor, CV_GRAY2BGR);
        // Draw detected blobs as blue circles.
        cv::drawKeypoints(tempColor, m_keyPoints, ret, cv::Scalar(255, 0, 0),
                          cv::DrawMatchesFlags::DRAW_RICH_KEYPOINTS);

        return ret;
    }
    cv::Mat const &SBDBlobExtractor::getDebugBlobImage() {
        if (m_debugBlobImageDirty) {
            m_debugBlobImage = generateDebugBlobImage();
            m_debugBlobImageDirty = false;
        }
        return m_debugBlobImage;
    }

    cv::Mat const &SBDBlobExtractor::getDebugExtraImage() {
        m_extraImage = m_keypointDetailer->getDebugImage().clone();
        return m_extraImage;
    }

} // namespace vbtracker
} // namespace osvr
<|endoftext|>