|
use anyhow::Result; |
|
use clap::Parser; |
|
|
|
use usls::{ |
|
models::YOLO, Annotator, DataLoader, Device, Options, Viewer, Vision, YOLOScale, YOLOTask, |
|
YOLOVersion, COCO_SKELETONS_16, |
|
}; |
|
|
|
#[derive(Parser, Clone)] |
|
#[command(author, version, about, long_about = None)] |
|
pub struct Args { |
|
|
|
#[arg(long)] |
|
pub model: Option<String>, |
|
|
|
|
|
#[arg(long, default_value_t = String::from("../../ultralytics/assets/bus.jpg"))] |
|
pub source: String, |
|
|
|
|
|
#[arg(long, value_enum, default_value_t = YOLOTask::Detect)] |
|
pub task: YOLOTask, |
|
|
|
|
|
#[arg(long, value_enum, default_value_t = YOLOVersion::V8)] |
|
pub ver: YOLOVersion, |
|
|
|
|
|
#[arg(long, value_enum, default_value_t = YOLOScale::N)] |
|
pub scale: YOLOScale, |
|
|
|
|
|
#[arg(long, default_value_t = 1)] |
|
pub batch_size: usize, |
|
|
|
|
|
#[arg(long, default_value_t = 224)] |
|
pub width_min: isize, |
|
|
|
|
|
#[arg(long, default_value_t = 640)] |
|
pub width: isize, |
|
|
|
|
|
#[arg(long, default_value_t = 1024)] |
|
pub width_max: isize, |
|
|
|
|
|
#[arg(long, default_value_t = 224)] |
|
pub height_min: isize, |
|
|
|
|
|
#[arg(long, default_value_t = 640)] |
|
pub height: isize, |
|
|
|
|
|
#[arg(long, default_value_t = 1024)] |
|
pub height_max: isize, |
|
|
|
|
|
#[arg(long, default_value_t = 80)] |
|
pub nc: usize, |
|
|
|
|
|
#[arg(long)] |
|
pub confs: Vec<f32>, |
|
|
|
|
|
#[arg(long)] |
|
pub trt: bool, |
|
|
|
|
|
#[arg(long)] |
|
pub cuda: bool, |
|
|
|
|
|
#[arg(long)] |
|
pub coreml: bool, |
|
|
|
|
|
#[arg(long)] |
|
pub half: bool, |
|
|
|
|
|
#[arg(long, default_value_t = 0)] |
|
pub device_id: usize, |
|
|
|
|
|
#[arg(long)] |
|
pub profile: bool, |
|
|
|
|
|
#[arg(long)] |
|
pub no_contours: bool, |
|
|
|
|
|
#[arg(long)] |
|
pub view: bool, |
|
|
|
|
|
#[arg(long)] |
|
pub nosave: bool, |
|
} |
|
|
|
fn main() -> Result<()> { |
|
let args = Args::parse(); |
|
|
|
|
|
if args.profile { |
|
tracing_subscriber::fmt() |
|
.with_max_level(tracing::Level::INFO) |
|
.init(); |
|
} |
|
|
|
|
|
let path = match &args.model { |
|
None => format!( |
|
"yolo/{}-{}-{}.onnx", |
|
args.ver.name(), |
|
args.scale.name(), |
|
args.task.name() |
|
), |
|
Some(x) => x.to_string(), |
|
}; |
|
|
|
|
|
let saveout = match &args.model { |
|
None => format!( |
|
"{}-{}-{}", |
|
args.ver.name(), |
|
args.scale.name(), |
|
args.task.name() |
|
), |
|
Some(x) => { |
|
let p = std::path::PathBuf::from(&x); |
|
p.file_stem().unwrap().to_str().unwrap().to_string() |
|
} |
|
}; |
|
|
|
|
|
let device = if args.cuda { |
|
Device::Cuda(args.device_id) |
|
} else if args.trt { |
|
Device::Trt(args.device_id) |
|
} else if args.coreml { |
|
Device::CoreML(args.device_id) |
|
} else { |
|
Device::Cpu(args.device_id) |
|
}; |
|
|
|
|
|
let options = Options::new() |
|
.with_model(&path)? |
|
.with_yolo_version(args.ver) |
|
.with_yolo_task(args.task) |
|
.with_device(device) |
|
.with_trt_fp16(args.half) |
|
.with_ixx(0, 0, (1, args.batch_size as _, 4).into()) |
|
.with_ixx(0, 2, (args.height_min, args.height, args.height_max).into()) |
|
.with_ixx(0, 3, (args.width_min, args.width, args.width_max).into()) |
|
.with_confs(if args.confs.is_empty() { |
|
&[0.2, 0.15] |
|
} else { |
|
&args.confs |
|
}) |
|
.with_nc(args.nc) |
|
.with_find_contours(!args.no_contours) |
|
|
|
|
|
|
|
|
|
.with_profile(args.profile); |
|
|
|
|
|
let mut model = YOLO::new(options)?; |
|
|
|
|
|
let dl = DataLoader::new(&args.source)? |
|
.with_batch(model.batch() as _) |
|
.build()?; |
|
|
|
|
|
let annotator = Annotator::default() |
|
.with_skeletons(&COCO_SKELETONS_16) |
|
.without_masks(true) |
|
.with_bboxes_thickness(3) |
|
.with_keypoints_name(false) |
|
.with_saveout_subs(&["YOLO"]) |
|
.with_saveout(&saveout); |
|
|
|
|
|
let mut viewer = if args.view { |
|
Some(Viewer::new().with_delay(5).with_scale(1.).resizable(true)) |
|
} else { |
|
None |
|
}; |
|
|
|
|
|
for (xs, _paths) in dl { |
|
let ys = model.forward(&xs, args.profile)?; |
|
let images_plotted = annotator.plot(&xs, &ys, !args.nosave)?; |
|
|
|
|
|
match &mut viewer { |
|
Some(viewer) => viewer.imshow(&images_plotted)?, |
|
None => continue, |
|
} |
|
|
|
|
|
match &mut viewer { |
|
Some(viewer) => { |
|
if !viewer.is_open() || viewer.is_key_pressed(usls::Key::Escape) { |
|
break; |
|
} |
|
} |
|
None => continue, |
|
} |
|
|
|
|
|
if !args.nosave { |
|
match &mut viewer { |
|
Some(viewer) => viewer.write_batch(&images_plotted)?, |
|
None => continue, |
|
} |
|
} |
|
} |
|
|
|
|
|
if !args.nosave { |
|
if let Some(viewer) = &mut viewer { |
|
viewer.finish_write()?; |
|
} |
|
} |
|
|
|
Ok(()) |
|
} |
|
|