Files
hcie-rust-v3.05/hcie-io/examples/tune_mae_inner_glow_kra_grid.rs
T
2026-07-09 02:59:53 +03:00

266 lines
10 KiB
Rust

#![allow(dead_code, unused_imports, unused_variables, unused_macros)]
/// Grid-search inner glow MAE tuning against Krita ground-truth PNGs.
///
/// Expects `_tmp/inner_glow_png_kra_set/` with transparent PNGs exported
/// from `_tmp/inner_glow_kra_set/`.
use std::io::Write;
const GT_DIR: &str = "_tmp/inner_glow_png_kra_set";
static mut CANVAS: u32 = 800;
const RESULTS_FILE: &str = "inner_glow_kra_grid_results.txt";
struct Sample {
name: String,
choke: f32,
size: f32,
opacity: f32,
noise: f32,
color: [u8; 4],
blend_mode: String,
ref_pixels: Vec<u8>,
content_mask: Vec<bool>,
}
fn main() {
let entries: Vec<_> = std::fs::read_dir(GT_DIR)
.unwrap_or_else(|e| panic!("cannot read {}: {}", GT_DIR, e))
.filter_map(|e| e.ok())
.filter(|e| {
let s = e.file_name();
let s = s.to_string_lossy();
s.starts_with("ig_sp") && s.ends_with(".png")
})
.collect();
println!("Found {} PNGs in {}", entries.len(), GT_DIR);
let mut samples: Vec<Sample> = Vec::new();
let mut skipped_noise = 0;
for entry in &entries {
let path = entry.path();
let stem = path.file_stem().unwrap().to_string_lossy().to_string();
let params = match parse_filename(&stem) {
Some(p) => p,
None => { eprintln!(" SKIP: {}", stem); continue; }
};
let (choke, size, opacity, noise, blend, color) = params;
if noise > 0.0 {
skipped_noise += 1;
continue;
}
let img = match image::open(&path) {
Ok(i) => i,
Err(e) => { eprintln!(" SKIP {}: {}", stem, e); continue; }
};
let rgba = img.to_rgba8();
let (src_w, src_h) = rgba.dimensions();
if src_w != src_h {
eprintln!(" SKIP {}: non-square dimensions {}x{}", stem, src_w, src_h);
continue;
}
unsafe {
CANVAS = src_w;
}
let src_raw = rgba.as_raw();
let mut cropped_raw = vec![0u8; unsafe { (CANVAS * CANVAS * 4) as usize }];
cropped_raw.copy_from_slice(src_raw);
let npix = unsafe { (CANVAS * CANVAS) as usize };
let mut shape_mask = vec![false; npix];
for i in 0..npix {
let is_bg = cropped_raw[i * 4] == 128 && cropped_raw[i * 4 + 1] == 128 && cropped_raw[i * 4 + 2] == 128;
if !is_bg {
shape_mask[i] = true;
}
}
let mut content_mask = vec![false; npix];
for i in 0..npix {
// Glow is inside the shape (shape_mask is true) and differs from pure magenta.
let is_pure_shape = cropped_raw[i * 4] == 255 && cropped_raw[i * 4 + 1] == 0 && cropped_raw[i * 4 + 2] == 252;
content_mask[i] = shape_mask[i] && !is_pure_shape;
}
samples.push(Sample {
name: stem,
choke, size, opacity, noise, color, blend_mode: blend,
ref_pixels: cropped_raw,
content_mask,
});
}
println!("Loaded {} samples (skipped {} with noise>0)", samples.len(), skipped_noise);
if samples.is_empty() {
eprintln!("No valid samples found.");
std::process::exit(1);
}
let content_count = samples[0].content_mask.iter().filter(|&&b| b).count();
println!("Content pixels per image: {} / {} ({:.1}%)",
content_count, npix(), content_count as f64 / npix() as f64 * 100.0);
let rect_alpha: Vec<u8> = {
let n = npix();
let mut a = vec![0u8; n];
for i in 0..n {
let is_bg = samples[0].ref_pixels[i * 4] == 128 && samples[0].ref_pixels[i * 4 + 1] == 128 && samples[0].ref_pixels[i * 4 + 2] == 128;
if !is_bg {
a[i] = 255;
}
}
a
};
let blur_factors: &[f32] = &[1.5, 2.0, 2.5, 3.0, 4.0, 5.0, 6.0, 8.0, 10.0, 12.0];
let passes_list: &[i32] = &[1, 2, 3, 4, 5];
let alpha_scales: &[f32] = &[0.0, 0.25, 0.5, 0.75, 1.0];
let glow_boosts: &[f32] = &[0.5, 0.75, 1.0, 1.25, 1.5, 2.0];
let spread_positions: &[bool] = &[false, true]; // false = after blur, true = before blur
let total_combos = blur_factors.len() * passes_list.len() * alpha_scales.len()
* glow_boosts.len() * spread_positions.len();
println!("Grid: {} combos x {} samples = {} evals", total_combos, samples.len(), total_combos * samples.len());
std::io::stdout().flush().ok();
let mut best_avg_mae = f64::MAX;
let mut best_params = (0.0f32, 0i32, 0.0f32, 0.0f32, false);
let mut eval_count = 0usize;
for &blur_factor in blur_factors {
for &passes in passes_list {
for &alpha_scale in alpha_scales {
for &glow_boost in glow_boosts {
for &spread_before_blur in spread_positions {
let mut total_mae = 0.0f64;
for s in &samples {
let glow = hcie_fx::tuned::generate_glow_tuned(
&rect_alpha, unsafe { CANVAS }, unsafe { CANVAS },
s.choke, s.size, s.color, true,
blur_factor, passes, alpha_scale, glow_boost, spread_before_blur,
);
let blend = blend_mode_from_name(&s.blend_mode);
let px = npix();
// Start with 50% gray background
let mut out = vec![128u8; px * 4];
for i in 0..px {
out[i * 4 + 3] = 255;
}
// 1. Composite magenta shape on top using Normal blend
for i in 0..px {
if rect_alpha[i] > 0 {
out[i * 4] = 255;
out[i * 4 + 1] = 0;
out[i * 4 + 2] = 252;
out[i * 4 + 3] = 255;
}
}
// 2. Composite inner glow over the shape using actual blend mode
for i in 0..px {
let sa = glow[i * 4 + 3];
if sa == 0 { continue; }
let dst = [out[i * 4], out[i * 4 + 1], out[i * 4 + 2], out[i * 4 + 3]];
let src = [glow[i * 4], glow[i * 4 + 1], glow[i * 4 + 2], sa];
let blended = hcie_blend::blend_pixels(dst, src, blend, s.opacity);
out[i * 4..i * 4 + 4].copy_from_slice(&blended);
}
let mut diff_sum = 0.0f64;
let mut content_pixels = 0u64;
for i in 0..px {
if !s.content_mask[i] { continue; }
content_pixels += 1;
for c in 0..4 {
diff_sum += (out[i * 4 + c] as i32 - s.ref_pixels[i * 4 + c] as i32).abs() as f64;
}
}
if content_pixels > 0 {
total_mae += diff_sum / content_pixels as f64 / 4.0;
}
}
eval_count += 1;
let avg_mae = total_mae / samples.len() as f64;
if avg_mae < best_avg_mae {
best_avg_mae = avg_mae;
best_params = (blur_factor, passes, alpha_scale, glow_boost, spread_before_blur);
println!("*** [{}/{}] NEW BEST avg MAE = {:.4} (bf={:.1}, p={}, as={:.2}, gb={:.2}, spb={})",
eval_count, total_combos, best_avg_mae, blur_factor, passes, alpha_scale, glow_boost, spread_before_blur);
std::io::stdout().flush().ok();
} else if eval_count % 500 == 0 {
println!(" [{}/{}] avg={:.4} best={:.4}", eval_count, total_combos, avg_mae, best_avg_mae);
std::io::stdout().flush().ok();
}
}
}
}
}
}
let (bbf, bp, bas, bgb, bspb) = best_params;
println!("\n=== RESULT ===");
println!("Best avg MAE = {:.4}", best_avg_mae);
println!(" blur_factor = {:.1}", bbf);
println!(" passes = {}", bp);
println!(" alpha_scale = {:.2}", bas);
println!(" glow_boost = {:.2}", bgb);
println!(" spread_before_blur= {}", bspb);
if let Ok(mut file) = std::fs::OpenOptions::new()
.create(true).write(true).truncate(true).open(RESULTS_FILE)
{
let _ = writeln!(file,
"blur_factor={:.1} passes={} alpha_scale={:.2} glow_boost={:.2} spread_before_blur={}\n\
avg_mae={:.4} samples={} combos={}",
bbf, bp, bas, bgb, bspb, best_avg_mae, samples.len(), total_combos);
}
}
fn npix() -> usize { unsafe { (CANVAS * CANVAS) as usize } }
fn parse_filename(name: &str) -> Option<(f32, f32, f32, f32, String, [u8; 4])> {
// ig_sp{choke}_s{size}_op{opacity}_n{noise}_{blend}_{color}
let parts: Vec<&str> = name.split('_').collect();
if parts.len() < 7 || parts[0] != "ig" { return None; }
let choke = parts[1].strip_prefix("sp")?.parse::<f32>().ok()?;
let size = parts[2].strip_prefix('s')?.parse::<f32>().ok()?;
let opacity = parts[3].strip_prefix("op")?.parse::<f32>().ok()? / 100.0;
let noise = parts[4].strip_prefix('n')?.parse::<f32>().ok()?;
let blend = parts[5].to_string();
let color = color_from_name(parts[6]);
Some((choke, size, opacity, noise, blend, color))
}
fn color_from_name(name: &str) -> [u8; 4] {
match name {
"red" => [220, 30, 30, 255],
"green" => [30, 180, 30, 255],
"blue" => [30, 60, 220, 255],
"white" => [220, 220, 220, 255],
"yellow" => [255, 220, 30, 255],
"cyan" => [30, 220, 220, 255],
"magenta" => [255, 30, 180, 255],
_ => [220, 30, 30, 255],
}
}
fn blend_mode_from_name(name: &str) -> hcie_blend::BlendMode {
match name {
"multiply" => hcie_blend::BlendMode::Multiply,
"screen" => hcie_blend::BlendMode::Screen,
"normal" => hcie_blend::BlendMode::Normal,
_ => hcie_blend::BlendMode::Normal,
}
}