<commit_before>/**
 * @file parser_test.cpp
 * @author Chase Geigle
 */

#include <iostream>

#include "cpptoml.h"
#include "logging/logger.h"
#include "parser/io/ptb_reader.h"
#include "parser/sr_parser.h"
#include "parser/trees/visitors/annotation_remover.h"
#include "parser/trees/visitors/empty_remover.h"
#include "parser/trees/visitors/unary_chain_remover.h"
#include "parser/trees/visitors/multi_transformer.h"
#include "parser/trees/visitors/head_finder.h"
#include "parser/trees/visitors/binarizer.h"
#include "parser/trees/visitors/debinarizer.h"
#include "parser/trees/visitors/transition_finder.h"
#include "parser/trees/visitors/leaf_node_finder.h"
#include "parser/trees/visitors/sequence_extractor.h"
#include "util/filesystem.h"
#include "util/progress.h"

using namespace meta;

std::string two_digit(uint8_t num)
{
    std::stringstream ss;
    ss << std::setw(2) << std::setfill('0') << static_cast<int>(num);
    return ss.str();
}

int main(int argc, char** argv)
{

    if (argc < 3)
    {
        std::cerr << "Usage: " << argv[0] << " config.toml output_file" << std::endl;
        return 1;
    }

    logging::set_cerr_logging();

    auto config = cpptoml::parse_file(argv[1]);

    auto prefix = config.get_as<std::string>("prefix");
    if (!prefix)
    {
        LOG(fatal) << "Global configuration must have a prefix key" << ENDLG;
        return 1;
    }

    auto parser_grp = config.get_group("parser");
    if (!parser_grp)
    {
        LOG(fatal) << "Configuration must contain a [parser] group" << ENDLG;
        return 1;
    }

    auto parser_prefix = parser_grp->get_as<std::string>("prefix");
    if (!parser_prefix)
    {
        LOG(fatal)
            << "[parser] group must contain a prefix to store model files"
            << ENDLG;
        return 1;
    }

    auto treebank = parser_grp->get_as<std::string>("treebank");
    if (!treebank)
    {
        LOG(fatal) << "[parser] group must contain a treebank path" << ENDLG;
        return 1;
    }

    auto corpus = parser_grp->get_as<std::string>("corpus");
    if (!corpus)
    {
        LOG(fatal) << "[parser] group must contain a corpus" << ENDLG;
        return 1;
    }

    auto test_sections = parser_grp->get_array("test-sections");
    if (!test_sections)
    {
        LOG(fatal) << "[parser] group must contain test-sections" << ENDLG;
        return 1;
    }

    auto section_size = parser_grp->get_as<int64_t>("section-size");
    if (!section_size)
    {
        LOG(fatal) << "[parser] group must contain section-size" << ENDLG;
        return 1;
    }

    std::string path = *prefix + "/" + *treebank + "/treebank-3/parsed/mrg/"
                       + *corpus;

    parser::multi_transformer<parser::annotation_remover, parser::empty_remover,
                              parser::unary_chain_remover> transformer;

    std::vector<sequence::sequence> testing;
    {
        auto begin = test_sections->at(0)->as<int64_t>()->value();
        auto end = test_sections->at(1)->as<int64_t>()->value();
        printing::progress progress(" > Reading testing data: ",
                                    (end - begin + 1) * *section_size);
        for (uint8_t i = begin; i <= end; ++i)
        {
            auto folder = two_digit(i);
            for (uint8_t j = 0; j <= *section_size; ++j)
            {
                progress((i - begin) * 99 + j);
                auto file = *corpus + "_" + folder + two_digit(j) + ".mrg";
                auto filename = path + "/" + folder + "/" + file;
                auto trees = parser::io::extract_trees(filename);
                for (auto& tree : trees)
                {
                    parser::sequence_extractor seq_ex;
                    tree.transform(transformer);
                    tree.visit(seq_ex);
                    testing.emplace_back(seq_ex.sequence());
                }
            }
        }
    }
    LOG(info) << testing.size() << " test examples" << ENDLG;

    parser::sr_parser parser{*parser_prefix};

    std::ofstream output{argv[2]};
    for (const auto& seq : testing)
    {
        std::cout << "Parsing:";
        for (const auto& obs : seq)
            std::cout << " " <<  obs.symbol();
        std::cout << "\n";

        auto tree = parser.parse(seq);
        tree.pretty_print(std::cout);
        std::cout << "\n\n";
        output << tree << "\n";
    }

    return 0;
}
<commit_msg>Simplify the parser test application.<commit_after>/**
 * @file parser_test.cpp
 * @author Chase Geigle
 */

#include <iostream>

#include "cpptoml.h"
#include "logging/logger.h"
#include "parser/io/ptb_reader.h"
#include "parser/sr_parser.h"
#include "parser/trees/visitors/empty_remover.h"
#include "parser/trees/visitors/sequence_extractor.h"
#include "util/filesystem.h"
#include "util/progress.h"

using namespace meta;

std::string two_digit(uint8_t num)
{
    std::stringstream ss;
    ss << std::setw(2) << std::setfill('0') << static_cast<int>(num);
    return ss.str();
}

