mii/usls
1
0
Fork 0
mirror of https://github.com/mii443/usls.git synced 2025-08-22 15:45:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

Using rayon to accelerate DB post-processing

Browse Source
This commit is contained in:
jamjamjon
2025-01-13 22:42:43 +08:00
parent 8c56d5e2a5
commit f511a475ac
8 changed files with 171 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
examples/db/README.md
View File

@ -1,10 +1,13 @@
## Quick Start
```shell
cargo run -r --example db
cargo run -r -F cuda --example db -- --device cuda --dtype fp16
```
## Results
![](https://github.com/jamjamjon/assets/releases/download/db/demo-paper.png)
![](https://github.com/jamjamjon/assets/releases/download/db/demo-slanted-text-number.png)
![](https://github.com/jamjamjon/assets/releases/download/db/demo-table-en.png)
![](https://github.com/jamjamjon/assets/releases/download/db/demo-table-ch.png)
![](https://github.com/jamjamjon/assets/releases/download/db/demo-sign.png)

56
examples/db/main.rs
View File

@ -4,9 +4,33 @@ use usls::{models::DB, Annotator, DataLoader, Options};
#[derive(argh::FromArgs)]
/// Example
struct Args {
/// model file
#[argh(option)]
model: Option<String>,
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
/// dtype
#[argh(option, default = "String::from(\"auto\")")]
dtype: String,
/// show bboxes
#[argh(option, default = "false")]
show_bboxes: bool,
/// show mbrs
#[argh(option, default = "false")]
show_mbrs: bool,
/// show bboxes confidence
#[argh(option, default = "false")]
show_bboxes_conf: bool,
/// show mbrs confidence
#[argh(option, default = "false")]
show_mbrs_conf: bool,
}
fn main() -> Result<()> {
@ -14,23 +38,26 @@ fn main() -> Result<()> {
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
.init();
let args: Args = argh::from_env();
// build model
let options = Options::ppocr_det_v4_server_ch()
.with_model_device(args.device.as_str().try_into()?)
.commit()?;
let mut model = DB::new(options)?;
let options = match &args.model {
Some(m) => Options::db().with_model_file(m),
None => Options::ppocr_det_v4_ch().with_model_dtype(args.dtype.as_str().try_into()?),
};
let mut model = DB::new(
options
.with_model_device(args.device.as_str().try_into()?)
.commit()?,
)?;
// load image
let x = DataLoader::try_read_batch(&[
"images/table.png",
"images/table1.jpg",
"images/table2.png",
"images/table-ch.jpg",
"images/db.png",
"images/table.png",
"images/table-ch.jpg",
"images/street.jpg",
"images/slanted-text-number.jpg",
])?;
// run
@ -38,12 +65,19 @@ fn main() -> Result<()> {
// annotate
let annotator = Annotator::default()
.without_bboxes(true)
.without_mbrs(true)
.without_bboxes(!args.show_bboxes)
.without_mbrs(!args.show_mbrs)
.without_bboxes_name(true)
.without_mbrs_name(true)
.without_bboxes_conf(!args.show_bboxes_conf)
.without_mbrs_conf(!args.show_mbrs_conf)
.with_polygons_alpha(60)
.with_contours_color([255, 105, 180, 255])
.with_saveout(model.spec());
annotator.annotate(&x, &y);
// summary
model.summary();
Ok(())
}

7
examples/fast/main.rs
View File

@ -41,12 +41,11 @@ fn main() -> Result<()> {
// load image
let x = DataLoader::try_read_batch(&[
"images/table.png",
"images/table1.jpg",
"images/table2.png",
"images/table-ch.jpg",
"images/db.png",
"images/table.png",
"images/table-ch.jpg",
"images/street.jpg",
"images/slanted-text-number.jpg",
])?;
// run

7
examples/slanet/main.rs
View File

@ -11,6 +11,10 @@ struct Args {
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
/// dtype
#[argh(option, default = "String::from(\"auto\")")]
dtype: String,
}
fn main() -> Result<()> {
@ -24,6 +28,7 @@ fn main() -> Result<()> {
// build model
let options = Options::slanet_lcnet_v2_mobile_ch()
.with_model_device(args.device.as_str().try_into()?)
.with_model_dtype(args.dtype.as_str().try_into()?)
.commit()?;
let mut model = SLANet::new(options)?;
@ -32,7 +37,7 @@ fn main() -> Result<()> {
// run
let ys = model.forward(&xs)?;
println!("{:?}", ys);
// println!("{:?}", ys);
// annotate
let annotator = Annotator::default()

10
examples/svtr/main.rs
View File

@ -7,6 +7,10 @@ struct Args {
/// device
#[argh(option, default = "String::from(\"cpu:0\")")]
device: String,
/// dtype
#[argh(option, default = "String::from(\"auto\")")]
dtype: String,
}
fn main() -> Result<()> {
@ -19,9 +23,9 @@ fn main() -> Result<()> {
// build model
let options = Options::ppocr_rec_v4_ch()
// svtr_v2_teacher_ch()
// .with_batch_size(2)
// repsvtr_ch()
.with_model_device(args.device.as_str().try_into()?)
.with_model_dtype(args.dtype.as_str().try_into()?)
.commit()?;
let mut model = SVTR::new(options)?;
@ -37,7 +41,7 @@ fn main() -> Result<()> {
println!("{paths:?}: {:?}", ys)
}
//summary
// summary
model.summary();
Ok(())

2
src/models/db/config.rs
View File

@ -33,8 +33,6 @@ impl crate::Options {
Self::db()
.with_image_mean(&[0.798, 0.785, 0.772])
.with_image_std(&[0.264, 0.2749, 0.287])
// .with_binary_thresh(0.3)
// .with_class_confs(&[0.1])
}
pub fn db_mobilenet_v3_large() -> Self {

191
src/models/db/impl.rs
View File

@ -2,8 +2,9 @@ use aksr::Builder;
use anyhow::Result;
use image::DynamicImage;
use ndarray::Axis;
use rayon::prelude::*;
use crate::{elapsed, DynConf, Engine, Mbr, Ops, Options, Polygon, Processor, Ts, Xs, Ys, Y};
use crate::{elapsed, Bbox, DynConf, Engine, Mbr, Ops, Options, Polygon, Processor, Ts, Xs, Ys, Y};
#[derive(Debug, Builder)]
pub struct DB {
@ -73,102 +74,120 @@ impl DB {
Ok(ys)
}
pub fn summary(&mut self) {
self.ts.summary();
}
pub fn postprocess(&mut self, xs: Xs) -> Result<Ys> {
let mut ys = Vec::new();
for (idx, luma) in xs[0].axis_iter(Axis(0)).enumerate() {
let mut y_bbox = Vec::new();
let mut y_polygons: Vec<Polygon> = Vec::new();
let mut y_mbrs: Vec<Mbr> = Vec::new();
let ys: Vec<Y> = xs[0]
.axis_iter(Axis(0))
.into_par_iter()
.enumerate()
.filter_map(|(idx, luma)| {
// input image
let (image_height, image_width) = self.processor.image0s_size[idx];
// input image
let (image_height, image_width) = self.processor.image0s_size[idx];
// reshape
let ratio = self.processor.scale_factors_hw[idx][0];
let v = luma
.as_slice()?
.par_iter()
.map(|x| {
if x <= &self.binary_thresh {
0u8
} else {
(*x * 255.0) as u8
}
})
.collect::<Vec<_>>();
// reshape
let ratio = self.processor.scale_factors_hw[idx][0];
let v = luma
.into_owned()
.into_raw_vec_and_offset()
.0
.iter()
.map(|x| {
if x <= &self.binary_thresh {
0u8
} else {
(*x * 255.0) as u8
}
})
.collect::<Vec<_>>();
let luma = Ops::resize_luma8_u8(
&v,
self.width as _,
self.height as _,
image_width as _,
image_height as _,
true,
"Bilinear",
)
.ok()?;
let mask_im: image::ImageBuffer<image::Luma<_>, Vec<_>> =
image::ImageBuffer::from_raw(image_width as _, image_height as _, luma)?;
let luma = Ops::resize_luma8_u8(
&v,
self.width as _,
self.height as _,
image_width as _,
image_height as _,
true,
"Bilinear",
)?;
let mask_im: image::ImageBuffer<image::Luma<_>, Vec<_>> =
match image::ImageBuffer::from_raw(image_width as _, image_height as _, luma) {
None => continue,
Some(x) => x,
};
// contours
let contours: Vec<imageproc::contours::Contour<i32>> =
imageproc::contours::find_contours_with_threshold(&mask_im, 1);
// contours
let contours: Vec<imageproc::contours::Contour<i32>> =
imageproc::contours::find_contours_with_threshold(&mask_im, 1);
// loop
let (y_polygons, y_bbox, y_mbrs): (Vec<Polygon>, Vec<Bbox>, Vec<Mbr>) = contours
.par_iter()
.filter_map(|contour| {
if contour.border_type == imageproc::contours::BorderType::Hole
&& contour.points.len() <= 2
{
return None;
}
// loop
for contour in contours.iter() {
if contour.border_type == imageproc::contours::BorderType::Hole
&& contour.points.len() <= 2
{
continue;
}
let polygon = Polygon::default()
.with_points_imageproc(&contour.points)
.with_id(0);
let delta =
polygon.area() * ratio.round() as f64 * self.unclip_ratio as f64
/ polygon.perimeter();
let polygon = Polygon::default().with_points_imageproc(&contour.points);
let delta = polygon.area() * ratio.round() as f64 * self.unclip_ratio as f64
/ polygon.perimeter();
let polygon = polygon
.unclip(delta, image_width as f64, image_height as f64)
.resample(50)
// .simplify(6e-4)
.convex_hull()
.verify();
// TODO: optimize
let polygon = polygon
.unclip(delta, image_width as f64, image_height as f64)
.resample(50)
// .simplify(6e-4)
.convex_hull()
.verify();
polygon.bbox().and_then(|bbox| {
if bbox.height() < self.min_height || bbox.width() < self.min_width {
return None;
}
let confidence = polygon.area() as f32 / bbox.area();
if confidence < self.confs[0] {
return None;
}
let bbox = bbox.with_confidence(confidence).with_id(0);
let mbr = polygon
.mbr()
.map(|mbr| mbr.with_confidence(confidence).with_id(0));
if let Some(bbox) = polygon.bbox() {
if bbox.height() < self.min_height || bbox.width() < self.min_width {
continue;
}
let confidence = polygon.area() as f32 / bbox.area();
if confidence < self.confs[0] {
continue;
}
y_bbox.push(bbox.with_confidence(confidence).with_id(0));
Some((polygon, bbox, mbr))
})
})
.fold(
|| (Vec::new(), Vec::new(), Vec::new()),
|mut acc, (polygon, bbox, mbr)| {
acc.0.push(polygon);
acc.1.push(bbox);
if let Some(mbr) = mbr {
acc.2.push(mbr);
}
acc
},
)
.reduce(
|| (Vec::new(), Vec::new(), Vec::new()),
|mut acc, (polygons, bboxes, mbrs)| {
acc.0.extend(polygons);
acc.1.extend(bboxes);
acc.2.extend(mbrs);
acc
},
);
if let Some(mbr) = polygon.mbr() {
y_mbrs.push(mbr.with_confidence(confidence).with_id(0));
}
y_polygons.push(polygon.with_id(0));
} else {
continue;
}
}
ys.push(
Y::default()
.with_bboxes(&y_bbox)
.with_polygons(&y_polygons)
.with_mbrs(&y_mbrs),
);
}
Some(
Y::default()
.with_bboxes(&y_bbox)
.with_polygons(&y_polygons)
.with_mbrs(&y_mbrs),
)
})
.collect();
Ok(ys.into())
}
pub fn summary(&mut self) {
self.ts.summary();
}
}

1
src/models/svtr/config.rs
View File

@ -4,6 +4,7 @@ impl crate::Options {
Self::default()
.with_model_name("svtr")
.with_model_ixx(0, 0, (1, 1, 8).into())
.with_model_ixx(0, 1, 3.into())
.with_model_ixx(0, 2, 48.into())
.with_model_ixx(0, 3, (320, 960, 1600).into())
.with_resize_mode(crate::ResizeMode::FitHeight)