Skip to content

Commit

Permalink
Fix yolo-seg bug
Browse files Browse the repository at this point in the history
  • Loading branch information
@jamjamjon
jamjamjon committed May 12, 2024
1 parent 76e4aab commit 862415a
Show file tree
Hide file tree
Showing 4 changed files with 60 additions and 55 deletions.
6 changes: 3 additions & 3 deletions examples/yolov8/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,11 @@ use usls::{coco, models::YOLO, Annotator, DataLoader, Options};
fn main() -> Result<(), Box<dyn std::error::Error>> {
// build model
let options = Options::default()
.with_model("yolov8m-dyn.onnx")?
// .with_model("yolov8m-dyn.onnx")?
// .with_model("yolov8m-dyn-f16.onnx")?
// .with_model("yolov8m-pose-dyn.onnx")?
// .with_model("yolov8m-cls-dyn.onnx")?
// .with_model("yolov8m-seg-dyn.onnx")?
.with_model("yolov8m-seg-dyn.onnx")?
// .with_model("yolov8m-obb-dyn.onnx")?
// .with_model("yolov8m-oiv7-dyn.onnx")?
// .with_trt(0)
Expand All @@ -20,7 +20,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
.with_confs(&[0.4, 0.15]) // class 0: 0.4, others: 0.15
.with_names2(&coco::KEYPOINTS_NAMES_17)
// .with_dry_run(10)
.with_profile(true);
.with_profile(false);
let mut model = YOLO::new(options)?;

// build dataloader
Expand Down
1 change: 0 additions & 1 deletion src/core/ort_tensor_attr.rs
View file

This file was deleted.

54 changes: 31 additions & 23 deletions src/models/yolo.rs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use anyhow::Result;
use clap::ValueEnum;
use image::DynamicImage;
use image::{DynamicImage, ImageBuffer};
use ndarray::{s, Array, Axis, IxDyn};
use regex::Regex;

Expand Down Expand Up @@ -341,29 +341,35 @@ impl YOLO {
let mask = coefs.dot(&proto).into_shape((nh, nw, 1))?; // (nh, nw, n)

// build image from ndarray
let mask_im = ops::build_dyn_image_from_raw(
mask.into_raw_vec(),
nw as u32,
nh as u32,
);
let mask: ImageBuffer<image::Luma<_>, Vec<f32>> =
match ImageBuffer::from_raw(
nw as u32,
nh as u32,
mask.clone().into_raw_vec(),
) {
Some(buf) => buf,
None => continue,
};
let mask = image::DynamicImage::from(mask);

// rescale masks
// rescale
let mask_original = ops::descale_mask(
mask_im,
mask,
nw as f32,
nh as f32,
image_width,
image_height,
);

// crop mask with bbox
let mut mask_original = mask_original.into_luma8();

// crop mask
for y in 0..image_height as usize {
for x in 0..image_width as usize {
if x < bbox.xmin() as usize
|| x > bbox.xmax() as usize
|| y < bbox.ymin() as usize
|| y > bbox.ymax() as usize
// || mask_original.get_pixel(x as u32, y as u32).0 < [127]
{
mask_original.put_pixel(
x as u32,
Expand All @@ -375,23 +381,25 @@ impl YOLO {
}

// get masks from image
let mut masks: Vec<Polygon> = Vec::new();
let contours: Vec<imageproc::contours::Contour<i32>> =
imageproc::contours::find_contours_with_threshold(
&mask_original,
1,
0,
);
contours.iter().for_each(|contour| {
if contour.points.len() > 2 {
masks.push(
Polygon::default()
.with_id(bbox.id())
.with_points_imageproc(&contour.points)
.with_name(bbox.name().cloned()),
);
}
});
y_polygons.extend(masks);
let polygon = match contours
.iter()
.map(|x| {
Polygon::default()
.with_id(bbox.id())
.with_points_imageproc(&x.points)
.with_name(bbox.name().cloned())
})
.max_by(|x, y| x.area().total_cmp(&y.area()))
{
None => continue,
Some(x) => x,
};
y_polygons.push(polygon);
}
y = y.with_polygons(&y_polygons);
}
Expand Down
54 changes: 26 additions & 28 deletions src/models/yolop.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,19 +122,19 @@ impl YOLOPv2 {
let mask_da = mask_da.into_luma8();
let mut y_polygons: Vec<Polygon> = Vec::new();
let contours: Vec<imageproc::contours::Contour<i32>> =
imageproc::contours::find_contours_with_threshold(&mask_da, 1);
contours.iter().for_each(|contour| {
if contour.border_type == imageproc::contours::BorderType::Outer
&& contour.points.len() > 2
{
y_polygons.push(
Polygon::default()
.with_id(0)
.with_points_imageproc(&contour.points)
.with_name(Some("Drivable area".to_string())),
);
}
});
imageproc::contours::find_contours_with_threshold(&mask_da, 0);
if let Some(polygon) = contours
.iter()
.map(|x| {
Polygon::default()
.with_id(0)
.with_points_imageproc(&x.points)
.with_name(Some("Drivable area".to_string()))
})
.max_by(|x, y| x.area().total_cmp(&y.area()))
{
y_polygons.push(polygon);
};

// Lane line
let x_ll = x_ll
Expand All @@ -156,21 +156,19 @@ impl YOLOPv2 {
);
let mask_ll = mask_ll.into_luma8();
let contours: Vec<imageproc::contours::Contour<i32>> =
imageproc::contours::find_contours_with_threshold(&mask_ll, 1);
let mut masks: Vec<Polygon> = Vec::new();
contours.iter().for_each(|contour| {
if contour.border_type == imageproc::contours::BorderType::Outer
&& contour.points.len() > 2
{
masks.push(
Polygon::default()
.with_id(1)
.with_points_imageproc(&contour.points)
.with_name(Some("Lane line".to_string())),
);
}
});
y_polygons.extend(masks);
imageproc::contours::find_contours_with_threshold(&mask_ll, 0);
if let Some(polygon) = contours
.iter()
.map(|x| {
Polygon::default()
.with_id(1)
.with_points_imageproc(&x.points)
.with_name(Some("Lane line".to_string()))
})
.max_by(|x, y| x.area().total_cmp(&y.area()))
{
y_polygons.push(polygon);
};

// save
ys.push(
Expand Down

0 comments on commit 862415a

Please sign in to comment.