int main(int argc, char** argv)
{

    if (argc < 3)
    {
        std::cerr << "Usage: " << argv[0] << " config.toml output_file" << std::endl;
        return 1;
    }

    logging::set_cerr_logging();

    auto config = cpptoml::parse_file(argv[1]);

    auto prefix = config.get_as<std::string>("prefix");
    if (!prefix)
    {
        LOG(fatal) << "Global configuration must have a prefix key" << ENDLG;
        return 1;
    }

    auto parser_grp = config.get_group("parser");
    if (!parser_grp)
    {
        LOG(fatal) << "Configuration must contain a [parser] group" << ENDLG;
        return 1;
    }

    auto parser_prefix = parser_grp->get_as<std::string>("prefix");
    if (!parser_prefix)
    {
        LOG(fatal)
            << "[parser] group must contain a prefix to store model files"
            << ENDLG;
        return 1;
    }

    auto treebank = parser_grp->get_as<std::string>("treebank");
    if (!treebank)
    {
        LOG(fatal) << "[parser] group must contain a treebank path" << ENDLG;
        return 1;
    }

    auto corpus = parser_grp->get_as<std::string>("corpus");
    if (!corpus)
    {
        LOG(fatal) << "[parser] group must contain a corpus" << ENDLG;
        return 1;
    }

    auto test_sections = parser_grp->get_array("test-sections");
    if (!test_sections)
    {
        LOG(fatal) << "[parser] group must contain test-sections" << ENDLG;
        return 1;
    }

    auto section_size = parser_grp->get_as<int64_t>("section-size");
    if (!section_size)
    {
        LOG(fatal) << "[parser] group must contain section-size" << ENDLG;
        return 1;
    }

    std::string path = *prefix + "/" + *treebank + "/treebank-3/parsed/mrg/"
                       + *corpus;

    parser::empty_remover transformer;

    std::vector<sequence::sequence> testing;
    {
        auto begin = test_sections->at(0)->as<int64_t>()->value();
        auto end = test_sections->at(1)->as<int64_t>()->value();
        printing::progress progress(" > Reading testing data: ",
                                    (end - begin + 1) * *section_size);
        for (uint8_t i = begin; i <= end; ++i)
        {
            auto folder = two_digit(i);
            for (uint8_t j = 0; j <= *section_size; ++j)
            {
                progress((i - begin) * 99 + j);
                auto file = *corpus + "_" + folder + two_digit(j) + ".mrg";
                auto filename = path + "/" + folder + "/" + file;
                auto trees = parser::io::extract_trees(filename);
                for (auto& tree : trees)
                {
                    parser::sequence_extractor seq_ex;
                    tree.transform(transformer);
                    tree.visit(seq_ex);
                    testing.emplace_back(seq_ex.sequence());
                }
            }
        }
    }
    LOG(info) << testing.size() << " test examples" << ENDLG;

    parser::sr_parser parser{*parser_prefix};

    std::ofstream output{argv[2]};
    for (const auto& seq : testing)
    {
        std::cout << "Parsing:";
        for (const auto& obs : seq)
            std::cout << " " <<  obs.symbol();
        std::cout << "\n";

        auto tree = parser.parse(seq);
        tree.pretty_print(std::cout);
        std::cout << "\n\n";
        output << tree << "\n";
    }

    return 0;
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2012 Carsten Burstedde, Donna Calhoun
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

 * Redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.
 * Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include "amr_forestclaw.H"
#include "no_solver_user.H"
#include "amr_utils.H"

void no_solver_linker(fclaw2d_domain_t* domain)
{
    link_problem_setup(domain,no_solver_setprob);

    /* Initialize data but don't do anything */

    fclaw2d_solver_functions_t* sf = get_solver_functions(domain);
    sf->f_patch_initialize = &no_solver_patch_initialize;
    sf->f_patch_single_step_update = &no_solver_update;

    fclaw2d_output_functions* of = get_output_functions(domain);
    of->f_patch_write_header = &vtk_write_header;
    of->f_patch_write_output = &vtk_write_output;

    link_regrid_functions(domain,no_solver_patch_tag4refinement,
                          no_solver_patch_tag4coarsening);
}

void no_solver_setprob(fclaw2d_domain_t* domain)
{
    set_maptype_();
}

void no_solver_patch_initialize(fclaw2d_domain_t *domain,
                                fclaw2d_patch_t *this_patch,
                                int this_block_idx,
                                int this_patch_idx)
{

    set_block_(&this_block_idx);

    // Global parameters
    const amr_options_t *gparms = get_domain_parms(domain);
    int mx = gparms->mx;
    int my = gparms->my;
    int mbc = gparms->mbc;
    int meqn = gparms->meqn;

    // Parameters specific to this patch
    ClawPatch *cp = get_clawpatch(this_patch);
    double xlower = cp->xlower();
    double ylower = cp->ylower();
    double dx = cp->dx();
    double dy = cp->dy();

    // Pointers needed for the fortran
    double* q = cp->q();

    int mpirank = domain->mpirank;

    initialize_(mx,my,meqn,mbc,xlower,ylower,dx,dy,q,mpirank);
}

double no_solver_update(fclaw2d_domain_t *domain,
                        fclaw2d_patch_t *this_patch,
                        int this_block_idx,
                        int this_patch_idx,
                        double t,
                        double dt)
{
    ClawPatch *cp = get_clawpatch(this_patch);

    // save the current time step for time interpolation.  Otherwise, we get
    // unitialized values.
    cp->save_current_step();  // Save for time interpolation
    return 1.0;
}

/* -----------------------------------------------------------------
   Default routine for tagging patches for refinement and coarsening
   ----------------------------------------------------------------- */
fclaw_bool no_solver_patch_tag4refinement(fclaw2d_domain_t *domain,
                                          fclaw2d_patch_t *this_patch,
                                          int this_block_idx, int this_patch_idx,
                                          int initflag)
{
    /* ----------------------------------------------------------- */
    // Global parameters
    const amr_options_t *gparms = get_domain_parms(domain);
    int mx = gparms->mx;
    int my = gparms->my;
    int mbc = gparms->mbc;
    int meqn = gparms->meqn;

    /* ----------------------------------------------------------- */
    // Patch specific parameters
    ClawPatch *cp = get_clawpatch(this_patch);
    double xlower = cp->xlower();
    double ylower = cp->ylower();
    double dx = cp->dx();
    double dy = cp->dy();

    /* ------------------------------------------------------------ */
    // Pointers needed to pass to Fortran
    double* q = cp->q();

    int tag_patch;  // == 0 or 1
    no_solver_tag4refinement_(mx,my,mbc,meqn,xlower,ylower,dx,dy,q,initflag,tag_patch);
    return tag_patch == 1;
}

fclaw_bool no_solver_patch_tag4coarsening(fclaw2d_domain_t *domain,
                                          fclaw2d_patch_t *this_patch,
                                          int blockno,
                                          int patchno)
{
    // This might come in handy if we want to debug a coarsening routine without
    // worrying about solvers.

    /* ----------------------------------------------------------- */
    // Global parameters
    const amr_options_t *gparms = get_domain_parms(domain);
    int mx = gparms->mx;
    int my = gparms->my;
    int mbc = gparms->mbc;
    int meqn = gparms->meqn;

    /* ----------------------------------------------------------- */
    // Patch specific parameters
    ClawPatch *cp = get_clawpatch(this_patch);
    double xlower = cp->xlower();
    double ylower = cp->ylower();
    double dx = cp->dx();
    double dy = cp->dy();

    /* ------------------------------------------------------------ */
    // Pointers needed to pass to Fortran
    double* qcoarse = cp->q();

    int tag_patch;  // == 0 or 1
    no_solver_tag4coarsening_(mx,my,mbc,meqn,xlower,ylower,dx,dy,qcoarse,tag_patch);
    return tag_patch == 0;
}

void vtk_write_header(fclaw2d_domain_t* domain, int iframe, int ngrids)
{
    const amr_options_t *gparms = get_domain_parms(domain);
    double time = get_domain_time(domain);

    printf("Matlab output Frame %d  at time %16.8e\n\n",iframe,time);

    // Write out header file containing global information for 'iframe'
    int meqn = gparms->meqn;
    int maux = 0;
    write_tfile_(iframe,time,meqn,ngrids,maux);

    // This opens file 'fort.qXXXX' for replace (where XXXX = <zero padding><iframe>, e.g. 0001,
    // 0010, 0114), and closes the file.
    new_qfile_(iframe);
}

void vtk_write_output(fclaw2d_domain_t *domain, fclaw2d_patch_t *this_patch,
                      int this_block_idx, int this_patch_idx,
                      int iframe,int num,int level)
{
    // In case this is needed by the setaux routine
    set_block_(&this_block_idx);

    /* ----------------------------------------------------------- */
    // Global parameters
    const amr_options_t *gparms = get_domain_parms(domain);
    int mx = gparms->mx;
    int my = gparms->my;
    int mbc = gparms->mbc;
    int meqn = gparms->meqn;

    /* ----------------------------------------------------------- */
    // Patch specific parameters
    ClawPatch *cp = get_clawpatch(this_patch);
    double xlower = cp->xlower();
    double ylower = cp->ylower();
    double dx = cp->dx();
    double dy = cp->dy();

    /* ------------------------------------------------------------ */
    // Pointers needed to pass to Fortran
    double* q = cp->q();

    // Other input arguments
    int maxmx = mx;
    int maxmy = my;

    /* ------------------------------------------------------------- */
    // This opens a file for append.  Now, the style is in the 'clawout' style.
    int matlab_level = level + 1;
    write_qfile_(maxmx,maxmy,meqn,mbc,mx,my,xlower,ylower,dx,dy,q,
                 iframe,num,matlab_level,this_block_idx);
}
<commit_msg>Changed name of output routine<commit_after>/*
Copyright (c) 2012 Carsten Burstedde, Donna Calhoun
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

 * Redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.
 * Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include "amr_forestclaw.H"
#include "no_solver_user.H"
#include "amr_utils.H"

void no_solver_linker(fclaw2d_domain_t* domain)
{
    link_problem_setup(domain,no_solver_setprob);

    /* Initialize data but don't do anything */

    fclaw2d_solver_functions_t* sf = get_solver_functions(domain);
    sf->f_patch_initialize = &no_solver_patch_initialize;
    sf->f_patch_single_step_update = &no_solver_update;

    fclaw2d_output_functions* of = get_output_functions(domain);
    of->f_patch_write_header = &matlab_parallel_write_header;
    of->f_patch_write_output = &matlab_parallel_write_output;

    link_regrid_functions(domain,no_solver_patch_tag4refinement,
                          no_solver_patch_tag4coarsening);
}

void no_solver_setprob(fclaw2d_domain_t* domain)
{
    set_maptype_();
}

void no_solver_patch_initialize(fclaw2d_domain_t *domain,
                                fclaw2d_patch_t *this_patch,
                                int this_block_idx,
                                int this_patch_idx)
{

    set_block_(&this_block_idx);

    // Global parameters
    const amr_options_t *gparms = get_domain_parms(domain);
    int mx = gparms->mx;
    int my = gparms->my;
    int mbc = gparms->mbc;
    int meqn = gparms->meqn;

    // Parameters specific to this patch
    ClawPatch *cp = get_clawpatch(this_patch);
    double xlower = cp->xlower();
    double ylower = cp->ylower();
    double dx = cp->dx();
    double dy = cp->dy();

    // Pointers needed for the fortran
    double* q = cp->q();

    int mpirank = domain->mpirank;

    initialize_(mx,my,meqn,mbc,xlower,ylower,dx,dy,q,mpirank);
}

double no_solver_update(fclaw2d_domain_t *domain,
                        fclaw2d_patch_t *this_patch,
                        int this_block_idx,
                        int this_patch_idx,
                        double t,
                        double dt)
{
    ClawPatch *cp = get_clawpatch(this_patch);

    // save the current time step for time interpolation.  Otherwise, we get
    // unitialized values.
    cp->save_current_step();  // Save for time interpolation
    return 1.0;
}

/* -----------------------------------------------------------------
   Default routine for tagging patches for refinement and coarsening
   ----------------------------------------------------------------- */
fclaw_bool no_solver_patch_tag4refinement(fclaw2d_domain_t *domain,
                                          fclaw2d_patch_t *this_patch,
                                          int this_block_idx, int this_patch_idx,
                                          int initflag)
{
    /* ----------------------------------------------------------- */
    // Global parameters
    const amr_options_t *gparms = get_domain_parms(domain);
    int mx = gparms->mx;
    int my = gparms->my;
    int mbc = gparms->mbc;
    int meqn = gparms->meqn;

    /* ----------------------------------------------------------- */
    // Patch specific parameters
    ClawPatch *cp = get_clawpatch(this_patch);
    double xlower = cp->xlower();
    double ylower = cp->ylower();
    double dx = cp->dx();
    double dy = cp->dy();

    /* ------------------------------------------------------------ */
    // Pointers needed to pass to Fortran
    double* q = cp->q();

    int tag_patch;  // == 0 or 1
    no_solver_tag4refinement_(mx,my,mbc,meqn,xlower,ylower,dx,dy,q,initflag,tag_patch);
    return tag_patch == 1;
}

fclaw_bool no_solver_patch_tag4coarsening(fclaw2d_domain_t *domain,
                                          fclaw2d_patch_t *this_patch,
                                          int blockno,
                                          int patchno)
{
    // This might come in handy if we want to debug a coarsening routine without
    // worrying about solvers.

    /* ----------------------------------------------------------- */
    // Global parameters
    const amr_options_t *gparms = get_domain_parms(domain);
    int mx = gparms->mx;
    int my = gparms->my;
    int mbc = gparms->mbc;
    int meqn = gparms->meqn;

    /* ----------------------------------------------------------- */
    // Patch specific parameters
    ClawPatch *cp = get_clawpatch(this_patch);
    double xlower = cp->xlower();
    double ylower = cp->ylower();
    double dx = cp->dx();
    double dy = cp->dy();

    /* ------------------------------------------------------------ */
    // Pointers needed to pass to Fortran
    double* qcoarse = cp->q();

    int tag_patch;  // == 0 or 1
    no_solver_tag4coarsening_(mx,my,mbc,meqn,xlower,ylower,dx,dy,qcoarse,tag_patch);
    return tag_patch == 0;
}

void matlab_parallel_write_header(fclaw2d_domain_t* domain, int iframe, int ngrids)
{
    if (domain->mpirank == 0)
    {
        const amr_options_t *gparms = get_domain_parms(domain);
        double time = get_domain_time(domain);

        printf("Matlab output Frame %d  at time %16.8e\n\n",iframe,time);

        // Write out header file containing global information for 'iframe'
        int meqn = gparms->meqn;
        int maux = 0;
        write_tfile_(iframe,time,meqn,ngrids,maux);

        // This opens file 'fort.qXXXX' for replace (where XXXX = <zero padding><iframe>, e.g. 0001,
        // 0010, 0114), and closes the file.
        new_qfile_(iframe);
    }
}

