Skip to content

Commit

Permalink
update
Browse files Browse the repository at this point in the history
  • Loading branch information
jamjamjon committed Jun 21, 2024
1 parent bbaef09 commit b8863af
Show file tree
Hide file tree
Showing 8 changed files with 231 additions and 114 deletions.
23 changes: 12 additions & 11 deletions src/core/engine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,7 @@ use prost::Message;
use std::collections::HashSet;

use crate::{
home_dir, onnx, ops::make_divisible, BaseTensor, Device, MinOptMax, Options, Ts, CHECK_MARK,
CROSS_MARK,
home_dir, onnx, ops::make_divisible, Device, MinOptMax, Options, Ts, CHECK_MARK, CROSS_MARK, X,
};

/// Ort Tensor Attrs: name, data_type, dims
Expand Down Expand Up @@ -431,21 +430,21 @@ impl OrtEngine {
// Ok(ys)
// }

pub fn execute(&mut self, xs: Vec<BaseTensor>) -> Result<Vec<BaseTensor>> {
pub fn execute(&mut self, xs: Vec<X>) -> Result<Vec<X>> {
// dtype alignment
let mut xs_: Vec<ort::SessionInputValue<'_>> = Vec::new();
let t_pre = std::time::Instant::now();
for (idtype, x) in self.inputs_attrs.dtypes.iter().zip(xs) {
for (idtype, x) in self.inputs_attrs.dtypes.iter().zip(xs.iter()) {
let x_ = match &idtype {
TensorElementType::Float32 => ort::Value::from_array(x.0.view())?.into_dyn(),
TensorElementType::Float32 => ort::Value::from_array(x.data().view())?.into_dyn(),
TensorElementType::Float16 => {
ort::Value::from_array(x.0.mapv(f16::from_f32).view())?.into_dyn()
ort::Value::from_array(x.data().mapv(f16::from_f32).view())?.into_dyn()
}
TensorElementType::Int32 => {
ort::Value::from_array(x.0.mapv(|x_| x_ as i32).view())?.into_dyn()
ort::Value::from_array(x.data().mapv(|x_| x_ as i32).view())?.into_dyn()
}
TensorElementType::Int64 => {
ort::Value::from_array(x.0.mapv(|x_| x_ as i64).view())?.into_dyn()
ort::Value::from_array(x.data().mapv(|x_| x_ as i64).view())?.into_dyn()
}
_ => todo!(),
};
Expand All @@ -461,12 +460,13 @@ impl OrtEngine {
self.ts.add_or_push(1, t_run);

// oputput
let mut ys = Vec::new();
let mut ys = xs;
let t_post = std::time::Instant::now();
for (dtype, name) in self
for ((idx, dtype), name) in self
.outputs_attrs
.dtypes
.iter()
.enumerate()
.zip(self.outputs_attrs.names.iter())
{
let y = &outputs[name.as_str()];
Expand All @@ -485,7 +485,8 @@ impl OrtEngine {
.into_owned(),
_ => todo!(),
};
ys.push(BaseTensor::from(y_));
// ys.push(X::from(y_));
ys[idx].tensor = y_;
}
let t_post = t_post.elapsed();
self.ts.add_or_push(2, t_post);
Expand Down
4 changes: 2 additions & 2 deletions src/core/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,10 +9,10 @@ mod min_opt_max;
pub mod onnx;
pub mod ops;
mod options;
mod tensor;
mod tokenizer_stream;
mod ts;
mod vision;
mod x;

pub use annotator::Annotator;
pub use dataloader::DataLoader;
Expand All @@ -23,9 +23,9 @@ pub use logits_sampler::LogitsSampler;
pub use metric::Metric;
pub use min_opt_max::MinOptMax;
pub use options::Options;
pub use tensor::BaseTensor;
pub use tokenizer_stream::TokenizerStream;
pub use ts::Ts;
pub use vision::Vision;
pub use x::X;

// pub type BaseTensor = ndarray::Array<f32, ndarray::IxDyn>;
47 changes: 0 additions & 47 deletions src/core/tensor.rs

This file was deleted.

15 changes: 7 additions & 8 deletions src/core/vision.rs
Original file line number Diff line number Diff line change
@@ -1,26 +1,25 @@
use crate::{BaseTensor, Options, Y};
use crate::{Options, X, Y};

pub trait Vision: Sized {
type Input; // DynamicImage
// type TensorType; // TODO: make it fixed?

/// Creates a new instance of the model with the given options.
fn new(options: Options) -> anyhow::Result<Self>;

/// Preprocesses the input data.
fn preprocess(&self, xs: &[Self::Input]) -> anyhow::Result<Vec<BaseTensor>>;
fn preprocess(&self, xs: &[Self::Input]) -> anyhow::Result<Vec<X>>;

/// Executes the model on the preprocessed data.
fn inference(&mut self, xs: Vec<BaseTensor>) -> anyhow::Result<Vec<BaseTensor>>;
fn inference(&mut self, xs: Vec<X>) -> anyhow::Result<Vec<X>>;

/// Postprocesses the model's output.
fn postprocess(&self, xs: Vec<BaseTensor>, xs0: &[Self::Input]) -> anyhow::Result<Vec<Y>>;
fn postprocess(&self, xs: Vec<X>) -> anyhow::Result<Vec<Y>>;

/// Executes the full pipeline.
fn run(&mut self, xs: &[Self::Input]) -> anyhow::Result<Vec<Y>> {
let ys = self.preprocess(xs)?;
let ys = self.inference(ys)?;
let ys = self.postprocess(ys, xs)?;
let ys = self.postprocess(ys)?;
Ok(ys)
}

Expand All @@ -35,11 +34,11 @@ pub trait Vision: Sized {
let t_exe = t_exe.elapsed();

let t_post = std::time::Instant::now();
let ys = self.postprocess(ys, xs)?;
let ys = self.postprocess(ys)?;
let t_post = t_post.elapsed();

if profile {
println!("> Pre: {t_pre:?} | Execution: {t_exe:?} | Post: {t_post:?}");
println!("> Preprocess: {t_pre:?} | Execution: {t_exe:?} | Postprocess: {t_post:?}");
}

Ok(ys)
Expand Down
145 changes: 145 additions & 0 deletions src/core/x.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,145 @@
use anyhow::Result;
use image::DynamicImage;
use ndarray::{Array, Dim, IxDyn, IxDynImpl};

use crate::ops;

#[derive(Debug, Clone, Default)]
pub struct X {
pub tensor: Array<f32, IxDyn>,
images: Option<Vec<DynamicImage>>,
}

impl From<Array<f32, IxDyn>> for X {
fn from(tensor: Array<f32, IxDyn>) -> Self {
Self {
tensor,
..Default::default()
}
}
}

impl From<Vec<DynamicImage>> for X {
fn from(images: Vec<DynamicImage>) -> Self {
Self {
images: Some(images.to_vec()),
..Default::default()
}
}
}

impl X {
pub fn load(images: &[DynamicImage]) -> Self {
Self {
images: Some(images.to_vec()),
..Default::default()
}
}

pub fn zeros(shape: &[usize]) -> Self {
Self {
tensor: Array::zeros(Dim(IxDynImpl::from(shape.to_vec()))),
..Default::default()
}
}

pub fn data(&self) -> &Array<f32, IxDyn> {
&self.tensor
}

pub fn images(&self) -> Option<&Vec<DynamicImage>> {
self.images.as_ref()
}

pub fn shape(&self) -> &[usize] {
self.tensor.shape()
}

pub fn dims(&self) -> &[usize] {
self.tensor.shape()
}

pub fn normalize(mut self, min_: f32, max_: f32) -> Self {
self.tensor = (self.tensor - min_) / (max_ - min_);
self
}

pub fn standardize(mut self, mean: &[f32], std: &[f32]) -> Self {
let mean = Array::from_shape_vec((1, mean.len(), 1, 1), mean.to_vec()).unwrap();
let std = Array::from_shape_vec((1, std.len(), 1, 1), std.to_vec()).unwrap();
self.tensor = (self.tensor - mean) / std;
self
}

pub fn letterbox(
mut self,
height: u32,
width: u32,
filter: &str,
bg: Option<u8>,
) -> Result<Self> {
match &self.images {
None => anyhow::bail!("No images found"),
Some(xs) => {
self.tensor = ops::letterbox(xs, height, width, filter, bg)?;
Ok(self)
}
}
}

pub fn resize(mut self, height: u32, width: u32, filter: &str) -> Result<Self> {
match &self.images {
None => anyhow::bail!("No images found"),
Some(xs) => {
self.tensor = ops::resize(xs, height, width, filter)?;
Ok(self)
}
}
}
}

// #[derive(Debug, Clone)]
// pub struct BaseTensor(pub Array<f32, IxDyn>);

// impl From<Array<f32, IxDyn>> for BaseTensor {
// fn from(x: Array<f32, IxDyn>) -> Self {
// Self(x)
// }
// }

// // TODO: from_dynamic_image

// impl BaseTensor {
// pub fn zeros(shape: &[usize]) -> Self {
// Self(Array::zeros(Dim(IxDynImpl::from(shape.to_vec()))))
// }

// pub fn data(&self) -> &Array<f32, IxDyn> {
// &self.0
// }

// pub fn shape(&self) -> &[usize] {
// self.0.shape()
// }

// pub fn dims(&self) -> &[usize] {
// self.0.shape()
// }

// pub fn normalize(mut self, min_: f32, max_: f32) -> Self {
// self.0 = (self.0 - min_) / (max_ - min_);
// self
// }

// pub fn standardize(mut self, mean: &[f32], std: &[f32]) -> Self {
// // let shape = self.0.shape();
// let mean = Array::from_shape_vec((1, mean.len(), 1, 1), mean.to_vec()).unwrap();
// let std = Array::from_shape_vec((1, std.len(), 1, 1), std.to_vec()).unwrap();
// self.0 = (self.0 - mean) / std;
// self
// }

// pub fn into_image() {
// todo!()
// }
// }
39 changes: 19 additions & 20 deletions src/models/db.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use crate::{ops, BaseTensor, DynConf, Mbr, MinOptMax, Options, OrtEngine, Polygon, Vision, Y};
use crate::{ops, DynConf, Mbr, MinOptMax, Options, OrtEngine, Polygon, Vision, X, Y};
use anyhow::Result;
use image::DynamicImage;
use ndarray::Axis;
Expand All @@ -18,7 +18,6 @@ pub struct DB {

impl Vision for DB {
type Input = DynamicImage;
// type TensorType = BaseTensor;

fn new(options: Options) -> Result<Self> {
let mut engine = OrtEngine::new(&options)?;
Expand Down Expand Up @@ -47,30 +46,26 @@ impl Vision for DB {
})
}

fn preprocess(&self, xs: &[Self::Input]) -> Result<Vec<BaseTensor>> {
let xs_ = ops::letterbox(
xs,
self.height.opt as u32,
self.width.opt as u32,
"bilinear",
Some(114),
)?;
// let xs_ = ops::normalize(xs_, 0., 255.);
// let xs_ = ops::standardize(xs_, &[0.485, 0.456, 0.406], &[0.229, 0.224, 0.225]);
// Ok(vec![xs_])
let xs_ = BaseTensor::from(xs_)
fn preprocess(&self, xs: &[Self::Input]) -> Result<Vec<X>> {
let xs = X::load(xs)
.letterbox(
self.height.opt as u32,
self.width.opt as u32,
"bilinear",
Some(114),
)?
.normalize(0., 255.)
.standardize(&[0.485, 0.456, 0.406], &[0.229, 0.224, 0.225]);
Ok(vec![xs_])
Ok(vec![xs])
}

fn inference(&mut self, xs: Vec<BaseTensor>) -> Result<Vec<BaseTensor>> {
fn inference(&mut self, xs: Vec<X>) -> Result<Vec<X>> {
self.engine.execute(xs)
}

fn postprocess(&self, xs: Vec<BaseTensor>, xs0: &[Self::Input]) -> Result<Vec<Y>> {
fn postprocess(&self, xs: Vec<X>) -> Result<Vec<Y>> {
let mut ys = Vec::new();
for (idx, luma) in xs[0].0.axis_iter(Axis(0)).enumerate() {
for (idx, luma) in xs[0].data().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();
Expand All @@ -90,8 +85,12 @@ impl Vision for DB {
ops::build_dyn_image_from_raw(v, self.height() as u32, self.width() as u32);

// input image
let image_width = xs0[idx].width() as f32;
let image_height = xs0[idx].height() as f32;
// let image_width = xs0[idx].width() as f32;
// let image_height = xs0[idx].height() as f32;
let (image_width, image_height) = match xs[0].images() {
None => anyhow::bail!("No images found"),
Some(images) => (images[idx].width() as f32, images[idx].height() as f32),
};

// rescale mask image
let (ratio, w_mask, h_mask) =
Expand Down
Loading

0 comments on commit b8863af

Please sign in to comment.