forked from Keelrleek/SoulStream
KeelrLeek
a667b4b28e
The previous .git directory was found gutted (only info/ remained), so history is unrecoverable locally. This commit restores version control with the tree exactly as it stands: the working app, the preview-tracks launcher on the artist hero, the scrollable-track-previews design bundle (docs + approved mockups), and the Phase 0 validation artifacts.
656 lines
25 KiB
Rust
656 lines
25 KiB
Rust
//! Derive an artist wordmark from album art for artists that have no provider
|
||
//! (Fanart.tv / TheAudioDB) logo. Many covers print the band name as clean text;
|
||
//! on-device OCR (Windows.Media.Ocr) locates that line, and (increment 2) it is
|
||
//! segmented off a near-uniform background into a transparent PNG that the
|
||
//! existing logo tracer vectorizes and tints like any provider logo.
|
||
//!
|
||
//! Fallible by design: any failure (non-Windows build, no OCR, no matching text)
|
||
//! leaves the caller on the plain text-name fallback.
|
||
|
||
/// Axis-aligned box in source-image pixels.
|
||
#[derive(Clone, Copy, Debug)]
|
||
pub(crate) struct Rect {
|
||
pub x: f32,
|
||
pub y: f32,
|
||
pub w: f32,
|
||
pub h: f32,
|
||
}
|
||
|
||
/// One recognized word with its own bounding box.
|
||
#[derive(Clone, Debug)]
|
||
pub(crate) struct OcrWord {
|
||
pub text: String,
|
||
pub rect: Rect,
|
||
}
|
||
|
||
/// One recognized line of text with its bounding box (union of its words) plus
|
||
/// the individual words — so a wordmark can be boxed to JUST the words that spell
|
||
/// the artist name, excluding a track number / remix credit sharing the line.
|
||
#[derive(Clone, Debug)]
|
||
pub(crate) struct OcrLine {
|
||
pub text: String,
|
||
pub rect: Rect,
|
||
pub words: Vec<OcrWord>,
|
||
}
|
||
|
||
/// Lowercase alphanumerics only — matches the normalization used elsewhere so a
|
||
/// cover's "CAMPING IN ALASKA" lines up with the artist name "Camping In Alaska".
|
||
fn normalize(s: &str) -> String {
|
||
s.to_lowercase()
|
||
.chars()
|
||
.filter(|c| c.is_alphanumeric())
|
||
.collect()
|
||
}
|
||
|
||
// Segmentation tuning (RGB distance space, 0..~441).
|
||
const BG_NEAR_DIST: f32 = 46.0; // a pixel this close to the dominant colour counts as background
|
||
const BG_MIN_DOMINANCE: f32 = 0.5; // background must fill at least this share of the crop, else it's a photo → reject
|
||
const ALPHA_LO: f32 = 46.0; // distance-from-bg below which a pixel is background (transparent)
|
||
const ALPHA_HI: f32 = 120.0; // distance-from-bg above which a pixel is full ink (opaque)
|
||
|
||
fn dist(a: [f32; 3], b: [f32; 3]) -> f32 {
|
||
((a[0] - b[0]).powi(2) + (a[1] - b[1]).powi(2) + (a[2] - b[2]).powi(2)).sqrt()
|
||
}
|
||
|
||
fn smoothstep(lo: f32, hi: f32, x: f32) -> f32 {
|
||
let t = ((x - lo) / (hi - lo)).clamp(0.0, 1.0);
|
||
t * t * (3.0 - 2.0 * t)
|
||
}
|
||
|
||
/// A wordmark lifted off a cover, plus a quality score for picking the best cover.
|
||
pub(crate) struct Wordmark {
|
||
pub image: image::RgbaImage,
|
||
pub score: f32,
|
||
}
|
||
|
||
/// The band-name line on a cover plus the vertical band it may occupy. `clamp_top`
|
||
/// / `clamp_bottom` are the midpoints to the nearest OTHER text lines (or image
|
||
/// edges), so the extraction crop can NEVER reach another line — the album title
|
||
/// and any other cover text are structurally excluded from the logo.
|
||
#[derive(Clone, Copy, Debug)]
|
||
pub(crate) struct BandLocation {
|
||
pub rect: Rect,
|
||
pub clamp_top: f32,
|
||
pub clamp_bottom: f32,
|
||
}
|
||
|
||
impl BandLocation {
|
||
/// Scale to another resolution of the same cover (locate at 1000px where OCR
|
||
/// is reliable, extract at 1800px for quality).
|
||
pub fn scaled(&self, sx: f32, sy: f32) -> BandLocation {
|
||
BandLocation {
|
||
rect: Rect {
|
||
x: self.rect.x * sx,
|
||
y: self.rect.y * sy,
|
||
w: self.rect.w * sx,
|
||
h: self.rect.h * sy,
|
||
},
|
||
clamp_top: self.clamp_top * sy,
|
||
clamp_bottom: self.clamp_bottom * sy,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Segment the band-name text at `loc` off its (near-uniform) background into a
|
||
/// transparent RGBA image. The crop is hard-clamped to `loc`'s band so no other
|
||
/// text line can enter. Returns `None` when the background is too busy to segment
|
||
/// (text over a photo) or the result is empty — those stay on the text fallback.
|
||
pub(crate) fn extract_wordmark(
|
||
image: &image::DynamicImage,
|
||
loc: &BandLocation,
|
||
) -> Option<Wordmark> {
|
||
use image::GenericImageView as _;
|
||
let rect = loc.rect;
|
||
let (iw, ih) = image.dimensions();
|
||
// Pad the located word box — vertical for ascenders/descenders + tracer room —
|
||
// but never past the clamp to a neighbouring line.
|
||
let pad_x = (rect.w * 0.05).max(6.0);
|
||
let pad_y = (rect.h * 0.45).max(6.0);
|
||
let x0 = (rect.x - pad_x).max(0.0) as u32;
|
||
let y0 = (rect.y - pad_y).max(loc.clamp_top).max(0.0) as u32;
|
||
let x1 = ((rect.x + rect.w + pad_x).min(iw as f32)) as u32;
|
||
let y1 = ((rect.y + rect.h + pad_y)
|
||
.min(loc.clamp_bottom)
|
||
.min(ih as f32)) as u32;
|
||
if x1 <= x0 + 4 || y1 <= y0 + 4 {
|
||
return None;
|
||
}
|
||
let crop = image.crop_imm(x0, y0, x1 - x0, y1 - y0).to_rgba8();
|
||
let (w, h) = crop.dimensions();
|
||
|
||
// Background = the dominant colour, via a coarse 16³ RGB histogram over the
|
||
// whole crop. Robust to sparse texture (faint line-art on black) in a way a
|
||
// ring mean is not — the text is a minority, so the mode is the background.
|
||
let mut hist = [0u32; 16 * 16 * 16];
|
||
for p in crop.pixels() {
|
||
let (r, g, b) = (p.0[0] >> 4, p.0[1] >> 4, p.0[2] >> 4);
|
||
hist[(r as usize) * 256 + (g as usize) * 16 + b as usize] += 1;
|
||
}
|
||
let mode = hist
|
||
.iter()
|
||
.enumerate()
|
||
.max_by_key(|(_, c)| **c)
|
||
.map(|(i, _)| i)
|
||
.unwrap_or(0);
|
||
let bg = [
|
||
((mode / 256) as f32) * 16.0 + 8.0,
|
||
(((mode / 16) % 16) as f32) * 16.0 + 8.0,
|
||
((mode % 16) as f32) * 16.0 + 8.0,
|
||
];
|
||
// Dominance guard: what share of the crop is near that background colour?
|
||
// A clean wordmark has a dominant background (text is the minority); a photo
|
||
// has no single dominant colour, so this share is low → reject.
|
||
let total = (w * h) as f32;
|
||
let near = crop
|
||
.pixels()
|
||
.filter(|p| dist([p.0[0] as f32, p.0[1] as f32, p.0[2] as f32], bg) < BG_NEAR_DIST)
|
||
.count() as f32;
|
||
let dominance = near / total;
|
||
if dominance < BG_MIN_DOMINANCE {
|
||
return None; // busy background — text over a photo, not segmentable
|
||
}
|
||
|
||
// Alpha ramp by distance from the background colour, tracking per-row ink so
|
||
// we can trim adjacent text lines the padding caught (the album title above /
|
||
// below the band name), which sit in their own ink band across an empty gap.
|
||
let mut full = image::RgbaImage::new(w, h);
|
||
let mut row_ink = vec![0u32; h as usize];
|
||
for y in 0..h {
|
||
for x in 0..w {
|
||
let p = crop.get_pixel(x, y).0;
|
||
let d = dist([p[0] as f32, p[1] as f32, p[2] as f32], bg);
|
||
let a = smoothstep(ALPHA_LO, ALPHA_HI, d);
|
||
let alpha = (a * 255.0).round() as u8;
|
||
// Flat white ink (not the cover's own colour). The frontend logo pass
|
||
// only traces the ALPHA silhouette and tints it to the accent, and it
|
||
// only does so when the mark is monochrome — a cover's varied/print
|
||
// colour scored below that bar and fell back to a small untinted PNG.
|
||
// A single flat colour guarantees the traced, accent-tinted SVG path.
|
||
full.put_pixel(x, y, image::Rgba([255, 255, 255, alpha]));
|
||
if alpha > 128 {
|
||
row_ink[y as usize] += 1;
|
||
}
|
||
}
|
||
}
|
||
// The wordmark is CENTRED on the located box. A single line of text never has a
|
||
// fully transparent row running through it — some glyph always spans every row
|
||
// (an i-dot's little gap is filled by neighbouring letters). So the first fully
|
||
// transparent row above/below the centre is the wordmark's TRUE edge; whatever
|
||
// lies past that transparent gap is separated content — a rule bar above, the
|
||
// album-title ghost below — and is dropped. We crop to the centre content, then
|
||
// re-add clean transparent padding.
|
||
let empty = |r: usize| (row_ink[r] as f32 / w as f32) < 0.004;
|
||
// A near-fully-inked row is a solid rule/bar touching the text with no gap —
|
||
// also a boundary (glyph rows always leave inter-letter gaps, never ~full).
|
||
let solid = |r: usize| row_ink[r] as f32 / w as f32 > 0.85;
|
||
let center = (((rect.y + rect.h / 2.0) as u32).saturating_sub(y0)).min(h - 1) as usize;
|
||
// 2 transparent rows = a real gap (bridges a 1px anti-aliasing fluke).
|
||
let gap_stop = 2usize;
|
||
// Safety cap so a pathological run can't swallow the whole cover.
|
||
let max_half = (rect.h * 1.5).round().max(6.0) as usize;
|
||
let lo = center.saturating_sub(max_half);
|
||
let hi = (center + max_half).min(h as usize - 1);
|
||
let mut top = center;
|
||
{
|
||
let (mut r, mut gap) = (center, 0usize);
|
||
while r > lo {
|
||
r -= 1;
|
||
if solid(r) {
|
||
break;
|
||
}
|
||
if empty(r) {
|
||
gap += 1;
|
||
if gap >= gap_stop {
|
||
break;
|
||
}
|
||
} else {
|
||
gap = 0;
|
||
top = r;
|
||
}
|
||
}
|
||
}
|
||
let mut bottom = center;
|
||
{
|
||
let (mut r, mut gap) = (center, 0usize);
|
||
while r + 1 <= hi {
|
||
r += 1;
|
||
if solid(r) {
|
||
break;
|
||
}
|
||
if empty(r) {
|
||
gap += 1;
|
||
if gap >= gap_stop {
|
||
break;
|
||
}
|
||
} else {
|
||
gap = 0;
|
||
bottom = r;
|
||
}
|
||
}
|
||
}
|
||
// Re-add transparent breathing room (~6% of content height) for the tracer/blur.
|
||
let pad = (((bottom - top + 1) as f32) * 0.06)
|
||
.round()
|
||
.clamp(3.0, 40.0) as usize;
|
||
let band_top = top.saturating_sub(pad) as u32;
|
||
let band_bottom = (bottom + pad).min(h as usize - 1) as u32;
|
||
let band_h = band_bottom - band_top + 1;
|
||
if band_h < 6 {
|
||
return None;
|
||
}
|
||
|
||
// Crop to the wordmark's row band and tally ink/partial for scoring.
|
||
let mut out = image::RgbaImage::new(w, band_h);
|
||
let mut ink = 0u32;
|
||
let mut partial = 0u32;
|
||
for (dy, y) in (band_top..=band_bottom).enumerate() {
|
||
for x in 0..w {
|
||
let px = *full.get_pixel(x, y);
|
||
out.put_pixel(x, dy as u32, px);
|
||
if px.0[3] > 200 {
|
||
ink += 1;
|
||
} else if px.0[3] > 20 {
|
||
partial += 1;
|
||
}
|
||
}
|
||
}
|
||
let coverage = ink as f32 / (w * band_h) as f32;
|
||
if coverage < 0.008 || coverage > 0.6 {
|
||
return None;
|
||
}
|
||
// Recipe D: median despeckle removes the ragged spurs a lo-fi print leaves on
|
||
// the letter edges before the frontend supersamples/blurs/traces it.
|
||
let out = median_alpha(&out);
|
||
// Score: reward anti-aliasing (edge quality), a title-like wide aspect, and
|
||
// a healthy ink coverage; higher wins across candidate covers.
|
||
let aa = (partial as f32 / (ink as f32 + 1.0)).min(0.6);
|
||
let aspect = ((w as f32 / band_h as f32) / 10.0).min(0.6);
|
||
let cover_score = 1.0 - (coverage - 0.15).abs();
|
||
Some(Wordmark {
|
||
image: out,
|
||
score: aa + aspect + cover_score,
|
||
})
|
||
}
|
||
|
||
// ---- Experimental alpha post-processing (for the quality comparison) ----
|
||
|
||
/// 3×3 median on the alpha channel — despeckles the ragged spurs a noisy print
|
||
/// leaves on the letter edges, without rounding real corners much.
|
||
fn median_alpha(src: &image::RgbaImage) -> image::RgbaImage {
|
||
let (w, h) = src.dimensions();
|
||
let mut out = src.clone();
|
||
for y in 0..h {
|
||
for x in 0..w {
|
||
// The previous implementation allocated a Vec for every output
|
||
// pixel. A fixed nine-byte buffer computes the identical median.
|
||
let mut n = [0u8; 9];
|
||
let mut len = 0usize;
|
||
for dy in -1..=1 {
|
||
for dx in -1..=1 {
|
||
let nx = x as i32 + dx;
|
||
let ny = y as i32 + dy;
|
||
if nx >= 0 && ny >= 0 && (nx as u32) < w && (ny as u32) < h {
|
||
n[len] = src.get_pixel(nx as u32, ny as u32).0[3];
|
||
len += 1;
|
||
}
|
||
}
|
||
}
|
||
n[..len].sort_unstable();
|
||
let mut p = *src.get_pixel(x, y);
|
||
p.0[3] = n[len / 2];
|
||
out.put_pixel(x, y, p);
|
||
}
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Otsu-binarized copy (black text on white / white on black, auto-thresholded on
|
||
/// luma). Experimental preprocessing to test whether the OCR reads stylized
|
||
/// display type better with the colour and gradients flattened away.
|
||
pub(crate) fn binarize(image: &image::DynamicImage) -> image::DynamicImage {
|
||
use image::GenericImageView as _;
|
||
let gray = image.to_luma8();
|
||
// Otsu threshold.
|
||
let mut hist = [0u32; 256];
|
||
for p in gray.pixels() {
|
||
hist[p.0[0] as usize] += 1;
|
||
}
|
||
let total: u32 = hist.iter().sum();
|
||
let sum: f64 = hist
|
||
.iter()
|
||
.enumerate()
|
||
.map(|(i, &c)| i as f64 * c as f64)
|
||
.sum();
|
||
let (mut sum_b, mut w_b, mut max, mut thresh) = (0.0f64, 0u32, 0.0f64, 128u8);
|
||
for (t, &c) in hist.iter().enumerate() {
|
||
w_b += c;
|
||
if w_b == 0 {
|
||
continue;
|
||
}
|
||
let w_f = total - w_b;
|
||
if w_f == 0 {
|
||
break;
|
||
}
|
||
sum_b += t as f64 * c as f64;
|
||
let m_b = sum_b / w_b as f64;
|
||
let m_f = (sum - sum_b) / w_f as f64;
|
||
let between = w_b as f64 * w_f as f64 * (m_b - m_f).powi(2);
|
||
if between > max {
|
||
max = between;
|
||
thresh = t as u8;
|
||
}
|
||
}
|
||
// Decide polarity so text is dark on light (OCR's expectation) by majority.
|
||
let (w, h) = image.dimensions();
|
||
let mut out = image::RgbaImage::new(w, h);
|
||
for (x, y, p) in gray.enumerate_pixels().map(|(x, y, p)| (x, y, p.0[0])) {
|
||
let v = if p > thresh { 255 } else { 0 };
|
||
out.put_pixel(x, y, image::Rgba([v, v, v, 255]));
|
||
}
|
||
image::DynamicImage::ImageRgba8(out)
|
||
}
|
||
|
||
/// OCR the cover twice — once as-is and once Otsu-binarized — and merge the line
|
||
/// sets. Colour/gradient display type (the cyan "CRAZY P" on Paper Years) is
|
||
/// often unreadable raw but recognized once flattened to black-on-white; the raw
|
||
/// pass still catches anything the threshold destroys. Both passes share pixel
|
||
/// coordinates (the binarized copy is the same dimensions), so word boxes line up.
|
||
pub(crate) fn ocr_lines_multi(image: &image::DynamicImage) -> Result<Vec<OcrLine>, String> {
|
||
let mut lines = ocr_lines(image)?;
|
||
if let Ok(mut bin) = ocr_lines(&binarize(image)) {
|
||
lines.append(&mut bin);
|
||
}
|
||
Ok(lines)
|
||
}
|
||
|
||
/// All text lines OCR finds on the image, in source-pixel coordinates.
|
||
pub(crate) fn ocr_lines(image: &image::DynamicImage) -> Result<Vec<OcrLine>, String> {
|
||
#[cfg(windows)]
|
||
{
|
||
windows_ocr::ocr_lines(image)
|
||
}
|
||
#[cfg(not(windows))]
|
||
{
|
||
let _ = image;
|
||
Err("cover-logo OCR is only available on Windows".into())
|
||
}
|
||
}
|
||
|
||
/// A substring match must be at least this share of its OCR line's letters, or
|
||
/// the line is a tracklist / credit that merely MENTIONS the name (e.g.
|
||
/// "A1. Disc Odyssey Crazy P remix" — "crazyp" is 25% of it) rather than a
|
||
/// wordmark that IS the name. Precision-first: only "just the band name" lines
|
||
/// (optionally with a hyphen/initial) clear this bar, never a whole credit line.
|
||
const NAME_LINE_MIN_RATIO: f32 = 0.6;
|
||
|
||
/// Locate the band-name line and the vertical band it may occupy — clamped at
|
||
/// the midpoint to the nearest OTHER text line above/below (or the image edge),
|
||
/// so extraction can never include the album title or any other cover text.
|
||
/// Scores every line so a bare-name line beats one that merely embeds the name,
|
||
/// and rejects lines where the name is a minority (a tracklist / remix credit).
|
||
pub(crate) fn locate_wordmark(
|
||
image: &image::DynamicImage,
|
||
artist: &str,
|
||
) -> Result<Option<BandLocation>, String> {
|
||
Ok(locate_wordmark_named(image, artist)?.map(|(loc, _, _)| loc))
|
||
}
|
||
|
||
/// Within one OCR line, find the consecutive run of WORDS whose normalized text
|
||
/// spells the artist name exactly, and return the union box of just those words.
|
||
/// This lifts "Crazy P" cleanly out of "A1. Disc Odyssey Crazy P remix" or
|
||
/// "CRAZY P THE PAPER YEARS" — the track number, remix credit and album title
|
||
/// are never inside the returned box. `None` if no run spells the name.
|
||
fn word_run_box(line: &OcrLine, want: &str) -> Option<Rect> {
|
||
// Only real word tokens: a "///", "A1." or "-" normalizes to empty and must
|
||
// NOT be allowed into the run's box, or the decorative mark next to the name
|
||
// gets extracted with it (Turnover's "/// TURNOVER").
|
||
let toks: Vec<&OcrWord> = line
|
||
.words
|
||
.iter()
|
||
.filter(|w| !normalize(&w.text).is_empty())
|
||
.collect();
|
||
for i in 0..toks.len() {
|
||
let mut acc = String::new();
|
||
for j in i..toks.len() {
|
||
acc.push_str(&normalize(&toks[j].text));
|
||
if acc.len() > want.len() {
|
||
break;
|
||
}
|
||
if acc == want {
|
||
let (mut x0, mut y0, mut x1, mut y1) = (f32::MAX, f32::MAX, f32::MIN, f32::MIN);
|
||
for w in &toks[i..=j] {
|
||
x0 = x0.min(w.rect.x);
|
||
y0 = y0.min(w.rect.y);
|
||
x1 = x1.max(w.rect.x + w.rect.w);
|
||
y1 = y1.max(w.rect.y + w.rect.h);
|
||
}
|
||
return Some(Rect {
|
||
x: x0,
|
||
y: y0,
|
||
w: x1 - x0,
|
||
h: y1 - y0,
|
||
});
|
||
}
|
||
}
|
||
}
|
||
None
|
||
}
|
||
|
||
/// Like [`locate_wordmark`] but also returns the matched OCR text and its score
|
||
/// (1.0 = a clean word-run of the name; <1 = a line the name only partly fills).
|
||
/// For diagnostics — "which line on which cover did we lift?".
|
||
pub(crate) fn locate_wordmark_named(
|
||
image: &image::DynamicImage,
|
||
artist: &str,
|
||
) -> Result<Option<(BandLocation, String, f32)>, String> {
|
||
let want = normalize(artist);
|
||
if want.is_empty() {
|
||
return Ok(None);
|
||
}
|
||
// Two-pass OCR (raw + binarized) so stylized colour wordmarks are read.
|
||
let lines = ocr_lines_multi(image)?;
|
||
|
||
// Prefer a precise word-run match (the box is JUST the name); among those,
|
||
// take the widest — a cover's big printed wordmark over a tiny corner label.
|
||
// Fall back to a whole-line substring only if no cover word-run matches, and
|
||
// only when the name fills most of the line (else it's a tracklist/credit).
|
||
let mut word_best: Option<(usize, Rect)> = None;
|
||
let mut line_best: Option<(usize, f32)> = None;
|
||
for (i, line) in lines.iter().enumerate() {
|
||
if let Some(rect) = word_run_box(line, &want) {
|
||
if word_best.is_none_or(|(_, r)| rect.w > r.w) {
|
||
word_best = Some((i, rect));
|
||
}
|
||
continue;
|
||
}
|
||
let norm = normalize(&line.text);
|
||
if norm.is_empty() {
|
||
continue;
|
||
}
|
||
let score = if norm == want {
|
||
1.0
|
||
} else if want.len() >= 4 && norm.contains(&want) {
|
||
let ratio = want.len() as f32 / norm.len() as f32;
|
||
if ratio >= NAME_LINE_MIN_RATIO {
|
||
ratio
|
||
} else {
|
||
continue;
|
||
}
|
||
} else {
|
||
continue;
|
||
};
|
||
if line_best.is_none_or(|(_, s)| score > s) {
|
||
line_best = Some((i, score));
|
||
}
|
||
}
|
||
let (bi, band, matched_score) = if let Some((i, rect)) = word_best {
|
||
(i, rect, 1.0)
|
||
} else if let Some((i, s)) = line_best {
|
||
(i, lines[i].rect, s)
|
||
} else {
|
||
return Ok(None);
|
||
};
|
||
let matched_text = lines[bi].text.clone();
|
||
let mut clamp_top = 0.0f32;
|
||
let mut clamp_bottom = image.height() as f32;
|
||
for (i, line) in lines.iter().enumerate() {
|
||
if i == bi {
|
||
continue;
|
||
}
|
||
let r = line.rect;
|
||
let r_bottom = r.y + r.h;
|
||
if r_bottom <= band.y {
|
||
// A line above — stop halfway up to it.
|
||
clamp_top = clamp_top.max((r_bottom + band.y) / 2.0);
|
||
} else if r.y >= band.y + band.h {
|
||
// A line below — stop halfway down to it.
|
||
clamp_bottom = clamp_bottom.min((band.y + band.h + r.y) / 2.0);
|
||
}
|
||
}
|
||
Ok(Some((
|
||
BandLocation {
|
||
rect: band,
|
||
clamp_top,
|
||
clamp_bottom,
|
||
},
|
||
matched_text,
|
||
matched_score,
|
||
)))
|
||
}
|
||
|
||
#[cfg(windows)]
|
||
mod windows_ocr {
|
||
use image::{DynamicImage, ImageEncoder as _};
|
||
use windows::{
|
||
Graphics::Imaging::{BitmapDecoder, BitmapPixelFormat, SoftwareBitmap},
|
||
Media::Ocr::OcrEngine,
|
||
Storage::Streams::{DataWriter, InMemoryRandomAccessStream},
|
||
Win32::System::WinRT::{RoInitialize, RoUninitialize, RO_INIT_MULTITHREADED},
|
||
};
|
||
|
||
use super::{OcrLine, OcrWord, Rect};
|
||
|
||
struct WinRtGuard;
|
||
impl Drop for WinRtGuard {
|
||
fn drop(&mut self) {
|
||
unsafe { RoUninitialize() };
|
||
}
|
||
}
|
||
|
||
fn initialize_winrt() -> Result<Option<WinRtGuard>, String> {
|
||
match unsafe { RoInitialize(RO_INIT_MULTITHREADED) } {
|
||
Ok(()) => Ok(Some(WinRtGuard)),
|
||
// Thread already in a different apartment: WinRT is usable, we just
|
||
// must not uninitialize it.
|
||
Err(error) if error.code().0 as u32 == 0x8001_0106 => Ok(None),
|
||
Err(error) => Err(format!("Windows Runtime initialization failed: {error}")),
|
||
}
|
||
}
|
||
|
||
/// Load a DynamicImage into a Bgra8 SoftwareBitmap (the format OCR wants),
|
||
/// reusing the same PNG→stream→decoder path the face detector uses.
|
||
fn software_bitmap(image: &DynamicImage) -> Result<SoftwareBitmap, String> {
|
||
let rgba = image.to_rgba8();
|
||
let (width, height) = rgba.dimensions();
|
||
let mut png = Vec::new();
|
||
image::codecs::png::PngEncoder::new(&mut png)
|
||
.write_image(
|
||
rgba.as_raw(),
|
||
width,
|
||
height,
|
||
image::ExtendedColorType::Rgba8,
|
||
)
|
||
.map_err(|error| format!("could not encode OCR raster: {error}"))?;
|
||
|
||
let stream = InMemoryRandomAccessStream::new()
|
||
.map_err(|error| format!("could not allocate OCR stream: {error}"))?;
|
||
let writer = DataWriter::CreateDataWriter(&stream)
|
||
.map_err(|error| format!("could not create OCR writer: {error}"))?;
|
||
writer
|
||
.WriteBytes(&png)
|
||
.map_err(|error| format!("could not write OCR raster: {error}"))?;
|
||
writer
|
||
.StoreAsync()
|
||
.and_then(|op| op.get())
|
||
.map_err(|error| format!("could not store OCR raster: {error}"))?;
|
||
writer
|
||
.FlushAsync()
|
||
.and_then(|op| op.get())
|
||
.map_err(|error| format!("could not flush OCR raster: {error}"))?;
|
||
let _ = writer.DetachStream();
|
||
stream
|
||
.Seek(0)
|
||
.map_err(|error| format!("could not rewind OCR raster: {error}"))?;
|
||
let decoder = BitmapDecoder::CreateAsync(&stream)
|
||
.and_then(|op| op.get())
|
||
.map_err(|error| format!("could not decode OCR raster: {error}"))?;
|
||
let decoded = decoder
|
||
.GetSoftwareBitmapAsync()
|
||
.and_then(|op| op.get())
|
||
.map_err(|error| format!("could not prepare OCR bitmap: {error}"))?;
|
||
SoftwareBitmap::Convert(&decoded, BitmapPixelFormat::Bgra8)
|
||
.map_err(|error| format!("could not convert OCR bitmap: {error}"))
|
||
}
|
||
|
||
pub(crate) fn ocr_lines(image: &DynamicImage) -> Result<Vec<OcrLine>, String> {
|
||
let _winrt = initialize_winrt()?;
|
||
let bitmap = software_bitmap(image)?;
|
||
let engine = OcrEngine::TryCreateFromUserProfileLanguages()
|
||
.map_err(|error| format!("could not create OCR engine: {error}"))?;
|
||
let result = engine
|
||
.RecognizeAsync(&bitmap)
|
||
.and_then(|op| op.get())
|
||
.map_err(|error| format!("OCR failed: {error}"))?;
|
||
|
||
let mut lines = Vec::new();
|
||
for line in &result
|
||
.Lines()
|
||
.map_err(|error| format!("could not read OCR lines: {error}"))?
|
||
{
|
||
let text = line.Text().map(|t| t.to_string_lossy()).unwrap_or_default();
|
||
if text.trim().is_empty() {
|
||
continue;
|
||
}
|
||
// The line has no box of its own — union its words' boxes, and keep
|
||
// each word's box so a wordmark can be located word-precisely.
|
||
let (mut min_x, mut min_y, mut max_x, mut max_y) =
|
||
(f32::MAX, f32::MAX, f32::MIN, f32::MIN);
|
||
let mut word_boxes = Vec::new();
|
||
if let Ok(words) = line.Words() {
|
||
for word in &words {
|
||
if let Ok(r) = word.BoundingRect() {
|
||
min_x = min_x.min(r.X);
|
||
min_y = min_y.min(r.Y);
|
||
max_x = max_x.max(r.X + r.Width);
|
||
max_y = max_y.max(r.Y + r.Height);
|
||
let wtext = word.Text().map(|t| t.to_string_lossy()).unwrap_or_default();
|
||
word_boxes.push(OcrWord {
|
||
text: wtext,
|
||
rect: Rect {
|
||
x: r.X,
|
||
y: r.Y,
|
||
w: r.Width,
|
||
h: r.Height,
|
||
},
|
||
});
|
||
}
|
||
}
|
||
}
|
||
if word_boxes.is_empty() {
|
||
continue;
|
||
}
|
||
lines.push(OcrLine {
|
||
text: text.trim().to_string(),
|
||
rect: Rect {
|
||
x: min_x,
|
||
y: min_y,
|
||
w: max_x - min_x,
|
||
h: max_y - min_y,
|
||
},
|
||
words: word_boxes,
|
||
});
|
||
}
|
||
Ok(lines)
|
||
}
|
||
}
|