void matlab_parallel_write_output(fclaw2d_domain_t *domain, fclaw2d_patch_t *this_patch,
                                  int this_block_idx, int this_patch_idx,
                                  int iframe,int num,int level)
{
    // In case this is needed by the setaux routine
    set_block_(&this_block_idx);

    /* ----------------------------------------------------------- */
    // Global parameters
    const amr_options_t *gparms = get_domain_parms(domain);
    int mx = gparms->mx;
    int my = gparms->my;
    int mbc = gparms->mbc;
    int meqn = gparms->meqn;

    /* ----------------------------------------------------------- */
    // Patch specific parameters
    ClawPatch *cp = get_clawpatch(this_patch);
    double xlower = cp->xlower();
    double ylower = cp->ylower();
    double dx = cp->dx();
    double dy = cp->dy();

    /* ------------------------------------------------------------ */
    // Pointers needed to pass to Fortran
    double* q = cp->q();

    // Other input arguments
    int maxmx = mx;
    int maxmy = my;

    /* ------------------------------------------------------------- */
    // This opens a file for append.  Now, the style is in the 'clawout' style.
    int matlab_level = level + 1;

    write_qfile_(maxmx,maxmy,meqn,mbc,mx,my,xlower,ylower,dx,dy,q,
                 iframe,num,matlab_level,this_block_idx);
}
<|endoftext|>


<commit_before>#include "pegr/script/ScriptHelper.hpp"

#include <stdexcept>
#include <cassert>

#include "pegr/debug/DebugMacros.hpp"
#include "pegr/script/Script.hpp"

namespace pegr {
namespace Script {
namespace Helper {

void for_pairs(int table_idx, std::function<bool()> func, bool pops_value) {
    lua_State* l = Script::get_lua_state();
    table_idx = Script::absolute_idx(table_idx);
    lua_pushnil(l);
    while (lua_next(l, table_idx) != 0) {
        Script::Pop_Guard pg;
        if (!pops_value) {
            pg.on_push(1); // We will eat the value ourselves
        }
        // Assume we will need to eat the key ourselves
        // in case the function raises errors or breaks the loop
        pg.on_push(1);
        bool cont = func();
        if (!cont) {
            break;
        }
        // If we got to this point, the key no longer needs to be eaten by us
        pg.on_pop(1);
    }
}

/**
 * @brief Recursive helper for simple_deep_copy
 * @param l The Lua state to use
 * @param table_idx Absolute index for the table to be copied
 * @param dupe_idx Absolute index for a table that maps originals to copies
 */
void simple_deep_copy_helper(lua_State* l, int table_idx, int dupe_idx) {
    assert_balance(1);

    // Don't copy non-tables
    if (!lua_istable(l, table_idx)) {
        lua_pushvalue(l, table_idx);
        return;
    }

    // Check if we already made a copy before doing anything
    lua_pushvalue(l, table_idx);
    lua_rawget(l, dupe_idx);
    if (lua_istable(l, -1)) {
        return;
    }
    lua_pop(l, 1);

    lua_newtable(l); // TODO: preallocate enough space
    int copy_idx = lua_gettop(l);

    // Store this table as a copy of the table at table_idx, even though
    // we aren't finished making it yet. (If we try to wait until we
    // are finished, it's possible that we will be waiting forever.)
    lua_pushvalue(l, table_idx);
    lua_pushvalue(l, copy_idx);
    lua_rawset(l, dupe_idx);

    lua_pushnil(l);
    while (lua_next(l, table_idx) != 0) {
        int val_idx = lua_gettop(l);

        // Copy the reference to the key so there is one for iterating with
        lua_pushvalue(l, -2);

        simple_deep_copy_helper(l, val_idx, dupe_idx);

        // Set the key-value pair
        lua_settable(l, copy_idx);

        // Remove the original value reference
        lua_pop(l, 1);

        // (A reference to the key is still on the stack for iteration)
    }
}

void simple_deep_copy(int table_idx) {
    assert_balance(1);
    lua_State* l = Script::get_lua_state();
    table_idx = Script::absolute_idx(table_idx);

    /* This table maps original tables to the copied forms.
     * This allows for duplicate reference values to be preserved and
     * avoids crashing for recursive cases.
     */
    lua_newtable(l);
    int dupe_idx = lua_gettop(l);

    simple_deep_copy_helper(l, table_idx, dupe_idx);

    lua_remove(l, dupe_idx);
}

void run_simple_function(Script::Regref ref, int nresults) {
    assert_balance(nresults);
    Script::push_reference(ref);
    Script::run_function(0, nresults);
}

std::string to_string(int idx, const char* def, int max_recusions) {
    assert_balance(0);
    lua_State* l = Script::get_lua_state();
    idx = Script::absolute_idx(idx);

    Script::push_reference(Script::m_luaglob_tostring);

    lua_pushvalue(l, idx);
    for (int count = 0; count <= max_recusions; ++count) {
        std::size_t strsize;
        const char* strdata = lua_tolstring(l, -1, &strsize);

        if (!strdata) {
            lua_pushvalue(l, -2); // tostring() function
            lua_pushvalue(l, -2); // the pesky string
            Script::run_function(1, 1);
            lua_remove(l, -2); // pop off the old pesky string
            continue;
        }

        lua_pop(l, 2); // remove cached tostring and the value
        return std::string(strdata, strsize);
    }
    lua_pop(l, 2); // remove cached tostring and the value
    if (def) {
        return def;
    }
    throw std::runtime_error(
        "Maximum recusion depth reached when trying to convert a Lua "
        "value into a string through tostring()");
}

const char* GENERIC_TO_STRING_DEFAULT =
        "[MAX TOSTRING() RECURSION DEPTH EXCEEDED]";

} // namespace Helper
} // namespace Script
} // namespace pegr
<commit_msg>Case-sensitive function name<commit_after>#include "pegr/script/ScriptHelper.hpp"

#include <stdexcept>
#include <cassert>

#include "pegr/debug/DebugMacros.hpp"
#include "pegr/script/Script.hpp"

namespace pegr {
namespace Script {
namespace Helper {

void for_pairs(int table_idx, std::function<bool()> func, bool pops_value) {
    lua_State* l = Script::get_lua_state();
    table_idx = Script::absolute_idx(table_idx);
    lua_pushnil(l);
    while (lua_next(l, table_idx) != 0) {
        Script::Pop_Guard pg;
        if (!pops_value) {
            pg.on_push(1); // We will eat the value ourselves
        }
        // Assume we will need to eat the key ourselves
        // in case the function raises errors or breaks the loop
        pg.on_push(1);
        bool cont = func();
        if (!cont) {
            break;
        }
        // If we got to this point, the key no longer needs to be eaten by us
        pg.on_pop(1);
    }
}

/**
 * @brief Recursive helper for simple_deep_copy
 * @param l The Lua state to use
 * @param table_idx Absolute index for the table to be copied
 * @param dupe_idx Absolute index for a table that maps originals to copies
 */
void simple_deep_copy_helper(lua_State* l, int table_idx, int dupe_idx) {
    assert_balance(1);

    // Don't copy non-tables
    if (!lua_istable(l, table_idx)) {
        lua_pushvalue(l, table_idx);
        return;
    }

    // Check if we already made a copy before doing anything
    lua_pushvalue(l, table_idx);
    lua_rawget(l, dupe_idx);
    if (lua_istable(l, -1)) {
        return;
    }
    lua_pop(l, 1);

    lua_newtable(l); // TODO: preallocate enough space
    int copy_idx = lua_gettop(l);

    // Store this table as a copy of the table at table_idx, even though
    // we aren't finished making it yet. (If we try to wait until we
    // are finished, it's possible that we will be waiting forever.)
    lua_pushvalue(l, table_idx);
    lua_pushvalue(l, copy_idx);
    lua_rawset(l, dupe_idx);

    lua_pushnil(l);
    while (lua_next(l, table_idx) != 0) {
        int val_idx = lua_gettop(l);

        // Copy the reference to the key so there is one for iterating with
        lua_pushvalue(l, -2);

        simple_deep_copy_helper(l, val_idx, dupe_idx);

        // Set the key-value pair
        lua_settable(l, copy_idx);

        // Remove the original value reference
        lua_pop(l, 1);

        // (A reference to the key is still on the stack for iteration)
    }
}

void simple_deep_copy(int table_idx) {
    assert_balance(1);
    lua_State* l = Script::get_lua_state();
    table_idx = Script::absolute_idx(table_idx);

    /* This table maps original tables to the copied forms.
     * This allows for duplicate reference values to be preserved and
     * avoids crashing for recursive cases.
     */
    lua_newtable(l);
    int dupe_idx = lua_gettop(l);

    simple_deep_copy_helper(l, table_idx, dupe_idx);

    lua_remove(l, dupe_idx);
}

void run_simple_function(Script::Regref ref, int nresults) {
    assert_balance(nresults);
    Script::push_reference(ref);
    Script::run_function(0, nresults);
}

std::string to_string(int idx, const char* def, int max_recusions) {
    assert_balance(0);
    lua_State* l = Script::get_lua_state();
    idx = Script::absolute_idx(idx);

    Script::push_reference(Script::m_luaglob_tostring);

    lua_pushvalue(l, idx);
    for (int count = 0; count <= max_recusions; ++count) {
        std::size_t strsize;
        const char* strdata = lua_tolstring(l, -1, &strsize);

        if (!strdata) {
            lua_pushvalue(l, -2); // tostring() function
            lua_pushvalue(l, -2); // the pesky string
            Script::run_function(1, 1);
            lua_remove(l, -2); // pop off the old pesky string
            continue;
        }

        lua_pop(l, 2); // remove cached tostring and the value
        return std::string(strdata, strsize);
    }
    lua_pop(l, 2); // remove cached tostring and the value
    if (def) {
        return def;
    }
    throw std::runtime_error(
        "Maximum recusion depth reached when trying to convert a Lua "
        "value into a string through tostring()");
}

const char* GENERIC_TO_STRING_DEFAULT =
        "[MAX tostring() RECURSION DEPTH EXCEEDED]";

} // namespace Helper
} // namespace Script
} // namespace pegr
<|endoftext|>


