ICE_Dataset/CCPD_SPE_GEN/cpp_pipeline/main.cpp

#include "plate_swap.hpp"
#include <iostream>
#include <thread>
#include <mutex>
#include <atomic>
#include <chrono>
#include <filesystem>

namespace fs = std::filesystem;
using namespace std;
using namespace ice;

// ── 配置 ──
struct Config {
    vector<string> ccpd_dirs;
    string plate_type;
    string output_dir;
    string plate_pool_dir;
    int num_images = 500;
    unsigned workers = thread::hardware_concurrency();
    int seed = 42;
    int h_shift = 0, s_scale = 100, v_scale = 100;
};

// ── 参数解析 ──
Config parse_args(int argc, char* argv[]) {
    Config cfg;
    cfg.workers = thread::hardware_concurrency();

    for (int i = 1; i < argc; i++) {
        string arg = argv[i];
        if (arg == "--ccpd-dir" && i+1 < argc) {
            string dirs = argv[++i];
            size_t pos = 0;
            while (pos < dirs.size()) {
                size_t comma = dirs.find(',', pos);
                string d = (comma == string::npos) ? dirs.substr(pos) : dirs.substr(pos, comma-pos);
                // trim
                while (!d.empty() && d.front() == ' ') d.erase(0,1);
                while (!d.empty() && d.back() == ' ') d.pop_back();
                if (!d.empty()) cfg.ccpd_dirs.push_back(d);
                if (comma == string::npos) break;
                pos = comma + 1;
            }
        }
        else if (arg == "--plate-type" && i+1 < argc) cfg.plate_type = argv[++i];
        else if (arg == "--output" && i+1 < argc) cfg.output_dir = argv[++i];
        else if (arg == "--plate-pool" && i+1 < argc) cfg.plate_pool_dir = argv[++i];
        else if (arg == "--num" && i+1 < argc) cfg.num_images = stoi(argv[++i]);
        else if (arg == "--workers" && i+1 < argc) cfg.workers = stoi(argv[++i]);
        else if (arg == "--seed" && i+1 < argc) cfg.seed = stoi(argv[++i]);
        else if (arg == "--h-shift" && i+1 < argc) cfg.h_shift = stoi(argv[++i]);
        else if (arg == "--s-scale" && i+1 < argc) cfg.s_scale = stoi(argv[++i]);
        else if (arg == "--v-scale" && i+1 < argc) cfg.v_scale = stoi(argv[++i]);
    }

    return cfg;
}

// ── 收集并预解析 CCPD 图片（过滤无效，预存角点） ──
struct IndexedImage {
    string path;
    CCPDMeta meta;
};

vector<IndexedImage> collect_and_index(const vector<string>& dirs, mt19937& rng) {
    vector<IndexedImage> result;
    for (auto& d : dirs) {
        error_code ec;
        if (!fs::is_directory(d, ec)) {
            cerr << "[WARN] Dir not found: " << d << endl;
            continue;
        }
        for (auto& entry : fs::directory_iterator(d, ec)) {
            if (ec) break;
            if (!entry.is_regular_file()) continue;
            string ext = entry.path().extension().string();
            for (auto& c : ext) c = (char)tolower(c);
            if (ext != ".jpg" && ext != ".jpeg" && ext != ".png") continue;
            
            string fname = entry.path().filename().string();
            CCPDMeta meta = parse_ccpd_filename(fname);
            if (!meta.valid) continue;
            
            IndexedImage ii;
            ii.path = entry.path().string();
            ii.meta = meta;
            result.push_back(ii);
        }
    }
    shuffle(result.begin(), result.end(), rng);
    return result;
}

// ── 主函数 ──
int main(int argc, char* argv[]) {
    Config cfg = parse_args(argc, argv);

    if (cfg.ccpd_dirs.empty() || cfg.plate_type.empty()) {
        cerr << "Usage: plate_swap --ccpd-dir <dirs> --plate-type <type> [options]" << endl;
        cerr << "  --plate-pool <dir> Plate pool directory (required)" << endl;
        cerr << "  --num <N>          Number of images (default 500)" << endl;
        cerr << "  --workers <N>      Thread count (default all cores)" << endl;
        return 1;
    }

    cout << "=== ICE C++ Plate Swap ===" << endl;
    cout << "Type: " << cfg.plate_type << endl;
    cout << "Threads: " << cfg.workers << endl;
    cout << "Target: " << cfg.num_images << " images" << endl;

    if (cfg.plate_pool_dir.empty()) {
        cerr << "[ERROR] --plate-pool is required" << endl;
        return 1;
    }

    // Load plate pool
    auto t0 = chrono::steady_clock::now();
    cout << "Loading plate pool..." << flush;
    PlatePool pool;
    string bin_path = cfg.plate_pool_dir + "/pool.bin";
    if (fs::exists(bin_path))
        pool.load_bin(bin_path);
    else
        pool.load(cfg.plate_pool_dir);
    size_t pool_count = pool.count(cfg.plate_type);
    auto t1 = chrono::steady_clock::now();
    cout << " (" << pool_count << " available, " 
         << chrono::duration<double>(t1-t0).count() << "s)" << endl;
    if (pool_count == 0) {
        cerr << "[ERROR] No plates of type " << cfg.plate_type << " in pool" << endl;
        return 1;
    }

    // Collect & pre-parse CCPD (filter invalid, store corners)
    mt19937 rng(cfg.seed);
    cout << "Scanning & indexing CCPD..." << flush;
    vector<IndexedImage> images = collect_and_index(cfg.ccpd_dirs, rng);
    auto t2 = chrono::steady_clock::now();
    cout << " (" << images.size() << " valid, "
         << chrono::duration<double>(t2-t1).count() << "s)" << endl;
    if (images.empty()) {
        cerr << "[ERROR] No valid CCPD images found" << endl;
        return 1;
    }

    fs::create_directories(cfg.output_dir);

    SwapConfig swap_cfg;
    swap_cfg.plate_type = cfg.plate_type;
    swap_cfg.output_dir = cfg.output_dir;
    swap_cfg.h_shift = cfg.h_shift;
    swap_cfg.s_scale = cfg.s_scale;
    swap_cfg.v_scale = cfg.v_scale;

    // Parallel processing with pre-parsed metadata
    atomic<int> success{0}, failed{0};
    atomic<int> idx{0};
    int total = min((int)images.size(), cfg.num_images);
    mutex cout_mtx;

    auto t_proc = chrono::steady_clock::now();

    vector<thread> threads;
    for (unsigned t = 0; t < cfg.workers; t++) {
        threads.emplace_back([&](unsigned seed_offset) {
            mt19937 local_rng(cfg.seed + seed_offset * 10007);

            while (true) {
                int i = idx.fetch_add(1);
                if (i >= total) break;

                bool ok = process_single(images[i].path, swap_cfg, pool,
                                         local_rng, &images[i].meta);
                if (ok) {
                    int s = success.fetch_add(1) + 1;
                    if (s % 100 == 0 || s == total) {
                        lock_guard<mutex> lk(cout_mtx);
                        cout << "\rProgress: " << s << "/" << total << flush;
                    }
                } else {
                    failed.fetch_add(1);
                }
            }
        }, t);
    }

    for (auto& th : threads) th.join();

    auto t_done = chrono::steady_clock::now();
    double elapsed = chrono::duration<double>(t_done - t_proc).count();

    cout << "\rDone! Success: " << success << ", Failed: " << failed << endl;
    cout << "Process time: " << elapsed << "s  (" << (success/max(elapsed,0.01)) << " img/s)" << endl;
    cout << "Output: " << cfg.output_dir << endl;

    return 0;
}