AlphaExpansion.hpp

Go to the documentation of this file.

#pragma once
 
#include "core/EnergyModel.hpp"
#include "solvers/MaxFlowSolver.hpp"
#include <memory>
#include <functional>
#include <type_traits>
 
template <typename T>
class AlphaExpansion {
public:
    using SolverFactory = std::function<std::unique_ptr<MaxFlowSolver<T>>(int num_vars, int num_edges)>;
 
    AlphaExpansion(EnergyModel<T> &model, SolverFactory solver_factory)
        : model_(model), solver_factory_(std::move(solver_factory)) {
    }
 
    void set_on_iteration(std::function<void()> cb) { on_iteration_ = std::move(cb); }
 
    [[nodiscard]] bool perform_expansion_move(const int alpha_label) const {
        on_iteration_();
        const std::vector<int> active_nodes = model_.get_active_nodes(alpha_label);
        if (active_nodes.empty()) return false;
 
        std::vector<typename MaxFlowSolver<T>::Var> node_var_ids;
        T old_subgraph_energy = 0;
        const auto solver = build_expansion_graph(alpha_label, active_nodes, node_var_ids, old_subgraph_energy);
 
        const T new_subgraph_energy = solver->minimize();
 
        bool improved;
        if constexpr (std::is_floating_point_v<T>) {
            improved = (old_subgraph_energy - new_subgraph_energy > static_cast<T>(1e-5));
        } else {
            improved = (new_subgraph_energy < old_subgraph_energy);
        }
        if (!improved) return false;
 
        bool changed = false;
        for (const int node: active_nodes) {
            if (const typename MaxFlowSolver<T>::Var var = node_var_ids[node]; solver->get_var(var) == 0) {
                model_.set_label(node, alpha_label);
                changed = true;
            }
        }
        return changed;
    }
 
private:
    std::unique_ptr<MaxFlowSolver<T>> build_expansion_graph(const int alpha_label, const std::vector<int> &active_nodes,
                                                            std::vector<typename MaxFlowSolver<T>::Var> &node_var_ids,
                                                            T &old_subgraph_energy) const {
        const int num_active = active_nodes.size();
        if (num_active == 0) return nullptr;
 
        const int estimated_edges = num_active * 4;
        auto solver = solver_factory_(num_active, estimated_edges);
 
        node_var_ids.assign(model_.num_nodes(), -1);
 
        for (int i = 0; i < num_active; ++i) {
            int node = active_nodes[i];
            node_var_ids[node] = solver->add_variable();
        }
 
        T old_energy = 0;
 
        for (const int node: active_nodes) {
            const typename MaxFlowSolver<T>::Var var = node_var_ids[node];
            const int current_label = model_.get_label(node);
            const T e0 = model_.get_unary_cost(node, alpha_label);
            const T e1 = model_.get_unary_cost(node, current_label);
            solver->add_term1(var, e0, e1);
            old_energy += e1;
        }
 
        for (const int node_i: active_nodes) {
            const int current_label_i = model_.get_label(node_i);
            const typename MaxFlowSolver<T>::Var var_i = node_var_ids[node_i];
 
            for (const int node_j: model_.get_neighbors_range(node_i)) {
                if (node_var_ids[node_j] != -1) {
                    if (node_i < node_j) {
                        const typename MaxFlowSolver<T>::Var var_j = node_var_ids[node_j];
                        const int current_label_j = model_.get_label(node_j);
                        const T e00 = model_.get_pairwise_cost(node_i, node_j, alpha_label, alpha_label);
                        const T e01 = model_.get_pairwise_cost(node_i, node_j, alpha_label, current_label_j);
                        const T e10 = model_.get_pairwise_cost(node_i, node_j, current_label_i, alpha_label);
                        const T e11 = model_.get_pairwise_cost(node_i, node_j, current_label_i, current_label_j);
                        solver->add_term2(var_i, var_j, e00, e01, e10, e11);
                        old_energy += e11;
                    }
                } else {
                    const int current_label_j = model_.get_label(node_j);
                    const T e0 = model_.get_pairwise_cost(node_i, node_j, alpha_label, current_label_j);
                    const T e1 = model_.get_pairwise_cost(node_i, node_j, current_label_i, current_label_j);
                    solver->add_term1(var_i, e0, e1);
                    old_energy += e1;
                }
            }
        }
 
        old_subgraph_energy = old_energy;
        return solver;
    }
 
    EnergyModel<T> &model_;
    SolverFactory solver_factory_;
    std::function<void()> on_iteration_ = []{};
};