<commit_before>
# include "CHBD.hpp"

// -------------------------------------------------------------------------
// Constructor:
// -------------------------------------------------------------------------

CHBD::CHBD(const CommonParams& pin,
           const GetPot& input_params) : p(pin), c1(p), c2(p), cp(p),
                                         particles(p,input_params),
                                         k1(p), k2(p), k4(p)
{

    //	---------------------------------------
    // create the FFTw plans:
    //	---------------------------------------

    nxyz = p.nx*p.ny*p.nz;
    dummy = fftw_alloc_complex(nxyz);
    p_forward  = fftw_mpi_plan_dft_3d(p.NX,p.NY,p.NZ,dummy,dummy,
                 MPI_COMM_WORLD,FFTW_FORWARD,FFTW_MEASURE);
    p_backward = fftw_mpi_plan_dft_3d(p.NX,p.NY,p.NZ,dummy,dummy,
                 MPI_COMM_WORLD,FFTW_BACKWARD,FFTW_MEASURE);
    fftw_free(dummy);

    //	---------------------------------------
    // set needed parameters:
    //	---------------------------------------

    co = input_params("PFApp/co",0.5);
    w = input_params("PFApp/w",1.0);
    M = input_params("PFApp/M",1.0);
    kap = input_params("PFApp/kap",1.0);
    part_step_skip = input_params("PFApp/part_step_skip",1);
    eCH = input_params("PFApp/eCH",0.1);

}

// -------------------------------------------------------------------------
// Destructor:
// -------------------------------------------------------------------------

CHBD::~CHBD()
{

}

// -------------------------------------------------------------------------
// Initialize phase-field method:
// -------------------------------------------------------------------------

void CHBD::initPhaseField()
{

    //	---------------------------------------
    // initialize the particle suspension:
    //	---------------------------------------

    particles.initParticles();
    cp = particles.mapToGrid();

    //	---------------------------------------
    // initialize the concentration fields:
    //	---------------------------------------

    srand(time(NULL)*(p.rank+1));   // set the random seed
    for (int i=0; i<nxyz; i++) {
        double cpi = creal(cp.getValue(i));
        double r = (double)rand()/RAND_MAX;
        double val = co + 0.1*(r-0.5);
        val *= 1.0 - cpi;
        c1.setValue(i,val);
        c2.setValue(i,1-val);
    }

    //	---------------------------------------
    // initialize the fourier wave-vectors:
    //	---------------------------------------

    calculateKfields();

    //	---------------------------------------
    // Output the initial configuration:
    //	---------------------------------------

    current_step = 0;
    outputPhaseField();

    //	---------------------------------------
    // Sync the processors:
    //	---------------------------------------

    MPI::COMM_WORLD.Barrier();

}

// -------------------------------------------------------------------------
// Step forward in time the phase-field method:
// -------------------------------------------------------------------------

void CHBD::updatePhaseField()
{

    //	---------------------------------------
    // Update particles:
    //	---------------------------------------

    particles.setTimeStep(current_step);
    if (current_step%part_step_skip == 0) {
        particles.calcCapillaryForce(cp,c1,c2);
        if (p.rank == 0) particles.updateParticles();
        cp = particles.mapToGrid();
    }

    //	---------------------------------------
    // Update Cahn-Hilliard:
    //	---------------------------------------

    Sfield dfdc1 = 12.0*w*(c1*c1*c1 - c1*c1 + c1*c2*c2 + c1*cp);
    Sfield dfdc2 = 12.0*w*(c2*c2*c2 - c2*c2 + c2*c1*c1 + c2*cp);
    c1.fft(p_forward);
    c2.fft(p_forward);
    dfdc1.fft(p_forward);
    dfdc2.fft(p_forward);
    if (eCH > 0.0)
    {
        // update with electric field in z-direction
        c1 = (c1 - p.dt*k2*dfdc1 - p.dt*eCH*(*kz)*(*kz)*c1)\
             /(1.0 + kap*p.dt*k4);
        c2 = (c2 - p.dt*k2*dfdc2 - p.dt*eCH*(*kz)*(*kz)*c2)\
             /(1.0 + kap*p.dt*k4);
    }
    else
    {
        // update with no electric field
        c1 = (c1 - p.dt*k2*dfdc1)/(1.0 + kap*p.dt*k4);
        c2 = (c2 - p.dt*k2*dfdc2)/(1.0 + kap*p.dt*k4);
    }
    c1.ifft(p_backward);
    c2.ifft(p_backward);

    //	---------------------------------------
    // Sync the processors:
    //	---------------------------------------

    MPI::COMM_WORLD.Barrier();

}

// -------------------------------------------------------------------------
// Write output for the phase-field method:
// -------------------------------------------------------------------------

void CHBD::outputPhaseField()
{
    int iskip = p.iskip;
    int jskip = p.jskip;
    int kskip = p.kskip;
    Sfield phi = c1 - c2;
    phi.writeVTKFile("c",current_step,iskip,jskip,kskip);
    if (p.rank == 0) particles.outputParticles();
}

// -------------------------------------------------------------------------
// Calculate the wave-vector fields (k^1, k^2, k^4 for spectral solution):
// -------------------------------------------------------------------------

void CHBD::calculateKfields()
{
    Sfield kxf(p);
    Sfield kyf(p);
    Sfield kzf(p);
    for (int i=0; i<p.nx; i++) {
        for (int j=0; j<p.ny; j++) {
            for (int k=0; k<p.nz; k++) {
                int ii = i + p.xOff;
                int jj = j;
                int kk = k;
                int ndx = i*p.NY*p.NZ + j*p.NZ + k;
                double kx = ii*(2.0*M_PI/p.LX);
                double ky = jj*(2.0*M_PI/p.LY);
                double kz = kk*(2.0*M_PI/p.LZ);
                if (ii > p.NX/2) kx = (ii-p.NX)*(2.0*M_PI/p.LX);
                if (jj > p.NY/2) ky = (jj-p.NY)*(2.0*M_PI/p.LY);
                if (kk > p.NZ/2) kz = (kk-p.NZ)*(2.0*M_PI/p.LZ);
                double kx2 = kx*kx;
                double ky2 = ky*ky;
                double kz2 = kz*kz;
                double kijk2 = kx2 + ky2 + kz2;
                double kijk4 = kijk2*kijk2;
                kxf.setValue(ndx,kx);
                kyf.setValue(ndx,ky);
                kzf.setValue(ndx,kz);
                k2.setValue(ndx,kijk2);
                k4.setValue(ndx,kijk4);
            }
        }
    }
    k1.setXValues(kxf);
    k1.setYValues(kyf);
    k1.setZValues(kzf);
    kz = k1.getZValues();
}
<commit_msg>change to outputting all fields for CHBD model<commit_after>
# include "CHBD.hpp"

// -------------------------------------------------------------------------
// Constructor:
// -------------------------------------------------------------------------

CHBD::CHBD(const CommonParams& pin,
           const GetPot& input_params) : p(pin), c1(p), c2(p), cp(p),
                                         particles(p,input_params),
                                         k1(p), k2(p), k4(p)
{

    //	---------------------------------------
    // create the FFTw plans:
    //	---------------------------------------

    nxyz = p.nx*p.ny*p.nz;
    dummy = fftw_alloc_complex(nxyz);
    p_forward  = fftw_mpi_plan_dft_3d(p.NX,p.NY,p.NZ,dummy,dummy,
                 MPI_COMM_WORLD,FFTW_FORWARD,FFTW_MEASURE);
    p_backward = fftw_mpi_plan_dft_3d(p.NX,p.NY,p.NZ,dummy,dummy,
                 MPI_COMM_WORLD,FFTW_BACKWARD,FFTW_MEASURE);
    fftw_free(dummy);

    //	---------------------------------------
    // set needed parameters:
    //	---------------------------------------

    co = input_params("PFApp/co",0.5);
    w = input_params("PFApp/w",1.0);
    M = input_params("PFApp/M",1.0);
    kap = input_params("PFApp/kap",1.0);
    part_step_skip = input_params("PFApp/part_step_skip",1);
    eCH = input_params("PFApp/eCH",0.1);

}

// -------------------------------------------------------------------------
// Destructor:
// -------------------------------------------------------------------------

CHBD::~CHBD()
{

}

// -------------------------------------------------------------------------
// Initialize phase-field method:
// -------------------------------------------------------------------------

void CHBD::initPhaseField()
{

    //	---------------------------------------
    // initialize the particle suspension:
    //	---------------------------------------

    particles.initParticles();
    cp = particles.mapToGrid();

    //	---------------------------------------
    // initialize the concentration fields:
    //	---------------------------------------

    srand(time(NULL)*(p.rank+1));   // set the random seed
    for (int i=0; i<nxyz; i++) {
        double cpi = creal(cp.getValue(i));
        double r = (double)rand()/RAND_MAX;
        double val = co + 0.1*(r-0.5);
        val *= 1.0 - cpi;
        c1.setValue(i,val);
        c2.setValue(i,1-val);
    }

    //	---------------------------------------
    // initialize the fourier wave-vectors:
    //	---------------------------------------

    calculateKfields();

    //	---------------------------------------
    // Output the initial configuration:
    //	---------------------------------------

    current_step = 0;
    outputPhaseField();

    //	---------------------------------------
    // Sync the processors:
    //	---------------------------------------

    MPI::COMM_WORLD.Barrier();

}

// -------------------------------------------------------------------------
// Step forward in time the phase-field method:
// -------------------------------------------------------------------------

