From 862415a3dd9c822a83f642715ef60768cf86dca7 Mon Sep 17 00:00:00 2001 From: jamjamjon <505024985@qq.com> Date: Sun, 12 May 2024 17:09:01 +0800 Subject: [PATCH] Fix yolo-seg bug --- examples/yolov8/main.rs | 6 ++--- src/core/ort_tensor_attr.rs | 1 - src/models/yolo.rs | 54 +++++++++++++++++++++---------------- src/models/yolop.rs | 54 ++++++++++++++++++------------------- 4 files changed, 60 insertions(+), 55 deletions(-) delete mode 100644 src/core/ort_tensor_attr.rs diff --git a/examples/yolov8/main.rs b/examples/yolov8/main.rs index 93e7425..91c527f 100644 --- a/examples/yolov8/main.rs +++ b/examples/yolov8/main.rs @@ -3,11 +3,11 @@ use usls::{coco, models::YOLO, Annotator, DataLoader, Options}; fn main() -> Result<(), Box> { // 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) @@ -20,7 +20,7 @@ fn main() -> Result<(), Box> { .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 diff --git a/src/core/ort_tensor_attr.rs b/src/core/ort_tensor_attr.rs deleted file mode 100644 index 2e52d58..0000000 --- a/src/core/ort_tensor_attr.rs +++ /dev/null @@ -1 +0,0 @@ -impl OrtTensorAttr {} diff --git a/src/models/yolo.rs b/src/models/yolo.rs index ab4d2ec..6e56309 100644 --- a/src/models/yolo.rs +++ b/src/models/yolo.rs @@ -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; @@ -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, Vec> = + 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, @@ -375,23 +381,25 @@ impl YOLO { } // get masks from image - let mut masks: Vec = Vec::new(); let contours: Vec> = 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); } diff --git a/src/models/yolop.rs b/src/models/yolop.rs index 51b1bf9..b082e87 100644 --- a/src/models/yolop.rs +++ b/src/models/yolop.rs @@ -122,19 +122,19 @@ impl YOLOPv2 { let mask_da = mask_da.into_luma8(); let mut y_polygons: Vec = Vec::new(); let contours: Vec> = - 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 @@ -156,21 +156,19 @@ impl YOLOPv2 { ); let mask_ll = mask_ll.into_luma8(); let contours: Vec> = - imageproc::contours::find_contours_with_threshold(&mask_ll, 1); - let mut masks: Vec = 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(