Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion examples/apisr/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,8 @@ fn main() -> Result<()> {

// save
for y in ys {
if let Some(images) = y.images() {
let images = y.images();
if !images.is_empty() {
for image in images.iter() {
image.save(format!(
"{}.png",
Expand Down
2 changes: 1 addition & 1 deletion examples/dwpose/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ fn main() -> Result<()> {

// RTMPose infer
for (x, y_det) in xs.iter().zip(ys_det.iter()) {
let y = rtmpose.forward(x, y_det.hbbs())?;
let y = rtmpose.forward(x, Some(y_det.hbbs()))?;

// Annotate
annotator.annotate(x, &y)?.save(format!(
Expand Down
3 changes: 2 additions & 1 deletion examples/fastvlm/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,8 @@ fn main() -> Result<()> {
let ys = model.forward(&xs, &args.prompt)?;

for y in ys.iter() {
if let Some(texts) = y.texts() {
let texts = y.texts();
if !texts.is_empty() {
for text in texts {
println!("\n[User]: {}\n[Assistant]:{:?}", args.prompt, text);
}
Expand Down
6 changes: 4 additions & 2 deletions examples/moondream2/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -62,15 +62,17 @@ fn main() -> Result<()> {
Task::Caption(_) => {
println!("{}:", task);
for (i, y) in ys.iter().enumerate() {
if let Some(texts) = y.texts() {
let texts = y.texts();
if !texts.is_empty() {
println!("Image {}: {:?}\n", i, texts[0]);
}
}
}
Task::Vqa(query) => {
println!("Question: {}", query);
for (i, y) in ys.iter().enumerate() {
if let Some(texts) = y.texts() {
let texts = y.texts();
if !texts.is_empty() {
println!("Image {}: {:?}\n", i, texts[0]);
}
}
Expand Down
3 changes: 2 additions & 1 deletion examples/ram/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,8 @@ fn main() -> Result<()> {

// results
for (x, y) in xs.iter().zip(ys.iter()) {
if let Some(texts) = y.texts() {
let texts = y.texts();
if !texts.is_empty() {
println!("Image: {:?}", x);
println!("Texts: {:?}", texts);
}
Expand Down
2 changes: 1 addition & 1 deletion examples/rtmpose/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,7 @@ fn main() -> Result<()> {

// RTMPose infer
for (x, y_det) in xs.iter().zip(ys_det.iter()) {
let y = rtmpose.forward(x, y_det.hbbs())?;
let y = rtmpose.forward(x, Some(y_det.hbbs()))?;

// Annotate
annotator.annotate(x, &y)?.save(format!(
Expand Down
2 changes: 1 addition & 1 deletion examples/rtmw/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ fn main() -> Result<()> {

// RTMPose infer
for (x, y_det) in xs.iter().zip(ys_det.iter()) {
let y = rtmpose.forward(x, y_det.hbbs())?;
let y = rtmpose.forward(x, Some(y_det.hbbs()))?;

// Annotate
annotator.annotate(x, &y)?.save(format!(
Expand Down
3 changes: 2 additions & 1 deletion examples/smolvlm/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,8 @@ fn main() -> Result<()> {
let ys = model.forward(&xs, &args.prompt)?;

for y in ys.iter() {
if let Some(texts) = y.texts() {
let texts = y.texts();
if !texts.is_empty() {
for text in texts {
println!("[User]: {}\n\n[Assistant]:{:?}", args.prompt, text);
}
Expand Down
3 changes: 2 additions & 1 deletion examples/swin2sr/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,8 @@ fn main() -> Result<()> {

// save
for y in ys {
if let Some(images) = y.images() {
let images = y.images();
if !images.is_empty() {
for image in images.iter() {
image.save(format!(
"{}.png",
Expand Down
3 changes: 2 additions & 1 deletion examples/yolo-sam2/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,8 @@ fn main() -> Result<()> {
// run & annotate
let ys_det = yolo.forward(&xs)?;
for y_det in ys_det.iter() {
if let Some(hbbs) = y_det.hbbs() {
let hbbs = y_det.hbbs();
if !hbbs.is_empty() {
// collect hhbs
let mut prompt = SamPrompt::default();
for hbb in hbbs {
Expand Down
5 changes: 2 additions & 3 deletions examples/yolo-track/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -158,11 +158,10 @@ fn main() -> Result<()> {
}

let ys = model.forward(&xs)?;
// println!("ys: {:?}", ys);

if let Some(hbbs) = ys[0].hbbs() {
let hbbs = ys[0].hbbs();
if !hbbs.is_empty() {
let tracked_hbbs = tracker.update(hbbs)?;
println!("tracked_hbbs: {:?}", tracked_hbbs);

// annnotate
let image_annotated = annotator.annotate(&xs[0], &tracked_hbbs)?;
Expand Down
12 changes: 7 additions & 5 deletions src/models/yolo/impl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ use crate::{
/// A versatile deep learning model that can perform multiple computer vision tasks including:
/// - Object Detection
/// - Instance Segmentation
/// - Keypoint Detection
/// - Keypoint Detection
/// - Image Classification
/// - Oriented Object Detection
#[derive(Debug, Builder)]
Expand Down Expand Up @@ -728,8 +728,9 @@ impl YOLO {
// KeypointsDetection
if let Some(pred_kpts) = slice_kpts {
let kpt_step = self.layout.kpt_step().unwrap_or(3);
if let Some(hbbs) = y.hbbs() {
let y_kpts = hbbs
if !y.hbbs().is_empty() {
let y_kpts = y
.hbbs()
.into_par_iter()
.filter_map(|hbb| {
let uid = hbb.uid();
Expand Down Expand Up @@ -768,7 +769,7 @@ impl YOLO {

// InstanceSegmentation
if let Some(coefs) = slice_coefs {
if let Some(hbbs) = y.hbbs() {
if !y.hbbs().is_empty() {
let protos = &xs[1];
let proto = protos.slice(s![idx, .., .., ..]);
let proto_shape = proto.shape();
Expand All @@ -778,7 +779,8 @@ impl YOLO {
let proto_vec = proto.to_flat_vec::<f32>().ok()?;
let proto_tensor = Tensor::from_vec(proto_vec, (nm, mh * mw)).ok()?;

let y_masks = hbbs
let y_masks = y
.hbbs()
.into_par_iter()
.filter_map(|hbb| {
let coefs_slice = coefs.slice(s![hbb.uid(), ..]);
Expand Down
140 changes: 40 additions & 100 deletions src/results/y.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,123 +10,63 @@ use crate::{Hbb, Image, Keypoint, Mask, Obb, Polygon, Prob, Text};
///
#[derive(Builder, Clone, Default)]
pub struct Y {
texts: Option<Vec<Text>>,
probs: Option<Vec<Prob>>,
keypoints: Option<Vec<Keypoint>>,
keypointss: Option<Vec<Vec<Keypoint>>>,
hbbs: Option<Vec<Hbb>>,
obbs: Option<Vec<Obb>>,
polygons: Option<Vec<Polygon>>,
masks: Option<Vec<Mask>>,
images: Option<Vec<Image>>,
pub texts: Vec<Text>,
pub probs: Vec<Prob>,
pub keypoints: Vec<Keypoint>,
pub keypointss: Vec<Vec<Keypoint>>,
pub hbbs: Vec<Hbb>,
pub obbs: Vec<Obb>,
pub polygons: Vec<Polygon>,
pub masks: Vec<Mask>,
pub images: Vec<Image>,
}

impl std::fmt::Debug for Y {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let mut f = f.debug_struct("Y");
if let Some(xs) = &self.texts {
if !xs.is_empty() {
f.field("Texts", &xs);
}
let mut s = f.debug_struct("Y");

if !self.texts.is_empty() {
s.field("Texts", &self.texts);
}
if let Some(xs) = &self.probs {
f.field("Probs", &xs);
if !self.probs.is_empty() {
s.field("Probs", &self.probs);
}
if let Some(xs) = &self.hbbs {
if !xs.is_empty() {
f.field("Hbbs", &xs);
}
if !self.hbbs.is_empty() {
s.field("Hbbs", &self.hbbs);
}
if let Some(xs) = &self.obbs {
if !xs.is_empty() {
f.field("Obbs", &xs);
}
if !self.obbs.is_empty() {
s.field("Obbs", &self.obbs);
}
if let Some(xs) = &self.keypoints {
if !xs.is_empty() {
f.field("Kpts", &xs);
}
if !self.keypoints.is_empty() {
s.field("Kpts", &self.keypoints);
}
if let Some(xs) = &self.keypointss {
if !xs.is_empty() {
f.field("Kptss", &xs);
}
if !self.keypointss.is_empty() {
s.field("Kptss", &self.keypointss);
}
if let Some(xs) = &self.polygons {
if !xs.is_empty() {
f.field("Polys", &xs);
}
if !self.polygons.is_empty() {
s.field("Polys", &self.polygons);
}
if let Some(xs) = &self.masks {
if !xs.is_empty() {
f.field("Masks", &xs);
}
if !self.masks.is_empty() {
s.field("Masks", &self.masks);
}
if let Some(xs) = &self.images {
if !xs.is_empty() {
f.field("Images", &xs);
}
if !self.images.is_empty() {
s.field("Images", &self.images);
}
f.finish()

s.finish()
}
}

impl Y {
pub fn is_empty(&self) -> bool {
if let Some(xs) = &self.texts {
if !xs.is_empty() {
return false;
}
}

if let Some(xs) = &self.probs {
if !xs.is_empty() {
return false;
}
}

if let Some(xs) = &self.hbbs {
if !xs.is_empty() {
return false;
}
}

if let Some(xs) = &self.obbs {
if !xs.is_empty() {
return false;
}
}

if let Some(xs) = &self.keypoints {
if !xs.is_empty() {
return false;
}
}

if let Some(xs) = &self.keypointss {
if !xs.is_empty() {
return false;
}
}

if let Some(xs) = &self.polygons {
if !xs.is_empty() {
return false;
}
}

if let Some(xs) = &self.masks {
if !xs.is_empty() {
return false;
}
}

if let Some(xs) = &self.images {
if !xs.is_empty() {
return false;
}
}

true
self.texts.is_empty()
&& self.probs.is_empty()
&& self.hbbs.is_empty()
&& self.obbs.is_empty()
&& self.keypoints.is_empty()
&& self.keypointss.is_empty()
&& self.polygons.is_empty()
&& self.masks.is_empty()
&& self.images.is_empty()
}
}
28 changes: 7 additions & 21 deletions src/viz/drawable/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -61,27 +61,13 @@ pub trait Drawable {

impl Drawable for Y {
fn draw(&self, ctx: &DrawContext, canvas: &mut image::RgbaImage) -> anyhow::Result<()> {
if let Some(masks) = self.masks() {
masks.draw(ctx, canvas)?;
}
if let Some(polygons) = self.polygons() {
polygons.draw(ctx, canvas)?;
}
if let Some(hbbs) = self.hbbs() {
hbbs.draw(ctx, canvas)?;
}
if let Some(obbs) = self.obbs() {
obbs.draw(ctx, canvas)?;
}
if let Some(keypoints) = self.keypoints() {
keypoints.draw(ctx, canvas)?;
}
if let Some(keypointss) = self.keypointss() {
keypointss.draw(ctx, canvas)?;
}
if let Some(probs) = self.probs() {
probs.draw(ctx, canvas)?;
}
self.masks().draw(ctx, canvas)?;
self.polygons().draw(ctx, canvas)?;
self.hbbs().draw(ctx, canvas)?;
self.obbs().draw(ctx, canvas)?;
self.keypoints().draw(ctx, canvas)?;
self.keypointss().draw(ctx, canvas)?;
self.probs().draw(ctx, canvas)?;

Ok(())
}
Expand Down