void CHBD::updatePhaseField()
{

    //	---------------------------------------
    // Update particles:
    //	---------------------------------------

    particles.setTimeStep(current_step);
    if (current_step%part_step_skip == 0) {
        particles.calcCapillaryForce(cp,c1,c2);
        if (p.rank == 0) particles.updateParticles();
        cp = particles.mapToGrid();
    }

    //	---------------------------------------
    // Update Cahn-Hilliard:
    //	---------------------------------------

    Sfield dfdc1 = 12.0*w*(c1*c1*c1 - c1*c1 + c1*c2*c2 + c1*cp);
    Sfield dfdc2 = 12.0*w*(c2*c2*c2 - c2*c2 + c2*c1*c1 + c2*cp);
    c1.fft(p_forward);
    c2.fft(p_forward);
    dfdc1.fft(p_forward);
    dfdc2.fft(p_forward);
    if (eCH > 0.0)
    {
        // update with electric field in z-direction
        c1 = (c1 - p.dt*k2*dfdc1 - p.dt*eCH*(*kz)*(*kz)*c1)\
             /(1.0 + kap*p.dt*k4);
        c2 = (c2 - p.dt*k2*dfdc2 - p.dt*eCH*(*kz)*(*kz)*c2)\
             /(1.0 + kap*p.dt*k4);
    }
    else
    {
        // update with no electric field
        c1 = (c1 - p.dt*k2*dfdc1)/(1.0 + kap*p.dt*k4);
        c2 = (c2 - p.dt*k2*dfdc2)/(1.0 + kap*p.dt*k4);
    }
    c1.ifft(p_backward);
    c2.ifft(p_backward);

    //	---------------------------------------
    // Sync the processors:
    //	---------------------------------------

    MPI::COMM_WORLD.Barrier();

}

// -------------------------------------------------------------------------
// Write output for the phase-field method:
// -------------------------------------------------------------------------

void CHBD::outputPhaseField()
{
    int iskip = p.iskip;
    int jskip = p.jskip;
    int kskip = p.kskip;
    c1.writeVTKFile("c1",current_step,iskip,jskip,kskip);
    c2.writeVTKFile("c2",current_step,iskip,jskip,kskip);
    cp.writeVTKFile("cp",current_step,iskip,jskip,kskip);
    if (p.rank == 0) particles.outputParticles();
}

// -------------------------------------------------------------------------
// Calculate the wave-vector fields (k^1, k^2, k^4 for spectral solution):
// -------------------------------------------------------------------------

void CHBD::calculateKfields()
{
    Sfield kxf(p);
    Sfield kyf(p);
    Sfield kzf(p);
    for (int i=0; i<p.nx; i++) {
        for (int j=0; j<p.ny; j++) {
            for (int k=0; k<p.nz; k++) {
                int ii = i + p.xOff;
                int jj = j;
                int kk = k;
                int ndx = i*p.NY*p.NZ + j*p.NZ + k;
                double kx = ii*(2.0*M_PI/p.LX);
                double ky = jj*(2.0*M_PI/p.LY);
                double kz = kk*(2.0*M_PI/p.LZ);
                if (ii > p.NX/2) kx = (ii-p.NX)*(2.0*M_PI/p.LX);
                if (jj > p.NY/2) ky = (jj-p.NY)*(2.0*M_PI/p.LY);
                if (kk > p.NZ/2) kz = (kk-p.NZ)*(2.0*M_PI/p.LZ);
                double kx2 = kx*kx;
                double ky2 = ky*ky;
                double kz2 = kz*kz;
                double kijk2 = kx2 + ky2 + kz2;
                double kijk4 = kijk2*kijk2;
                kxf.setValue(ndx,kx);
                kyf.setValue(ndx,ky);
                kzf.setValue(ndx,kz);
                k2.setValue(ndx,kijk2);
                k4.setValue(ndx,kijk4);
            }
        }
    }
    k1.setXValues(kxf);
    k1.setYValues(kyf);
    k1.setZValues(kzf);
    kz = k1.getZValues();
}
<|endoftext|>


<commit_before>/*
 * Software License Agreement (BSD License)
 *
 *  Copyright (c) 2012, Willow Garage, Inc.
 *  All rights reserved.
 *
 *  Redistribution and use in source and binary forms, with or without
 *  modification, are permitted provided that the following conditions
 *  are met:
 *
 *   * Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above
 *     copyright notice, this list of conditions and the following
 *     disclaimer in the documentation and/or other materials provided
 *     with the distribution.
 *   * Neither the name of Willow Garage, Inc. nor the names of its
 *     contributors may be used to endorse or promote products derived
 *     from this software without specific prior written permission.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 *  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 *  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 *  FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 *  COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 *  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 *  BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 *  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 *  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 *  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 *  ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 *  POSSIBILITY OF SUCH DAMAGE.
 *
 */

#include <pcl/io/openni_grabber.h>
#include <pcl/visualization/cloud_viewer.h>
#include <boost/thread/thread.hpp>
#include <pcl/visualization/pcl_visualizer.h>
#include <pcl/io/io.h>
#include <boost/make_shared.hpp>
#include <pcl/common/time.h>
#include <pcl/features/integral_image_normal.h>
#include <pcl/features/normal_3d.h>
#include <pcl/ModelCoefficients.h>
#include <pcl/segmentation/organized_multi_plane_segmentation.h>
#include <pcl/filters/extract_indices.h>

typedef pcl::PointXYZRGBA PointT;

class OpenNIOrganizedMultiPlaneSegmentation
{
  private:
    boost::shared_ptr<pcl::visualization::PCLVisualizer> viewer;
    pcl::PointCloud<PointT>::ConstPtr prev_cloud;
    boost::mutex cloud_mutex;

  public:
    OpenNIOrganizedMultiPlaneSegmentation ()
    {

    }
    ~OpenNIOrganizedMultiPlaneSegmentation ()
    {
    }

    boost::shared_ptr<pcl::visualization::PCLVisualizer>
    cloudViewer (pcl::PointCloud<PointT>::ConstPtr cloud)
    {
      boost::shared_ptr < pcl::visualization::PCLVisualizer> viewer (new pcl::visualization::PCLVisualizer ("Viewer"));
      viewer->setBackgroundColor (0, 0, 0);
      pcl::visualization::PointCloudColorHandlerCustom<PointT> single_color (cloud, 0, 255, 0);
      viewer->addPointCloud < PointT > (cloud, single_color, "cloud");
      viewer->setPointCloudRenderingProperties (pcl::visualization::PCL_VISUALIZER_POINT_SIZE, 3, "cloud");
      viewer->setPointCloudRenderingProperties (pcl::visualization::PCL_VISUALIZER_OPACITY, 0.15, "cloud");
      viewer->addCoordinateSystem (1.0);
      viewer->initCameraParameters ();
      return (viewer);
    }

    void
    cloud_cb_ (const pcl::PointCloud<PointT>::ConstPtr& cloud)
    {
      if (!viewer->wasStopped ())
      {
        cloud_mutex.lock ();
        prev_cloud = cloud;
        cloud_mutex.unlock ();
      }
    }

    void
    run ()
    {
      pcl::Grabber* interface = new pcl::OpenNIGrabber ();

      boost::function<void
      (const pcl::PointCloud<PointT>::ConstPtr&)> f = boost::bind (&OpenNIOrganizedMultiPlaneSegmentation::cloud_cb_, this, _1);

      //make a viewer
      pcl::PointCloud<PointT>::Ptr init_cloud_ptr (new pcl::PointCloud<PointT>);
      viewer = cloudViewer (init_cloud_ptr);
      boost::signals2::connection c = interface->registerCallback (f);

      interface->start ();

      unsigned char red [6] = {255,   0,   0, 255, 255,   0};
      unsigned char grn [6] = {  0, 255,   0, 255,   0, 255};
      unsigned char blu [6] = {  0,   0, 255,   0, 255, 255};

      while (!viewer->wasStopped ())
      {
        viewer->spinOnce (100);

        if (prev_cloud && cloud_mutex.try_lock ())
        {
          printf ("Computing normals...\n");
          pcl::PointCloud < pcl::Normal > normal_cloud;
          double normal_start = pcl::getTime ();
          pcl::IntegralImageNormalEstimation<PointT, pcl::Normal> ne;
          ne.setNormalEstimationMethod (ne.SIMPLE_3D_GRADIENT);
          ne.setMaxDepthChangeFactor (0.02f);
          ne.setNormalSmoothingSize (10.0f);
          ne.setInputCloud (prev_cloud);
          ne.compute (normal_cloud);
          double normal_end = pcl::getTime ();
          std::cout << "Normal Estimation took " << double (normal_end - normal_start) << std::endl;

          printf ("Extracting planes...\n");
          std::vector < pcl::ModelCoefficients > plane_models;
          std::vector < pcl::PointIndices > planes_inliers;
          pcl::PointCloud<pcl::Label> labels;
          std::vector<pcl::PointIndices> label_indices;
          double plane_extract_start = pcl::getTime ();
          pcl::OrganizedMultiPlaneSegmentation<PointT, pcl::Normal, pcl::Label> mps;
          mps.setMinInliers (10000);
          mps.setAngularThreshold (0.017453 * 2.0); //3 degrees
          mps.setDistanceThreshold (0.02); //2cm
          mps.setInputNormals (boost::make_shared < pcl::PointCloud<pcl::Normal> > (normal_cloud));
          mps.setInputCloud (prev_cloud);
          mps.segment (plane_models, planes_inliers,labels,label_indices);
          double plane_extract_end = pcl::getTime ();
          std::cout << "Plane extraction took " << double (plane_extract_end - plane_extract_start) << std::endl;
          std::cout << "Frame took " << double (plane_extract_end - normal_start) << std::endl;

          char name [1024];
          pcl::PointCloud<PointT>::Ptr cluster (new pcl::PointCloud<PointT>);
          viewer->removeAllPointClouds(0);
          viewer->removeAllShapes ();
          if(!viewer->updatePointCloud < PointT > (boost::make_shared < pcl::PointCloud<PointT> > (*prev_cloud),"cloud")){
            viewer->addPointCloud < PointT > (boost::make_shared < pcl::PointCloud<PointT> > (*prev_cloud),"cloud");
          }

          //Draw Visualization
          for (size_t i = 0; i < planes_inliers.size (); i++)
          {
            Eigen::Vector4f centroid;
            pcl::compute3DCentroid (*prev_cloud, planes_inliers[i], centroid);
            pcl::PointXYZ pt1 = pcl::PointXYZ (centroid[0], centroid[1], centroid[2]);
            pcl::PointXYZ pt2 = pcl::PointXYZ (centroid[0] + (0.5 * plane_models[i].values[0]),
                                               centroid[1] + (0.5 * plane_models[i].values[1]),
                                               centroid[2] + (0.5 * plane_models[i].values[2]));
            sprintf (name, "normal_%d", (unsigned)i);
            viewer->addArrow (pt2, pt1, 1.0, 0, 0, name);

            pcl::copyPointCloud (*prev_cloud,planes_inliers[i],*cluster);
            sprintf (name, "plane_%02d", (int)i);
            pcl::visualization::PointCloudColorHandlerCustom <PointT> color (cluster, red[i], grn[i], blu[i]);
            viewer->addPointCloud (cluster, color, name);

          }
          cloud_mutex.unlock ();
        }
      }

      interface->stop ();
    }
};

