diff --git a/examples/svtr/README.md b/examples/svtr/README.md index 23b0bf8..6c94700 100644 --- a/examples/svtr/README.md +++ b/examples/svtr/README.md @@ -25,6 +25,7 @@ cargo run -r --example svtr ```shell [Texts] from the background, but also separate text instances which -[Texts] are closely jointed. Some examples are ilustrated in Fig.7. +[Texts] are closely jointed. Some examples are illustrated in Fig.7. [Texts] 你有这么高速运转的机械进入中国,记住我给出的原理 +[Texts] 110022345 ``` diff --git a/examples/svtr/main.rs b/examples/svtr/main.rs index 81a4d67..3ca77b1 100644 --- a/examples/svtr/main.rs +++ b/examples/svtr/main.rs @@ -15,10 +15,13 @@ fn main() -> Result<(), Box> { DataLoader::try_read("./examples/svtr/text1.png")?, DataLoader::try_read("./examples/svtr/text2.png")?, DataLoader::try_read("./examples/svtr/text3.png")?, + DataLoader::try_read("./examples/svtr/text4.png")?, ]; // run - model.run(&xs)?; + for text in model.run(&xs)?.into_iter() { + println!("[Texts] {text}") + } Ok(()) } diff --git a/examples/svtr/text4.png b/examples/svtr/text4.png new file mode 100644 index 0000000..9700cf6 Binary files /dev/null and b/examples/svtr/text4.png differ diff --git a/src/models/svtr.rs b/src/models/svtr.rs index 1bd88ae..0f8f8dd 100644 --- a/src/models/svtr.rs +++ b/src/models/svtr.rs @@ -41,46 +41,53 @@ impl SVTR { }) } - pub fn run(&mut self, xs: &[DynamicImage]) -> Result<()> { + pub fn run(&mut self, xs: &[DynamicImage]) -> Result> { let xs_ = ops::resize_with_fixed_height(xs, self.height.opt as u32, self.width.opt as u32, 0.0)?; let xs_ = ops::normalize(xs_, 0.0, 255.0); let ys: Vec> = self.engine.run(&[xs_])?; let ys = ys[0].to_owned(); - self.postprocess(&ys)?; - Ok(()) + + self.postprocess(&ys) } - pub fn postprocess(&self, xs: &Array) -> Result<()> { - for batch in xs.axis_iter(Axis(0)) { - let mut texts: Vec = Vec::new(); - for (i, seq) in batch.axis_iter(Axis(0)).enumerate() { - let (id, &confidence) = seq - .into_iter() - .enumerate() - .reduce(|max, x| if x.1 > max.1 { x } else { max }) - .unwrap(); - if id == 0 || confidence < self.confs[0] { - continue; - } - if i == 0 && id == self.vocab.len() - 1 { - continue; - } - texts.push(self.vocab[id].to_owned()); - } - texts.dedup(); + pub fn postprocess(&self, output: &Array) -> Result> { + let mut texts: Vec = Vec::new(); + for batch in output.axis_iter(Axis(0)) { + let preds = batch + .axis_iter(Axis(0)) + .filter_map(|x| { + x.into_iter() + .enumerate() + .max_by(|(_, x), (_, y)| x.total_cmp(y)) + }) + .collect::>(); - print!("[Texts] "); - if texts.is_empty() { - println!("Nothing detected!"); - } else { - for text in texts.into_iter() { - print!("{text}"); - } - println!(); - } + let text = preds + .iter() + .enumerate() + .fold(Vec::new(), |mut text_ids, (idx, (text_id, &confidence))| { + if *text_id == 0 || confidence < self.confs[0] { + return text_ids; + } + + if idx == 0 || idx == self.vocab.len() - 1 { + return text_ids; + } + + if *text_id != preds[idx - 1].0 { + text_ids.push(*text_id); + } + text_ids + }) + .into_iter() + .map(|idx| self.vocab[idx].to_owned()) + .collect::(); + + texts.push(text); } - Ok(()) + Ok(texts) } } +