feat: 特殊车牌生成管线 + C++/多进程/色温校正
This commit is contained in:
kirl committed
1 parent
3808b9f035
commit
3101992f51
29 files changed
+2336
-80
No files matched your search
@@ -0,0 +1,203 @@
|
||||
#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();
|
||||
bool fast = false;
|
||||
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 == "--fast") cfg.fast = true;
|
||||
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;
|
||||
cerr << " --fast Fast mode (skip degradation)" << 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;
|
||||
cout << "Fast: " << (cfg.fast ? "yes" : "no") << 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.fast_mode = cfg.fast;
|
||||
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;
|
||||
}
|
||||
Reference in new issue
Block a user