int
main ()
{
  OpenNIOrganizedMultiPlaneSegmentation multi_plane_app;
  multi_plane_app.run ();
  return 0;
}
<commit_msg>minor changes<commit_after>/*
 * Software License Agreement (BSD License)
 *
 *  Copyright (c) 2012, Willow Garage, Inc.
 *  All rights reserved.
 *
 *  Redistribution and use in source and binary forms, with or without
 *  modification, are permitted provided that the following conditions
 *  are met:
 *
 *   * Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above
 *     copyright notice, this list of conditions and the following
 *     disclaimer in the documentation and/or other materials provided
 *     with the distribution.
 *   * Neither the name of Willow Garage, Inc. nor the names of its
 *     contributors may be used to endorse or promote products derived
 *     from this software without specific prior written permission.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 *  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 *  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 *  FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 *  COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 *  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 *  BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 *  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 *  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 *  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 *  ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 *  POSSIBILITY OF SUCH DAMAGE.
 *
 */

#include <pcl/io/openni_grabber.h>
#include <pcl/visualization/cloud_viewer.h>
#include <boost/thread/thread.hpp>
#include <pcl/visualization/pcl_visualizer.h>
#include <pcl/io/io.h>
#include <boost/make_shared.hpp>
#include <pcl/common/time.h>
#include <pcl/features/integral_image_normal.h>
#include <pcl/features/normal_3d.h>
#include <pcl/ModelCoefficients.h>
#include <pcl/segmentation/organized_multi_plane_segmentation.h>
#include <pcl/filters/extract_indices.h>

typedef pcl::PointXYZRGBA PointT;

class OpenNIOrganizedMultiPlaneSegmentation
{
  private:
    boost::shared_ptr<pcl::visualization::PCLVisualizer> viewer;
    pcl::PointCloud<PointT>::ConstPtr prev_cloud;
    boost::mutex cloud_mutex;

  public:
    OpenNIOrganizedMultiPlaneSegmentation ()
    {

    }
    ~OpenNIOrganizedMultiPlaneSegmentation ()
    {
    }

    boost::shared_ptr<pcl::visualization::PCLVisualizer>
    cloudViewer (pcl::PointCloud<PointT>::ConstPtr cloud)
    {
      boost::shared_ptr < pcl::visualization::PCLVisualizer> viewer (new pcl::visualization::PCLVisualizer ("Viewer"));
      viewer->setBackgroundColor (0, 0, 0);
      pcl::visualization::PointCloudColorHandlerCustom<PointT> single_color (cloud, 0, 255, 0);
      viewer->addPointCloud<PointT> (cloud, single_color, "cloud");
      viewer->setPointCloudRenderingProperties (pcl::visualization::PCL_VISUALIZER_POINT_SIZE, 3, "cloud");
      viewer->setPointCloudRenderingProperties (pcl::visualization::PCL_VISUALIZER_OPACITY, 0.15, "cloud");
      viewer->addCoordinateSystem (1.0);
      viewer->initCameraParameters ();
      return (viewer);
    }

    void
    cloud_cb_ (const pcl::PointCloud<PointT>::ConstPtr& cloud)
    {
      if (!viewer->wasStopped ())
      {
        cloud_mutex.lock ();
        prev_cloud = cloud;
        cloud_mutex.unlock ();
      }
    }

    void
    removePreviousDataFromScreen (size_t prev_models_size)
    {
      char name[1024];
      for (size_t i = 0; i < prev_models_size; i++)
      {
        sprintf (name, "normal_%d", (unsigned)i);
        viewer->removeShape (name);

        sprintf (name, "plane_%02d", (int)i);
        viewer->removePointCloud (name);
      }
    }

    void
    run ()
    {
      pcl::Grabber* interface = new pcl::OpenNIGrabber ();

      boost::function<void(const pcl::PointCloud<PointT>::ConstPtr&)> f = boost::bind (&OpenNIOrganizedMultiPlaneSegmentation::cloud_cb_, this, _1);

      //make a viewer
      pcl::PointCloud<PointT>::Ptr init_cloud_ptr (new pcl::PointCloud<PointT>);
      viewer = cloudViewer (init_cloud_ptr);
      boost::signals2::connection c = interface->registerCallback (f);

      interface->start ();

      unsigned char red [6] = {255,   0,   0, 255, 255,   0};
      unsigned char grn [6] = {  0, 255,   0, 255,   0, 255};
      unsigned char blu [6] = {  0,   0, 255,   0, 255, 255};

      pcl::IntegralImageNormalEstimation<PointT, pcl::Normal> ne;
      ne.setNormalEstimationMethod (ne.SIMPLE_3D_GRADIENT);
      ne.setMaxDepthChangeFactor (0.02f);
      ne.setNormalSmoothingSize (10.0f);

      pcl::OrganizedMultiPlaneSegmentation<PointT, pcl::Normal, pcl::Label> mps;
      mps.setMinInliers (10000);
      mps.setAngularThreshold (0.017453 * 2.0); //3 degrees
      mps.setDistanceThreshold (0.02); //2cm

      std::vector<pcl::ModelCoefficients> plane_models;
      std::vector<pcl::PointIndices> planes_inliers;

      size_t prev_models_size = 0;
      char name[1024];

      while (!viewer->wasStopped ())
      {
        viewer->spinOnce (100);

        if (prev_cloud && cloud_mutex.try_lock ())
        {
          pcl::PointCloud<pcl::Normal>::Ptr normal_cloud (new pcl::PointCloud<pcl::Normal>);
          double normal_start = pcl::getTime ();
          ne.setInputCloud (prev_cloud);
          ne.compute (*normal_cloud);
          double normal_end = pcl::getTime ();
          std::cout << "Normal Estimation took " << double (normal_end - normal_start) << std::endl;

          double plane_extract_start = pcl::getTime ();
          mps.setInputNormals (normal_cloud);
          mps.setInputCloud (prev_cloud);
          mps.segment (plane_models, planes_inliers);
          double plane_extract_end = pcl::getTime ();
          std::cout << "Plane extraction took " << double (plane_extract_end - plane_extract_start) << std::endl;
          std::cout << "Frame took " << double (plane_extract_end - normal_start) << std::endl;

          pcl::PointCloud<PointT>::Ptr cluster (new pcl::PointCloud<PointT>);

          if (!viewer->updatePointCloud<PointT> (prev_cloud, "cloud"))
            viewer->addPointCloud<PointT> (prev_cloud, "cloud");

          removePreviousDataFromScreen (prev_models_size);
          //Draw Visualization
          for (size_t i = 0; i < planes_inliers.size (); i++)
          {
            Eigen::Vector4f centroid;
            pcl::compute3DCentroid (*prev_cloud, planes_inliers[i], centroid);
            pcl::PointXYZ pt1 = pcl::PointXYZ (centroid[0], centroid[1], centroid[2]);
            pcl::PointXYZ pt2 = pcl::PointXYZ (centroid[0] + (0.5 * plane_models[i].values[0]),
                                               centroid[1] + (0.5 * plane_models[i].values[1]),
                                               centroid[2] + (0.5 * plane_models[i].values[2]));
            sprintf (name, "normal_%d", (unsigned)i);
            viewer->addArrow (pt2, pt1, 1.0, 0, 0, name);

            pcl::copyPointCloud (*prev_cloud, planes_inliers[i], *cluster);
            sprintf (name, "plane_%02d", (int)i);
            pcl::visualization::PointCloudColorHandlerCustom <PointT> color (cluster, red[i], grn[i], blu[i]);
            viewer->addPointCloud (cluster, color, name);
          }
          prev_models_size = planes_inliers.size ();
          cloud_mutex.unlock ();
        }
      }

      interface->stop ();
    }
};

int
main ()
{
  OpenNIOrganizedMultiPlaneSegmentation multi_plane_app;
  multi_plane_app.run ();
  return 0;
}
<|endoftext|>


