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)) {
    }
 
    [[nodiscard]] bool perform_expansion_move(const int alpha_label) const {
        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;
        const auto solver = build_expansion_graph(alpha_label, active_nodes, node_var_ids);
 
        solver->minimize();
 
        bool changed = false;
        std::vector<int> proposed_labels = model_.get_labels();
 
        for (const int node: active_nodes) {
            if (const typename MaxFlowSolver<T>::Var var = node_var_ids[node]; solver->get_var(var) == 0) {
                proposed_labels[node] = alpha_label;
                changed = true;
            }
        }
 
        if (changed) {
            T old_energy = model_.evaluate_total_energy();
            T new_energy = model_.evaluate_total_energy(proposed_labels);
            bool improved = false;
            if constexpr (std::is_floating_point_v<T>) {
                improved = (old_energy - new_energy > static_cast<T>(1e-5));
            } else {
                improved = (new_energy < old_energy);
            }
            if (improved) {
                model_.set_labels(proposed_labels);
                return true;
            }
        }
 
        return false;
    }
 
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) 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();
        }
 
        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);
        }
 
        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(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);
                    }
                } 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);
                }
            }
        }
 
        return solver;
    }
 
    EnergyModel<T> &model_;
    SolverFactory solver_factory_;
};