From 12ec2e436adb146c384017ed5c6b5d962dcb7d4c Mon Sep 17 00:00:00 2001 From: wep21 Date: Tue, 16 Jun 2026 03:56:46 +0900 Subject: [PATCH] feat: add ecdetseg Signed-off-by: wep21 --- README.md | 1 + examples/image-segmentation/ecdetseg.rs | 62 +++++++ examples/image-segmentation/main.rs | 11 +- src/dataloader/hub.rs | 8 +- src/models/vision/ecdetseg/config.rs | 71 +++++++ src/models/vision/ecdetseg/impl.rs | 237 ++++++++++++++++++++++++ src/models/vision/ecdetseg/mod.rs | 4 + src/models/vision/mod.rs | 2 + 8 files changed, 392 insertions(+), 4 deletions(-) create mode 100644 examples/image-segmentation/ecdetseg.rs create mode 100644 src/models/vision/ecdetseg/config.rs create mode 100644 src/models/vision/ecdetseg/impl.rs create mode 100644 src/models/vision/ecdetseg/mod.rs diff --git a/README.md b/README.md index 50a0225c..cc20978b 100644 --- a/README.md +++ b/README.md @@ -214,6 +214,7 @@ cargo run -r -F openvino -F ort-load-dynamic --example yolo -- --device openvino | [YOLOE-v8/11-Prompt-Free](https://github.com/THU-MIG/yoloe) | Open-Set Detection And Segmentation | [demo](./examples/image-segmentation/yoloe_prompt_free) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | [YOLOE-26-Prompt-Free](https://github.com/ultralytics/ultralytics) | Open-Set Detection And Segmentation | [demo](./examples/image-segmentation/yoloe_prompt_free) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | [FastSAM](https://github.com/CASIA-IVA-Lab/FastSAM) | Instance Segmentation | [demo](./examples/image-segmentation) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | +| [EdgeCrafter (ECDetSeg)](https://github.com/Intellindust-AI-Lab/EdgeCrafter) | Instance Segmentation | [demo](./examples/image-segmentation) | ❓ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | | [SAM2](https://github.com/facebookresearch/segment-anything-2) | Segment Anything | [demo](./examples/image-segmentation) | ✅ | ❓ | ✅ | ❌ | ❌ | ❌ | ❌ | | [SAM3-Tracker](https://github.com/facebookresearch/segment-anything-3) | Segment Anything | [demo](./examples/image-segmentation) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | [BiRefNet - COD](https://github.com/ZhengPeng7/BiRefNet) | Camouflaged Object Detection | [demo](./examples/birefnet) | ✅ | ❓ | ✅ | ✅ | ✅ | ✅ | ✅ | diff --git a/examples/image-segmentation/ecdetseg.rs b/examples/image-segmentation/ecdetseg.rs new file mode 100644 index 00000000..233f0b03 --- /dev/null +++ b/examples/image-segmentation/ecdetseg.rs @@ -0,0 +1,62 @@ +use anyhow::Result; +use clap::Args; +use usls::{Config, DType, Device, Scale}; + +#[derive(Args, Debug)] +pub struct EcdetsegArgs { + /// Scale: s, m, l, x + #[arg(long, global = true, default_value = "m")] + pub scale: Scale, + + /// Optional local ONNX model path (overrides the released asset) + #[arg(long)] + pub model: Option, + + /// Dtype: fp32, fp16 + #[arg(long, default_value = "fp16")] + pub dtype: DType, + + /// Device: cpu, cuda:0, tensorrt:0, trt-rtx:0, coreml, etc. + #[arg(long, global = true, default_value = "cpu")] + pub device: Device, + + /// Processor device (for pre/post processing) + #[arg(long, global = true, default_value = "cpu")] + pub processor_device: Device, + + /// Batch size + #[arg(long, global = true, default_value_t = 1)] + pub batch: usize, + + /// Min batch size (TensorRT) + #[arg(long, global = true, default_value_t = 1)] + pub min_batch: usize, + + /// Max batch size (TensorRT) + #[arg(long, global = true, default_value_t = 4)] + pub max_batch: usize, + + /// num dry run + #[arg(long, global = true, default_value_t = 3)] + pub num_dry_run: usize, +} + +pub fn config(args: &EcdetsegArgs) -> Result { + let config = match &args.model { + Some(path) => Config::ecdetseg().with_model_file(path), + None => match args.scale { + Scale::S => Config::ecdetseg_s(), + Scale::M => Config::ecdetseg_m(), + Scale::L => Config::ecdetseg_l(), + Scale::X => Config::ecdetseg_x(), + _ => anyhow::bail!("Unsupported scale: {}. Try s, m, l, x.", args.scale), + }, + } + .with_model_dtype(args.dtype) + .with_model_device(args.device) + .with_batch_size_min_opt_max_all(args.min_batch, args.batch, args.max_batch) + .with_num_dry_run_all(args.num_dry_run) + .with_image_processor_device(args.processor_device); + + Ok(config) +} diff --git a/examples/image-segmentation/main.rs b/examples/image-segmentation/main.rs index 7ba51943..0e8a3c64 100644 --- a/examples/image-segmentation/main.rs +++ b/examples/image-segmentation/main.rs @@ -2,11 +2,13 @@ use anyhow::Result; use clap::{Parser, Subcommand}; use usls::{ models::{ - FastSAM, Sam3Prompt, Sam3Tracker, SamPrompt, YOLOEPromptFree, YOLOPv2, RFDETR, SAM, SAM2, + ECDetSeg, FastSAM, Sam3Prompt, Sam3Tracker, SamPrompt, YOLOEPromptFree, YOLOPv2, RFDETR, + SAM, SAM2, }, Annotator, Config, DataLoader, Model, Source, }; +mod ecdetseg; mod fastsam; mod rfdetr; mod sam; @@ -36,6 +38,7 @@ struct Cli { #[derive(Subcommand)] enum Commands { Rfdetr(rfdetr::RfdetrArgs), + Ecdetseg(ecdetseg::EcdetsegArgs), YoloePromptFree(yoloe_prompt_free::YoloePromptFreeArgs), Sam(sam::SamArgs), Sam2(sam2::Sam2Args), @@ -89,6 +92,12 @@ fn main() -> Result<()> { .commit()?; run::(config, &cli.source, &annotator) } + Commands::Ecdetseg(args) => { + let config = ecdetseg::config(args)? + .with_class_confs(&cli.confs) + .commit()?; + run::(config, &cli.source, &annotator) + } Commands::YoloePromptFree(args) => { let config = yoloe_prompt_free::config(args)? .with_class_confs(&cli.confs) diff --git a/src/dataloader/hub.rs b/src/dataloader/hub.rs index 1796818e..fd98967a 100644 --- a/src/dataloader/hub.rs +++ b/src/dataloader/hub.rs @@ -485,9 +485,11 @@ impl Hub { pb = pb.with_message(msg); } - // Create temporary file in system temp directory (more secure and reliable) - let mut temp_file = NamedTempFile::new() - .context("Failed to create temporary download file in system temp directory")?; + let parent_dir = dst_path + .parent() + .context("Invalid download destination: no parent directory found")?; + let mut temp_file = NamedTempFile::new_in(parent_dir) + .context("Failed to create temporary download file next to destination")?; let mut reader = resp.into_body().into_reader(); const BUFFER_SIZE: usize = 64 * 1024; diff --git a/src/models/vision/ecdetseg/config.rs b/src/models/vision/ecdetseg/config.rs new file mode 100644 index 00000000..50bc16bc --- /dev/null +++ b/src/models/vision/ecdetseg/config.rs @@ -0,0 +1,71 @@ +use crate::Task; + +const ECDETSEG_RELEASE: &str = "https://github.com/wep21/assets/releases/download/ecdetseg"; + +/// +/// > # ECDetSeg: EdgeCrafter Instance Segmentation +/// > +/// > Compact ViT-based real-time instance segmentation from EdgeCrafter, built on the +/// > RT-DETR / D-FINE detection family with an additional mask head. +/// > +/// > # Paper & Code +/// > +/// > - **GitHub**: [Intellindust-AI-Lab/EdgeCrafter](https://github.com/Intellindust-AI-Lab/EdgeCrafter) +/// > - **arXiv**: +/// > +/// > # Model Variants +/// > +/// > - **ecdetseg-s**: Small model for 80-class COCO instance segmentation +/// > - **ecdetseg-m**: Medium model for 80-class COCO instance segmentation +/// > - **ecdetseg-l**: Large model for 80-class COCO instance segmentation +/// > - **ecdetseg-x**: Extra large model for 80-class COCO instance segmentation +/// > +/// > # Implemented Features / Tasks +/// > +/// > - [X] **Instance Segmentation**: 80-class COCO instance segmentation +/// > +/// > # Precision / File naming +/// > +/// > Assets are hosted on the `wep21/assets` GitHub releases (tag `ecdetseg`). +/// > FP32 weights use `ecdetseg-{scale}.onnx`; FP16 weights follow the `-fp16.onnx` +/// > convention (`ecdetseg-{scale}-fp16.onnx`) and are selected automatically via +/// > [`crate::Config::with_model_dtype(DType::Fp16)`](crate::Config::with_model_dtype). +/// +/// Model configuration for `ECDetSeg` +/// +impl crate::Config { + /// Base configuration for ECDetSeg models. + /// + /// ECDetSeg shares the RT-DETR / D-FINE detection front-end (same `images` + + /// `orig_target_sizes` inputs and `labels`, `boxes`, `scores` outputs) and adds a + /// fourth `masks` output for instance segmentation. + pub fn ecdetseg() -> Self { + Self::rtdetr() + .with_name("ecdetseg") + .with_task(Task::InstanceSegmentation) + // EdgeCrafter eval stretches to the model size (letterbox degrades small models) + .with_resize_mode_type(crate::ResizeModeType::FitExact) + .with_image_mean([0.485, 0.456, 0.406]) + .with_image_std([0.229, 0.224, 0.225]) + } + + /// Small model for 80-class COCO instance segmentation. + pub fn ecdetseg_s() -> Self { + Self::ecdetseg().with_model_file(format!("{ECDETSEG_RELEASE}/ecdetseg-s.onnx")) + } + + /// Medium model for 80-class COCO instance segmentation. + pub fn ecdetseg_m() -> Self { + Self::ecdetseg().with_model_file(format!("{ECDETSEG_RELEASE}/ecdetseg-m.onnx")) + } + + /// Large model for 80-class COCO instance segmentation. + pub fn ecdetseg_l() -> Self { + Self::ecdetseg().with_model_file(format!("{ECDETSEG_RELEASE}/ecdetseg-l.onnx")) + } + + /// Extra large model for 80-class COCO instance segmentation. + pub fn ecdetseg_x() -> Self { + Self::ecdetseg().with_model_file(format!("{ECDETSEG_RELEASE}/ecdetseg-x.onnx")) + } +} diff --git a/src/models/vision/ecdetseg/impl.rs b/src/models/vision/ecdetseg/impl.rs new file mode 100644 index 00000000..9a2aa029 --- /dev/null +++ b/src/models/vision/ecdetseg/impl.rs @@ -0,0 +1,237 @@ +use aksr::Builder; +use anyhow::Result; +use ndarray::{s, Axis}; +use rayon::prelude::*; + +use crate::{ + inputs, Config, DynConf, Engine, Engines, FromConfig, Hbb, Image, ImageProcessor, Mask, Model, + Module, Ops, ResizeModeType, Xs, X, Y, +}; + +/// ECDetSeg: EdgeCrafter real-time instance segmentation. +/// +/// Detection front-end matches the RT-DETR / D-FINE family (inputs `images` + +/// `orig_target_sizes`, outputs `labels`, `boxes`, `scores`) with an additional +/// `masks` output decoded into per-instance masks like RF-DETR segmentation. +#[derive(Debug, Builder)] +pub struct ECDetSeg { + pub height: usize, + pub width: usize, + pub batch: usize, + pub names: Vec, + pub confs: DynConf, + pub processor: ImageProcessor, + pub spec: String, + pub classes_excluded: Vec, + pub classes_retained: Vec, +} + +impl Model for ECDetSeg { + type Input<'a> = &'a [Image]; + + fn batch(&self) -> usize { + self.batch + } + + fn spec(&self) -> &str { + &self.spec + } + + fn build(mut config: Config) -> Result<(Self, Engines)> { + let engine = Engine::from_config(config.take_module(&Module::Model)?)?; + let (batch, height, width) = ( + engine.batch().opt(), + engine.try_height().unwrap_or(&640.into()).opt(), + engine.try_width().unwrap_or(&640.into()).opt(), + ); + let spec = engine.spec().to_owned(); + let names: Vec = config.inference.class_names; + let confs = DynConf::new_or_default(&config.inference.class_confs, names.len()); + let classes_excluded = config.inference.classes_excluded; + let classes_retained = config.inference.classes_retained; + if !classes_excluded.is_empty() { + tracing::info!("classes_excluded: {classes_excluded:?}"); + } + if !classes_retained.is_empty() { + tracing::info!("classes_retained: {classes_retained:?}"); + } + let processor = ImageProcessor::from_config(config.image_processor)? + .with_image_width(width as _) + .with_image_height(height as _); + + let model = Self { + height, + width, + batch, + spec, + names, + confs, + processor, + classes_excluded, + classes_retained, + }; + + let engines = Engines::from(engine); + Ok((model, engines)) + } + + fn run(&mut self, engines: &mut Engines, images: Self::Input<'_>) -> Result> { + let x1 = crate::perf!("ECDetSeg::preprocess", self.processor.process(images)?); + let x2 = X::from(vec![self.height as f32, self.width as f32]) + .insert_axis(0)? + .repeat(0, self.batch)?; + let ys = crate::perf!( + "ECDetSeg::inference", + engines.run(&Module::Model, inputs![&x1, x2]?)? + ); + crate::perf!("ECDetSeg::postprocess", self.postprocess(&ys)) + } +} + +impl ECDetSeg { + fn postprocess(&self, outputs: &Xs) -> Result> { + let resize_mode = match self.processor.resize_mode_type() { + Some(ResizeModeType::Letterbox) => ResizeModeType::Letterbox, + Some(ResizeModeType::FitAdaptive) => ResizeModeType::FitAdaptive, + Some(ResizeModeType::FitExact) => ResizeModeType::FitExact, + Some(x) => anyhow::bail!("Unsupported resize mode for ECDetSeg postprocess: {x:?}. Supported: FitExact, FitAdaptive, Letterbox"), + _ => anyhow::bail!("No resize mode specified. Supported: FitExact, FitAdaptive, Letterbox"), + }; + + // Detection outputs match the RT-DETR / D-FINE export: labels (i64), + // boxes (xyxy in model-input space), scores (sigmoid'd). + let labels = outputs + .get::(0) + .ok_or_else(|| anyhow::anyhow!("Failed to get labels"))?; + let boxes = outputs + .get::(1) + .ok_or_else(|| anyhow::anyhow!("Failed to get bboxes"))?; + let scores = outputs + .get::(2) + .ok_or_else(|| anyhow::anyhow!("Failed to get scores"))?; + // Optional `masks` output: per-query mask logits at the prototype resolution. + let preds_masks = outputs.get::(3); + + let ys: Vec = labels + .axis_iter(Axis(0)) + .into_par_iter() + .zip(boxes.axis_iter(Axis(0)).into_par_iter()) + .zip(scores.axis_iter(Axis(0)).into_par_iter()) + .enumerate() + .filter_map(|(idx, ((labels, boxes), scores))| { + let info = &self.processor.images_transform_info[idx]; + let (image_height, image_width) = (info.height_src, info.width_src); + let dets: Vec<(Hbb, Option)> = scores + .iter() + .enumerate() + .filter_map(|(i, &score)| { + let class_id = labels[i] as usize; + if score < self.confs[class_id] { + return None; + } + + if !self.classes_excluded.is_empty() + && self.classes_excluded.contains(&class_id) + { + return None; + } + + if !self.classes_retained.is_empty() + && !self.classes_retained.contains(&class_id) + { + return None; + } + + // Map xyxy box from model-input space back to the source image. + let xyxy_raw = boxes.slice(s![i, ..]); + let (x1, y1, x2, y2) = match resize_mode { + ResizeModeType::FitExact => { + let scale_x = image_width as f32 / self.width as f32; + let scale_y = image_height as f32 / self.height as f32; + ( + xyxy_raw[0] * scale_x, + xyxy_raw[1] * scale_y, + xyxy_raw[2] * scale_x, + xyxy_raw[3] * scale_y, + ) + } + ResizeModeType::Letterbox => { + let ratio = info.height_scale; + let pad_w = info.width_pad; + let pad_h = info.height_pad; + ( + (xyxy_raw[0] - pad_w) / ratio, + (xyxy_raw[1] - pad_h) / ratio, + (xyxy_raw[2] - pad_w) / ratio, + (xyxy_raw[3] - pad_h) / ratio, + ) + } + ResizeModeType::FitAdaptive => { + let ratio = info.height_scale; + ( + xyxy_raw[0] / ratio, + xyxy_raw[1] / ratio, + xyxy_raw[2] / ratio, + xyxy_raw[3] / ratio, + ) + } + _ => unreachable!(), + }; + + let x1 = x1.max(0.0).min(image_width as _); + let y1 = y1.max(0.0).min(image_height as _); + let x2 = x2.max(0.0).min(image_width as _); + let y2 = y2.max(0.0).min(image_height as _); + let mut hbb = Hbb::default() + .with_xyxy(x1, y1, x2, y2) + .with_confidence(score) + .with_id(class_id); + if !self.names.is_empty() { + hbb = hbb.with_name(&self.names[class_id]); + } + + // Decode the per-instance mask (raw logits -> binarized at 0). + if let Some(preds_masks) = &preds_masks { + let mask = preds_masks.slice(s![idx, i, .., ..]); + let (mh, mw) = (mask.shape()[0], mask.shape()[1]); + let mask = mask.into_owned().into_raw_vec_and_offset().0; + let mask = Ops::resize_mask_with_mode( + mask, + mw, + mh, + image_width as _, + image_height as _, + self.width as _, + self.height as _, + resize_mode, + info, + crate::ResizeFilter::Bilinear, + ) + .ok()?; + + let mut mask = + Mask::new(&mask, image_width as _, image_height as _).ok()?; + mask = mask.with_id(class_id).with_confidence(score); + if !self.names.is_empty() { + mask = mask.with_name(&self.names[class_id]); + } + Some((hbb, Some(mask))) + } else { + Some((hbb, None)) + } + }) + .collect(); + + let (y_hbbs, y_masks): (Vec<_>, Vec<_>) = dets.into_iter().unzip(); + + Some( + Y::default() + .with_hbbs(&y_hbbs) + .with_masks(&y_masks.into_iter().flatten().collect::>()), + ) + }) + .collect(); + + Ok(ys) + } +} diff --git a/src/models/vision/ecdetseg/mod.rs b/src/models/vision/ecdetseg/mod.rs new file mode 100644 index 00000000..fbd2b752 --- /dev/null +++ b/src/models/vision/ecdetseg/mod.rs @@ -0,0 +1,4 @@ +mod config; +mod r#impl; + +pub use r#impl::*; diff --git a/src/models/vision/mod.rs b/src/models/vision/mod.rs index c35c7202..caebcec6 100644 --- a/src/models/vision/mod.rs +++ b/src/models/vision/mod.rs @@ -35,6 +35,7 @@ mod yolo; mod yolop; // Segmentation +mod ecdetseg; mod fastsam; mod mediapipe_segmenter; mod sam; @@ -106,6 +107,7 @@ pub use depth_pro::*; pub use dinov2::*; pub use dinov3::*; pub use dwpose::*; +pub use ecdetseg::*; pub use fast::*; pub use fastsam::*; pub use fastvit::*;