<commit_before>
/*
 * Copyright 2011 Google Inc.
 *
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#define WIN32_LEAN_AND_MEAN
#include <Windows.h>
#include <wincodec.h>
#include "SkAutoCoInitialize.h"
#include "SkImageDecoder.h"
#include "SkImageEncoder.h"
#include "SkIStream.h"
#include "SkMovie.h"
#include "SkStream.h"
#include "SkTScopedComPtr.h"

class SkImageDecoder_WIC : public SkImageDecoder {
protected:
    virtual bool onDecode(SkStream* stream, SkBitmap* bm, Mode mode);
};

bool SkImageDecoder_WIC::onDecode(SkStream* stream, SkBitmap* bm, Mode mode) {
    //Initialize COM.
    SkAutoCoInitialize scopedCo;
    if (!scopedCo.succeeded()) {
        return false;
    }

    HRESULT hr = S_OK;

    //Create Windows Imaging Component ImagingFactory.
    SkTScopedComPtr<IWICImagingFactory> piImagingFactory;
    if (SUCCEEDED(hr)) {
        hr = CoCreateInstance(
            CLSID_WICImagingFactory
            , NULL
            , CLSCTX_INPROC_SERVER
            , IID_PPV_ARGS(&piImagingFactory)
        );
    }

    //Convert SkStream to IStream.
    SkTScopedComPtr<IStream> piStream;
    if (SUCCEEDED(hr)) {
        hr = SkIStream::CreateFromSkStream(stream, false, &piStream);
    }

    //Make sure we're at the beginning of the stream.
    if (SUCCEEDED(hr)) {
        LARGE_INTEGER liBeginning = { 0 };
        hr = piStream->Seek(liBeginning, STREAM_SEEK_SET, NULL);
    }

    //Create the decoder from the stream content.
    SkTScopedComPtr<IWICBitmapDecoder> piBitmapDecoder;
    if (SUCCEEDED(hr)) {
        hr = piImagingFactory->CreateDecoderFromStream(
            piStream.get()                    //Image to be decoded
            , NULL                            //No particular vendor
            , WICDecodeMetadataCacheOnDemand  //Cache metadata when needed
            , &piBitmapDecoder                //Pointer to the decoder
        );
    }

    //Get the first frame from the decoder.
    SkTScopedComPtr<IWICBitmapFrameDecode> piBitmapFrameDecode;
    if (SUCCEEDED(hr)) {
        hr = piBitmapDecoder->GetFrame(0, &piBitmapFrameDecode);
    }

    //Get the BitmapSource interface of the frame.
    SkTScopedComPtr<IWICBitmapSource> piBitmapSourceOriginal;
    if (SUCCEEDED(hr)) {
        hr = piBitmapFrameDecode->QueryInterface(
            IID_PPV_ARGS(&piBitmapSourceOriginal)
        );
    }

    //Get the size of the bitmap.
    UINT width;
    UINT height;
    if (SUCCEEDED(hr)) {
        hr = piBitmapSourceOriginal->GetSize(&width, &height);
    }

    //Exit early if we're only looking for the bitmap bounds.
    if (SUCCEEDED(hr)) {
        bm->setConfig(SkBitmap::kARGB_8888_Config, width, height);
        if (SkImageDecoder::kDecodeBounds_Mode == mode) {
            return true;
        }
        if (!this->allocPixelRef(bm, NULL)) {
            return false;
        }
    }

    //Create a format converter.
    SkTScopedComPtr<IWICFormatConverter> piFormatConverter;
    if (SUCCEEDED(hr)) {
        hr = piImagingFactory->CreateFormatConverter(&piFormatConverter);
    }

    if (SUCCEEDED(hr)) {
        hr = piFormatConverter->Initialize(
            piBitmapSourceOriginal.get()      //Input bitmap to convert
            , GUID_WICPixelFormat32bppPBGRA   //Destination pixel format
            , WICBitmapDitherTypeNone         //Specified dither patterm
            , NULL                            //Specify a particular palette
            , 0.f                             //Alpha threshold
            , WICBitmapPaletteTypeCustom      //Palette translation type
        );
    }

    //Get the BitmapSource interface of the format converter.
    SkTScopedComPtr<IWICBitmapSource> piBitmapSourceConverted;
    if (SUCCEEDED(hr)) {
        hr = piFormatConverter->QueryInterface(
            IID_PPV_ARGS(&piBitmapSourceConverted)
        );
    }

    //Copy the pixels into the bitmap.
    if (SUCCEEDED(hr)) {
        SkAutoLockPixels alp(*bm);
        bm->eraseColor(0);
        const int stride = bm->rowBytes();
        hr = piBitmapSourceConverted->CopyPixels(
            NULL,                             //Get all the pixels
            stride,
            stride * height,
            reinterpret_cast<BYTE *>(bm->getPixels())
        );
    }

    return SUCCEEDED(hr);
}

/////////////////////////////////////////////////////////////////////////

SkImageDecoder* SkImageDecoder::Factory(SkStream* stream) {
    return SkNEW(SkImageDecoder_WIC);
}

/////////////////////////////////////////////////////////////////////////

SkMovie* SkMovie::DecodeStream(SkStream* stream) {
    return NULL;
}

/////////////////////////////////////////////////////////////////////////

class SkImageEncoder_WIC : public SkImageEncoder {
public:
    SkImageEncoder_WIC(Type t) : fType(t) {}

protected:
    virtual bool onEncode(SkWStream* stream, const SkBitmap& bm, int quality);

private:
    Type fType;
};

bool SkImageEncoder_WIC::onEncode(SkWStream* stream
                                , const SkBitmap& bitmapOrig
                                , int quality)
{
    GUID type;
    switch (fType) {
        case kJPEG_Type:
            type = GUID_ContainerFormatJpeg;
            break;
        case kPNG_Type:
            type = GUID_ContainerFormatPng;
            break;
        default:
            return false;
    }

    //Convert to 8888 if needed.
    const SkBitmap* bitmap;
    SkBitmap bitmapCopy;
    if (SkBitmap::kARGB_8888_Config == bitmapOrig.config()) {
        bitmap = &bitmapOrig;
    } else {
        if (!bitmapOrig.copyTo(&bitmapCopy, SkBitmap::kARGB_8888_Config)) {
            return false;
        }
        bitmap = &bitmapCopy;
    }

    //Initialize COM.
    SkAutoCoInitialize scopedCo;
    if (!scopedCo.succeeded()) {
        return false;
    }

    HRESULT hr = S_OK;

    //Create Windows Imaging Component ImagingFactory.
    SkTScopedComPtr<IWICImagingFactory> piImagingFactory;
    if (SUCCEEDED(hr)) {
        hr = CoCreateInstance(
            CLSID_WICImagingFactory
            , NULL
            , CLSCTX_INPROC_SERVER
            , IID_PPV_ARGS(&piImagingFactory)
        );
    }

    //Convert the SkWStream to an IStream.
    SkTScopedComPtr<IStream> piStream;
    if (SUCCEEDED(hr)) {
        hr = SkWIStream::CreateFromSkWStream(stream, &piStream);
    }

    //Create an encode of the appropriate type.
    SkTScopedComPtr<IWICBitmapEncoder> piEncoder;
    if (SUCCEEDED(hr)) {
        hr = piImagingFactory->CreateEncoder(type, NULL, &piEncoder);
    }

    if (SUCCEEDED(hr)) {
        hr = piEncoder->Initialize(piStream.get(), WICBitmapEncoderNoCache);
    }

    //Create a the frame.
    SkTScopedComPtr<IWICBitmapFrameEncode> piBitmapFrameEncode;
    SkTScopedComPtr<IPropertyBag2> piPropertybag;
    if (SUCCEEDED(hr)) {
        hr = piEncoder->CreateNewFrame(&piBitmapFrameEncode, &piPropertybag);
    }

    if (SUCCEEDED(hr)) {
        PROPBAG2 name = { 0 };
        name.dwType = PROPBAG2_TYPE_DATA;
        name.vt = VT_R4;
        name.pstrName = L"ImageQuality";

        VARIANT value;
        VariantInit(&value);
        value.vt = VT_R4;
        value.fltVal = (FLOAT)(quality / 100.0);

        //Ignore result code.
        //  This returns E_FAIL if the named property is not in the bag.
        //TODO(bungeman) enumerate the properties,
        //  write and set hr iff property exists.
        piPropertybag->Write(1, &name, &value);
    }
    if (SUCCEEDED(hr)) {
        hr = piBitmapFrameEncode->Initialize(piPropertybag.get());
    }

    //Set the size of the frame.
    const UINT width = bitmap->width();
    const UINT height = bitmap->height();
    if (SUCCEEDED(hr)) {
        hr = piBitmapFrameEncode->SetSize(width, height);
    }

    //Set the pixel format of the frame.
    const WICPixelFormatGUID formatDesired = GUID_WICPixelFormat32bppBGRA;
    WICPixelFormatGUID formatGUID = formatDesired;
    if (SUCCEEDED(hr)) {
        hr = piBitmapFrameEncode->SetPixelFormat(&formatGUID);
    }
    if (SUCCEEDED(hr)) {
        //Be sure the image format is the one requested.
        hr = IsEqualGUID(formatGUID, formatDesired) ? S_OK : E_FAIL;
    }

    //Write the pixels into the frame.
    if (SUCCEEDED(hr)) {
        SkAutoLockPixels alp(*bitmap);
        hr = piBitmapFrameEncode->WritePixels(
            height
            , bitmap->rowBytes()
            , bitmap->rowBytes()*height
            , reinterpret_cast<BYTE*>(bitmap->getPixels()));
    }

    if (SUCCEEDED(hr)) {
        hr = piBitmapFrameEncode->Commit();
    }

    if (SUCCEEDED(hr)) {
        hr = piEncoder->Commit();
    }

    return SUCCEEDED(hr);
}

SkImageEncoder* SkImageEncoder::Create(Type t) {
    switch (t) {
        case kJPEG_Type:
        case kPNG_Type:
            break;
        default:
            return NULL;
    }
    return SkNEW_ARGS(SkImageEncoder_WIC, (t));
}

<commit_msg>Update WIC Image decoder to unpremul colors before save<commit_after>
/*
 * Copyright 2011 Google Inc.
 *
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#define WIN32_LEAN_AND_MEAN
#include <Windows.h>
#include <wincodec.h>
#include "SkAutoCoInitialize.h"
#include "SkImageDecoder.h"
#include "SkImageEncoder.h"
#include "SkIStream.h"
#include "SkMovie.h"
#include "SkStream.h"
#include "SkTScopedComPtr.h"
#include "SkUnPreMultiply.h"

class SkImageDecoder_WIC : public SkImageDecoder {
protected:
    virtual bool onDecode(SkStream* stream, SkBitmap* bm, Mode mode);
};

bool SkImageDecoder_WIC::onDecode(SkStream* stream, SkBitmap* bm, Mode mode) {
    //Initialize COM.
    SkAutoCoInitialize scopedCo;
    if (!scopedCo.succeeded()) {
        return false;
    }

    HRESULT hr = S_OK;

    //Create Windows Imaging Component ImagingFactory.
    SkTScopedComPtr<IWICImagingFactory> piImagingFactory;
    if (SUCCEEDED(hr)) {
        hr = CoCreateInstance(
            CLSID_WICImagingFactory
            , NULL
            , CLSCTX_INPROC_SERVER
            , IID_PPV_ARGS(&piImagingFactory)
        );
    }

    //Convert SkStream to IStream.
    SkTScopedComPtr<IStream> piStream;
    if (SUCCEEDED(hr)) {
        hr = SkIStream::CreateFromSkStream(stream, false, &piStream);
    }

    //Make sure we're at the beginning of the stream.
    if (SUCCEEDED(hr)) {
        LARGE_INTEGER liBeginning = { 0 };
        hr = piStream->Seek(liBeginning, STREAM_SEEK_SET, NULL);
    }

    //Create the decoder from the stream content.
    SkTScopedComPtr<IWICBitmapDecoder> piBitmapDecoder;
    if (SUCCEEDED(hr)) {
        hr = piImagingFactory->CreateDecoderFromStream(
            piStream.get()                    //Image to be decoded
            , NULL                            //No particular vendor
            , WICDecodeMetadataCacheOnDemand  //Cache metadata when needed
            , &piBitmapDecoder                //Pointer to the decoder
        );
    }

    //Get the first frame from the decoder.
    SkTScopedComPtr<IWICBitmapFrameDecode> piBitmapFrameDecode;
    if (SUCCEEDED(hr)) {
        hr = piBitmapDecoder->GetFrame(0, &piBitmapFrameDecode);
    }

    //Get the BitmapSource interface of the frame.
    SkTScopedComPtr<IWICBitmapSource> piBitmapSourceOriginal;
    if (SUCCEEDED(hr)) {
        hr = piBitmapFrameDecode->QueryInterface(
            IID_PPV_ARGS(&piBitmapSourceOriginal)
        );
    }

    //Get the size of the bitmap.
    UINT width;
    UINT height;
    if (SUCCEEDED(hr)) {
        hr = piBitmapSourceOriginal->GetSize(&width, &height);
    }

    //Exit early if we're only looking for the bitmap bounds.
    if (SUCCEEDED(hr)) {
        bm->setConfig(SkBitmap::kARGB_8888_Config, width, height);
        if (SkImageDecoder::kDecodeBounds_Mode == mode) {
            return true;
        }
        if (!this->allocPixelRef(bm, NULL)) {
            return false;
        }
    }

    //Create a format converter.
    SkTScopedComPtr<IWICFormatConverter> piFormatConverter;
    if (SUCCEEDED(hr)) {
        hr = piImagingFactory->CreateFormatConverter(&piFormatConverter);
    }

    if (SUCCEEDED(hr)) {
        hr = piFormatConverter->Initialize(
            piBitmapSourceOriginal.get()      //Input bitmap to convert
            , GUID_WICPixelFormat32bppPBGRA   //Destination pixel format
            , WICBitmapDitherTypeNone         //Specified dither patterm
            , NULL                            //Specify a particular palette
            , 0.f                             //Alpha threshold
            , WICBitmapPaletteTypeCustom      //Palette translation type
        );
    }

    //Get the BitmapSource interface of the format converter.
    SkTScopedComPtr<IWICBitmapSource> piBitmapSourceConverted;
    if (SUCCEEDED(hr)) {
        hr = piFormatConverter->QueryInterface(
            IID_PPV_ARGS(&piBitmapSourceConverted)
        );
    }

    //Copy the pixels into the bitmap.
    if (SUCCEEDED(hr)) {
        SkAutoLockPixels alp(*bm);
        bm->eraseColor(0);
        const int stride = bm->rowBytes();
        hr = piBitmapSourceConverted->CopyPixels(
            NULL,                             //Get all the pixels
            stride,
            stride * height,
            reinterpret_cast<BYTE *>(bm->getPixels())
        );

        // Note: we don't need to premultiply here since we specified PBGRA
        bm->computeAndSetOpaquePredicate();
    }

    return SUCCEEDED(hr);
}

/////////////////////////////////////////////////////////////////////////

SkImageDecoder* SkImageDecoder::Factory(SkStream* stream) {
    return SkNEW(SkImageDecoder_WIC);
}

/////////////////////////////////////////////////////////////////////////

SkMovie* SkMovie::DecodeStream(SkStream* stream) {
    return NULL;
}

/////////////////////////////////////////////////////////////////////////

class SkImageEncoder_WIC : public SkImageEncoder {
public:
    SkImageEncoder_WIC(Type t) : fType(t) {}

protected:
    virtual bool onEncode(SkWStream* stream, const SkBitmap& bm, int quality);

private:
    Type fType;
};

bool SkImageEncoder_WIC::onEncode(SkWStream* stream
                                , const SkBitmap& bitmapOrig
                                , int quality)
{
    GUID type;
    switch (fType) {
        case kJPEG_Type:
            type = GUID_ContainerFormatJpeg;
            break;
        case kPNG_Type:
            type = GUID_ContainerFormatPng;
            break;
        default:
            return false;
    }

    //Convert to 8888 if needed.
    const SkBitmap* bitmap;
    SkBitmap bitmapCopy;
    if (SkBitmap::kARGB_8888_Config == bitmapOrig.config() && bitmapOrig.isOpaque()) {
        bitmap = &bitmapOrig;
    } else {
        if (!bitmapOrig.copyTo(&bitmapCopy, SkBitmap::kARGB_8888_Config)) {
            return false;
        }
        bitmap = &bitmapCopy;
    }

    // We cannot use PBGRA so we need to unpremultiply ourselves
    if (!bitmap->isOpaque()) {
        SkAutoLockPixels alp(*bitmap);

        uint8_t* pixels = reinterpret_cast<uint8_t*>(bitmap->getPixels());
        for (int y = 0; y < bitmap->height(); ++y) {
            for (int x = 0; x < bitmap->width(); ++x) {
                uint8_t* bytes = pixels + y * bitmap->rowBytes() + x * bitmap->bytesPerPixel();

                SkPMColor* src = reinterpret_cast<SkPMColor*>(bytes);
                SkColor* dst = reinterpret_cast<SkColor*>(bytes);

                *dst = SkUnPreMultiply::PMColorToColor(*src);
            }
        }
    }

    //Initialize COM.
    SkAutoCoInitialize scopedCo;
    if (!scopedCo.succeeded()) {
        return false;
    }

    HRESULT hr = S_OK;

    //Create Windows Imaging Component ImagingFactory.
    SkTScopedComPtr<IWICImagingFactory> piImagingFactory;
    if (SUCCEEDED(hr)) {
        hr = CoCreateInstance(
            CLSID_WICImagingFactory
            , NULL
            , CLSCTX_INPROC_SERVER
            , IID_PPV_ARGS(&piImagingFactory)
        );
    }

    //Convert the SkWStream to an IStream.
    SkTScopedComPtr<IStream> piStream;
    if (SUCCEEDED(hr)) {
        hr = SkWIStream::CreateFromSkWStream(stream, &piStream);
    }

    //Create an encode of the appropriate type.
    SkTScopedComPtr<IWICBitmapEncoder> piEncoder;
    if (SUCCEEDED(hr)) {
        hr = piImagingFactory->CreateEncoder(type, NULL, &piEncoder);
    }

    if (SUCCEEDED(hr)) {
        hr = piEncoder->Initialize(piStream.get(), WICBitmapEncoderNoCache);
    }

    //Create a the frame.
    SkTScopedComPtr<IWICBitmapFrameEncode> piBitmapFrameEncode;
    SkTScopedComPtr<IPropertyBag2> piPropertybag;
    if (SUCCEEDED(hr)) {
        hr = piEncoder->CreateNewFrame(&piBitmapFrameEncode, &piPropertybag);
    }

    if (SUCCEEDED(hr)) {
        PROPBAG2 name = { 0 };
        name.dwType = PROPBAG2_TYPE_DATA;
        name.vt = VT_R4;
        name.pstrName = L"ImageQuality";

        VARIANT value;
        VariantInit(&value);
        value.vt = VT_R4;
        value.fltVal = (FLOAT)(quality / 100.0);

        //Ignore result code.
        //  This returns E_FAIL if the named property is not in the bag.
        //TODO(bungeman) enumerate the properties,
        //  write and set hr iff property exists.
        piPropertybag->Write(1, &name, &value);
    }
    if (SUCCEEDED(hr)) {
        hr = piBitmapFrameEncode->Initialize(piPropertybag.get());
    }

    //Set the size of the frame.
    const UINT width = bitmap->width();
    const UINT height = bitmap->height();
    if (SUCCEEDED(hr)) {
        hr = piBitmapFrameEncode->SetSize(width, height);
    }

    //Set the pixel format of the frame.
    const WICPixelFormatGUID formatDesired = GUID_WICPixelFormat32bppBGRA;
    WICPixelFormatGUID formatGUID = formatDesired;
    if (SUCCEEDED(hr)) {
        hr = piBitmapFrameEncode->SetPixelFormat(&formatGUID);
    }
    if (SUCCEEDED(hr)) {
        //Be sure the image format is the one requested.
        hr = IsEqualGUID(formatGUID, formatDesired) ? S_OK : E_FAIL;
    }

    //Write the pixels into the frame.
    if (SUCCEEDED(hr)) {
        SkAutoLockPixels alp(*bitmap);
        hr = piBitmapFrameEncode->WritePixels(
            height
            , bitmap->rowBytes()
            , bitmap->rowBytes()*height
            , reinterpret_cast<BYTE*>(bitmap->getPixels()));
    }

    if (SUCCEEDED(hr)) {
        hr = piBitmapFrameEncode->Commit();
    }

    if (SUCCEEDED(hr)) {
        hr = piEncoder->Commit();
    }

    return SUCCEEDED(hr);
}

SkImageEncoder* SkImageEncoder::Create(Type t) {
    switch (t) {
        case kJPEG_Type:
        case kPNG_Type:
            break;
        default:
            return NULL;
    }
    return SkNEW_ARGS(SkImageEncoder_WIC, (t));
}

<|endoftext|>