285 lines
11 KiB
Rust
285 lines
11 KiB
Rust
|
#![allow(dead_code, unused_imports, unused_variables, unused_macros, unreachable_patterns)]
|
||
|
|
/// Grid-search drop shadow MAE tuning against Photoshop ground-truth PNGs.
|
||
|
|
///
|
||
|
|
/// Expects `_tmp/drop_shadow_png_set/` with transparent PNGs exported
|
||
|
|
/// from `_tmp/drop_shadow_psd_set/`.
|
||
|
|
///
|
||
|
|
/// Only pixels that differ from the transparent background are compared,
|
||
|
|
/// so the black rectangle body and the drop shadow pixels are included.
|
||
|
|
|
||
|
|
use std::io::Write;
|
||
|
|
|
||
|
|
const GT_DIR: &str = "_tmp/drop_shadow_png_set";
|
||
|
|
const CANVAS: u32 = 256;
|
||
|
|
const MARGIN: u32 = 64;
|
||
|
|
const RESULTS_FILE: &str = "drop_shadow_grid_results_v2.txt";
|
||
|
|
|
||
|
|
struct Sample {
|
||
|
|
name: String,
|
||
|
|
angle: f32,
|
||
|
|
distance: f32,
|
||
|
|
size: f32,
|
||
|
|
spread: f32,
|
||
|
|
opacity: f32,
|
||
|
|
noise: f32,
|
||
|
|
color: [u8; 4],
|
||
|
|
blend_mode: String,
|
||
|
|
ref_pixels: Vec<u8>,
|
||
|
|
/// Pixels that belong to the drop shadow (reference alpha > 0 but original shape alpha == 0).
|
||
|
|
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("ds_a") && 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 (angle, distance, size, spread, opacity, noise, blend_mode, 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 raw = rgba.into_raw();
|
||
|
|
let npix = (CANVAS * CANVAS) as usize;
|
||
|
|
let mut shape_mask = vec![false; npix];
|
||
|
|
for y in 0..CANVAS {
|
||
|
|
for x in 0..CANVAS {
|
||
|
|
if x >= MARGIN && x < CANVAS - MARGIN && y >= MARGIN && y < CANVAS - MARGIN {
|
||
|
|
shape_mask[(y * CANVAS + x) as usize] = true;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
let mut content_mask = vec![false; npix];
|
||
|
|
for i in 0..npix {
|
||
|
|
let a = raw[i * 4 + 3];
|
||
|
|
// Compare only shadow pixels: reference is non-transparent and outside the original shape.
|
||
|
|
content_mask[i] = a > 0 && !shape_mask[i];
|
||
|
|
}
|
||
|
|
|
||
|
|
samples.push(Sample {
|
||
|
|
name: stem,
|
||
|
|
angle, distance, size, spread, opacity, noise,
|
||
|
|
color, blend_mode,
|
||
|
|
ref_pixels: 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 y in 0..CANVAS {
|
||
|
|
for x in 0..CANVAS {
|
||
|
|
if x >= MARGIN && x < CANVAS - MARGIN && y >= MARGIN && y < CANVAS - MARGIN {
|
||
|
|
a[(y * CANVAS + x) as usize] = 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 spread_scales: &[f32] = &[0.0, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5];
|
||
|
|
|
||
|
|
let total_combos = blur_factors.len() * passes_list.len() * spread_scales.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);
|
||
|
|
let mut eval_count = 0usize;
|
||
|
|
|
||
|
|
let mut per_angle_best: std::collections::HashMap<i32, (f64, f32, i32, f32)> = std::collections::HashMap::new();
|
||
|
|
let angles: Vec<f32> = {
|
||
|
|
let mut set = std::collections::BTreeSet::new();
|
||
|
|
for s in &samples { set.insert(s.angle as i32); }
|
||
|
|
set.into_iter().map(|a| a as f32).collect()
|
||
|
|
};
|
||
|
|
|
||
|
|
for &blur_factor in blur_factors {
|
||
|
|
for &passes in passes_list {
|
||
|
|
for &spread_scale in spread_scales {
|
||
|
|
let mut total_mae = 0.0f64;
|
||
|
|
|
||
|
|
for s in &samples {
|
||
|
|
let shadow = hcie_fx::tuned::generate_shadow_tuned(
|
||
|
|
&rect_alpha, CANVAS, CANVAS,
|
||
|
|
s.angle, s.distance, s.spread, s.size,
|
||
|
|
s.color, false,
|
||
|
|
blur_factor, passes, spread_scale,
|
||
|
|
);
|
||
|
|
|
||
|
|
let blend = blend_mode_from_name(&s.blend_mode);
|
||
|
|
|
||
|
|
let px = npix();
|
||
|
|
let mut out = vec![0u8; px * 4];
|
||
|
|
// Black rectangle base
|
||
|
|
for i in 0..px {
|
||
|
|
if rect_alpha[i] > 0 {
|
||
|
|
out[i * 4] = 0;
|
||
|
|
out[i * 4 + 1] = 0;
|
||
|
|
out[i * 4 + 2] = 0;
|
||
|
|
out[i * 4 + 3] = 255;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
// Composite drop shadow under the shape using actual blend mode
|
||
|
|
for i in 0..px {
|
||
|
|
let sa = shadow[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 = [shadow[i * 4], shadow[i * 4 + 1], shadow[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 {
|
||
|
|
let mae = diff_sum / content_pixels as f64 / 4.0;
|
||
|
|
total_mae += mae;
|
||
|
|
|
||
|
|
let angle_key = s.angle as i32;
|
||
|
|
let entry = per_angle_best.entry(angle_key).or_insert((f64::MAX, 0.0, 0, 0.0));
|
||
|
|
if mae < entry.0 {
|
||
|
|
*entry = (mae, blur_factor, passes, spread_scale);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
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, spread_scale);
|
||
|
|
println!("*** [{}/{}] NEW BEST avg MAE = {:.4} (bf={:.1}, p={}, ss={:.2})",
|
||
|
|
eval_count, total_combos, best_avg_mae, blur_factor, passes, spread_scale);
|
||
|
|
std::io::stdout().flush().ok();
|
||
|
|
} else if eval_count % 100 == 0 {
|
||
|
|
println!(" [{}/{}] avg={:.4} best={:.4}", eval_count, total_combos, avg_mae, best_avg_mae);
|
||
|
|
std::io::stdout().flush().ok();
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
let (bbf, bp, bss) = best_params;
|
||
|
|
println!("\n=== RESULT ===");
|
||
|
|
println!("Best avg MAE = {:.4}", best_avg_mae);
|
||
|
|
println!(" blur_factor = {:.1}", bbf);
|
||
|
|
println!(" passes = {}", bp);
|
||
|
|
println!(" spread_scale = {:.2}", bss);
|
||
|
|
|
||
|
|
println!("\n=== PER-ANGLE BEST ===");
|
||
|
|
for angle in &angles {
|
||
|
|
let key = *angle as i32;
|
||
|
|
if let Some((mae, bf, p, ss)) = per_angle_best.get(&key) {
|
||
|
|
println!(" angle={:.0}: MAE={:.4} bf={:.1} p={} ss={:.2}",
|
||
|
|
angle, mae, bf, p, ss);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
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={} spread_scale={:.2}\n\
|
||
|
|
avg_mae={:.4} samples={} combos={}",
|
||
|
|
bbf, bp, bss, best_avg_mae, samples.len(), total_combos);
|
||
|
|
|
||
|
|
let _ = writeln!(file, "\n=== PER-ANGLE BEST ===");
|
||
|
|
for angle in &angles {
|
||
|
|
let key = *angle as i32;
|
||
|
|
if let Some((mae, bf, p, ss)) = per_angle_best.get(&key) {
|
||
|
|
let _ = writeln!(file, " angle={:.0}: MAE={:.4} bf={:.1} p={} ss={:.2}",
|
||
|
|
angle, mae, bf, p, ss);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn npix() -> usize { (CANVAS * CANVAS) as usize }
|
||
|
|
|
||
|
|
fn parse_filename(name: &str) -> Option<(f32, f32, f32, f32, f32, f32, String, [u8; 4])> {
|
||
|
|
// ds_a{angle}_d{distance}_s{size}_sp{spread}_op{opacity}_n{noise}_{blend}_{color}
|
||
|
|
let parts: Vec<&str> = name.split('_').collect();
|
||
|
|
if parts.len() < 9 || parts[0] != "ds" { return None; }
|
||
|
|
let angle = parts[1].strip_prefix('a')?.parse::<f32>().ok()?;
|
||
|
|
let distance = parts[2].strip_prefix('d')?.parse::<f32>().ok()?;
|
||
|
|
let size = parts[3].strip_prefix('s')?.parse::<f32>().ok()?;
|
||
|
|
let spread = parts[4].strip_prefix("sp")?.parse::<f32>().ok()?;
|
||
|
|
let opacity = parts[5].strip_prefix("op")?.parse::<f32>().ok()? / 100.0;
|
||
|
|
let noise = parts[6].strip_prefix('n')?.parse::<f32>().ok()?;
|
||
|
|
let blend = parts[7].to_string();
|
||
|
|
let color = color_from_name(parts[8]);
|
||
|
|
Some((angle, distance, size, spread, 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,
|
||
|
|
}
|
||
|
|
}
|