478 lines
13 KiB
Rust
478 lines
13 KiB
Rust
use qtrack::scene::{Participant, Scene};
|
|
use qtrack::Detection;
|
|
use opencv::{
|
|
core::{self, Mat},
|
|
highgui,
|
|
prelude::*,
|
|
videoio,
|
|
};
|
|
use nalgebra as na;
|
|
use clap::{Clap, AppSettings};
|
|
|
|
pub struct VideoWriter {
|
|
writer: Option<opencv::videoio::VideoWriter>,
|
|
size: Option<(i32, i32)>,
|
|
out_file: String
|
|
}
|
|
|
|
impl VideoWriter {
|
|
pub fn new<S: ToString>(out_file: S) -> Self {
|
|
Self {
|
|
writer: None,
|
|
size: None,
|
|
out_file: out_file.to_string()
|
|
}
|
|
}
|
|
|
|
pub fn release(&mut self) {
|
|
if let Some(mut w) = self.writer.take() {
|
|
w.release().unwrap();
|
|
}
|
|
}
|
|
|
|
fn reinit(&mut self, size: (i32, i32)) {
|
|
println!("reinit {:?}", size);
|
|
self.release();
|
|
|
|
self.size = Some(size);
|
|
self.writer = Some(opencv::videoio::VideoWriter::new(
|
|
&self.out_file,
|
|
opencv::videoio::VideoWriter::fourcc(b'X' as _, b'V' as _, b'I' as _, b'D' as _).unwrap(),
|
|
24.0,
|
|
opencv::core::Size::new(size.0, size.1),
|
|
true,
|
|
).unwrap());
|
|
}
|
|
|
|
pub fn feed(&mut self, m: &mut opencv::core::Mat) {
|
|
let size = (m.cols(), m.rows());
|
|
|
|
if self.writer.is_none() || self.size != Some(size) {
|
|
self.reinit(size);
|
|
}
|
|
|
|
self.writer.as_mut().unwrap().write(m).unwrap();
|
|
}
|
|
}
|
|
|
|
pub const NAMES: [&'static str; 8] = [
|
|
"person",
|
|
"bicycle",
|
|
"car",
|
|
"motorbike",
|
|
"aeroplane",
|
|
"bus",
|
|
"train",
|
|
"truck",
|
|
];
|
|
|
|
fn draw_pred(frame: &mut Mat, det: &Detection) -> opencv::Result<()> {
|
|
let rect = core::Rect::new(
|
|
(det.x - 2.0) as i32,
|
|
(det.y - 2.0) as i32,
|
|
4 as i32,
|
|
4 as i32
|
|
);
|
|
|
|
let color = match det.class {
|
|
0 => core::Scalar::new(200.0, 0.0, 0.0, 128.0),
|
|
1 => core::Scalar::new(255.0, 255.0, 0.0, 128.0),
|
|
2 => core::Scalar::new(0.0, 200.0, 200.0, 128.0),
|
|
3 => core::Scalar::new(200.0, 0.0, 200.0, 128.0),
|
|
4 => core::Scalar::new(100.0, 100.0, 0.0, 128.0),
|
|
5 => core::Scalar::new(200.0, 200.0, 200.0, 128.0),
|
|
6 => core::Scalar::new(0.0, 0.0, 0.0, 128.0),
|
|
7 => core::Scalar::new(0.0, 0.0, 0.0, 128.0),
|
|
8 => core::Scalar::new(0.0, 0.0, 0.0, 128.0),
|
|
_ => core::Scalar::new(0.0, 0.0, 0.0, 128.0),
|
|
};
|
|
|
|
// Draw a bounding box.
|
|
opencv::imgproc::rectangle(
|
|
frame,
|
|
rect,
|
|
color,
|
|
1,
|
|
opencv::imgproc::LINE_8,
|
|
0
|
|
)?;
|
|
|
|
// core::add_weighted(&mat1, 0.2, &mat2, 0.8, 0.0, frame, mat1.depth()?)?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn draw_track(frame: &mut Mat, p: &Participant, color: core::Scalar, ts_sec: f32, detail: bool, rectangles: bool, track: bool) -> opencv::Result<()> {
|
|
let x = p.object.pos.x;
|
|
let y = p.object.pos.y;
|
|
let pred_dist = p.object.predict_distance(0, ts_sec);
|
|
|
|
opencv::imgproc::circle(
|
|
frame,
|
|
core::Point::new(
|
|
x as i32,
|
|
y as i32,
|
|
),
|
|
4,
|
|
color,
|
|
opencv::imgproc::FILLED,
|
|
opencv::imgproc::LINE_8,
|
|
0
|
|
)?;
|
|
|
|
opencv::imgproc::put_text(
|
|
frame,
|
|
&format!("{}", p.id),
|
|
core::Point::new(x as i32, (y - 10.0) as i32),
|
|
opencv::imgproc::FONT_HERSHEY_SIMPLEX,
|
|
0.4,
|
|
core::Scalar::new(0.0, 0.0, 0.0, 255.0),
|
|
1,
|
|
opencv::imgproc::LINE_AA,
|
|
false
|
|
)?;
|
|
|
|
// let value = grant_tracker::meter::math::linearity(p.object.vel_hist.iter_points());
|
|
let speed = (p.object.vel).round();
|
|
|
|
opencv::imgproc::put_text(
|
|
frame,
|
|
// &format!("{:0.2}", value),
|
|
&format!("{}", speed),
|
|
core::Point::new(x as i32 - 5, y as i32 + 10),
|
|
opencv::imgproc::FONT_HERSHEY_SIMPLEX,
|
|
0.3,
|
|
core::Scalar::new(0.0, 128.0, 255.0, 255.0),
|
|
1,
|
|
opencv::imgproc::LINE_AA,
|
|
false
|
|
)?;
|
|
|
|
if rectangles {
|
|
let det = p.last_detection();
|
|
// // Draw a bounding box.
|
|
opencv::imgproc::rectangle(
|
|
frame,
|
|
core::Rect::new(
|
|
(det.x - det.w * 0.5) as i32,
|
|
(det.y - det.h * 0.5) as i32,
|
|
det.w as i32,
|
|
det.h as i32
|
|
),
|
|
core::Scalar::new(255.0, 255.0, 255.0, 128.0),
|
|
1,
|
|
opencv::imgproc::LINE_8,
|
|
0
|
|
)?;
|
|
}
|
|
|
|
if track {
|
|
for (idx, &(_, pt, _)) in p.object.history.iter().enumerate() {
|
|
opencv::imgproc::circle(
|
|
frame,
|
|
core::Point::new(
|
|
pt.x as i32,
|
|
pt.y as i32,
|
|
),
|
|
1,
|
|
core::Scalar::new(255.0, 128.0, 255.0, 0.0),
|
|
opencv::imgproc::FILLED,
|
|
opencv::imgproc::LINE_8,
|
|
0
|
|
)?;
|
|
}
|
|
}
|
|
|
|
if detail {
|
|
if p.object.predictor.has_quadratic {
|
|
let mut prev: Option<na::Complex<f32>> = None;
|
|
for x in -30 .. (pred_dist as i32 + 120) {
|
|
let x = x as f32;
|
|
let y = p.object.predictor.curvature * x * x;
|
|
let pt = na::Complex::new(x, y) * p.object.predictor.direction + p.object.predictor.extremum;
|
|
|
|
if let Some(prev) = prev {
|
|
opencv::imgproc::line(
|
|
frame,
|
|
core::Point::new(
|
|
prev.re as i32,
|
|
prev.im as i32,
|
|
),
|
|
core::Point::new(
|
|
pt.re as i32,
|
|
pt.im as i32,
|
|
),
|
|
core::Scalar::new(255.0, 255.0, 0.0, 0.0),
|
|
1,
|
|
opencv::imgproc::LINE_8,
|
|
0
|
|
)?;
|
|
}
|
|
|
|
prev = Some(pt);
|
|
}
|
|
|
|
opencv::imgproc::circle(
|
|
frame,
|
|
core::Point::new(
|
|
p.object.predictor.extremum.re as i32,
|
|
p.object.predictor.extremum.im as i32,
|
|
),
|
|
2,
|
|
core::Scalar::new(0.0, 255.0, 255.0, 128.0),
|
|
opencv::imgproc::FILLED,
|
|
opencv::imgproc::LINE_8,
|
|
0
|
|
)?;
|
|
}
|
|
}
|
|
|
|
if p.object.predictor.has_linear {
|
|
let line = p.object.predictor.direction;
|
|
|
|
// opencv::imgproc::line(
|
|
// frame,
|
|
// core::Point::new(
|
|
// (p.object.predictor.mean.re - line.re * 50.0) as i32,
|
|
// (p.object.predictor.mean.im - line.im * 50.0) as i32,
|
|
// ),
|
|
// core::Point::new(
|
|
// (p.object.predictor.mean.re + line.re * 50.0) as i32,
|
|
// (p.object.predictor.mean.im + line.im * 50.0) as i32,
|
|
// ),
|
|
// core::Scalar::new(128.0, 128.0, 0.0, 0.0),
|
|
// 1,
|
|
// opencv::imgproc::LINE_8,
|
|
// 0
|
|
// )?;
|
|
|
|
opencv::imgproc::line(
|
|
frame,
|
|
core::Point::new(
|
|
x as i32,
|
|
y as i32,
|
|
),
|
|
core::Point::new(
|
|
(x + line.re * 30.0) as i32,
|
|
(y + line.im * 30.0) as i32,
|
|
),
|
|
core::Scalar::new(128.0, 128.0, 0.0, 0.0),
|
|
1,
|
|
opencv::imgproc::LINE_8,
|
|
0
|
|
)?;
|
|
|
|
opencv::imgproc::circle(
|
|
frame,
|
|
core::Point::new(
|
|
p.object.predictor.mean.re as i32,
|
|
p.object.predictor.mean.im as i32,
|
|
),
|
|
2,
|
|
core::Scalar::new(255.0, 0.0, 255.0, 128.0),
|
|
opencv::imgproc::FILLED,
|
|
opencv::imgproc::LINE_8,
|
|
0
|
|
)?;
|
|
}
|
|
|
|
let pred = p.object.predict(0, ts_sec);
|
|
opencv::imgproc::circle(
|
|
frame,
|
|
core::Point::new(
|
|
pred.x as i32,
|
|
pred.y as i32,
|
|
),
|
|
2,
|
|
core::Scalar::new(0.0, 0.0, 255.0, 128.0),
|
|
opencv::imgproc::FILLED,
|
|
opencv::imgproc::LINE_8,
|
|
0
|
|
)?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn get_bounding_for_video(name: &str, width: f32, height: f32) -> Vec<na::Point2<f32>> {
|
|
match name {
|
|
_ => {
|
|
let padding = width * 0.01;
|
|
|
|
vec![
|
|
nalgebra::Point2::new(padding, padding),
|
|
nalgebra::Point2::new(width - padding - padding, padding),
|
|
nalgebra::Point2::new(width - padding - padding, height - padding - padding),
|
|
nalgebra::Point2::new(padding, height - padding - padding),
|
|
]
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clap)]
|
|
#[clap(version = "1.0", author = "Andrey T. <andrey@aidev.ru>")]
|
|
#[clap(setting = AppSettings::ColoredHelp)]
|
|
struct Opts {
|
|
#[clap(short, long)]
|
|
skip: Option<u32>,
|
|
|
|
#[clap(short, long)]
|
|
root: String,
|
|
|
|
#[clap(short, long)]
|
|
input: String,
|
|
}
|
|
|
|
fn main() -> Result<(), grant_tracker::error::Error> {
|
|
use std::io::BufRead;
|
|
|
|
let opts: Opts = Opts::parse();
|
|
|
|
let file_stem = std::path::Path::new(&opts.input).file_stem().unwrap().to_string_lossy();
|
|
let dets_file_name = opts.root + &file_stem + ".dets";
|
|
let dets_file = std::fs::File::open(dets_file_name)?;
|
|
|
|
let window = "video capture";
|
|
highgui::named_window(window, 1)?;
|
|
|
|
let mut cam = videoio::VideoCapture::from_file(&opts.input, videoio::CAP_ANY)?; // 0 is the default camera
|
|
|
|
let opened = videoio::VideoCapture::is_opened(&cam)?;
|
|
if !opened {
|
|
panic!("Unable to open default camera!");
|
|
}
|
|
|
|
let width = cam.get(videoio::CAP_PROP_FRAME_WIDTH)? as f32;
|
|
let height = cam.get(videoio::CAP_PROP_FRAME_HEIGHT)? as f32;
|
|
|
|
let mut scene = Scene::new(get_bounding_for_video(&file_stem, width, height));
|
|
|
|
// let mut points_frame = core::Mat::new_rows_cols_with_default(1024, 1024, core::CV_8UC3, core::Scalar::new(0.0, 0.0, 0.0, 0.0)).unwrap();
|
|
let mut frames = [core::Mat::default()];
|
|
let reader = std::io::BufReader::new(dets_file).lines();
|
|
let mut line_reader = reader.into_iter();
|
|
let mut frame_idx = 0usize;
|
|
let mut paused = true;
|
|
let mut skip = false;
|
|
let mut rectangles = false;
|
|
let mut details = true;
|
|
let mut tracks = false;
|
|
let mut ghosts = false;
|
|
let mut frame_skip: usize = 1;
|
|
let mut forwarding = opts.skip.unwrap_or(0);
|
|
|
|
loop {
|
|
match highgui::wait_key(1)? {
|
|
/* esc */ 27 => break,
|
|
/* spc */ 32 => paused = !paused,
|
|
/* n */ 110 => skip = true,
|
|
/* c */ 99 => details = !details,
|
|
/* t */ 116 => tracks = !tracks,
|
|
/* b */ 98 => rectangles = !rectangles,
|
|
/* z */ 122 => ghosts = !ghosts,
|
|
/* rgt */ 83 => forwarding = 100,
|
|
/* 1-9 */ code @ 49 ..= 57 => frame_skip = code as usize - 48,
|
|
-1 => (),
|
|
key => println!("key {}", key),
|
|
}
|
|
|
|
if paused && frame_idx != 0 {
|
|
if !skip {
|
|
continue;
|
|
} else {
|
|
skip = false;
|
|
}
|
|
}
|
|
|
|
frame_idx += 1;
|
|
|
|
if cam.read(&mut frames[0]).is_err() {
|
|
break;
|
|
}
|
|
|
|
let (fwidth, fheight) = (frames[0].cols(), frames[0].rows());
|
|
|
|
if fwidth == 0 || fheight == 0 {
|
|
break;
|
|
}
|
|
|
|
let ts_sec =(cam.get(videoio::CAP_PROP_POS_MSEC)? / 1000.0) as f32;
|
|
let frame = &mut frames[0];
|
|
|
|
// opencv::imgproc::rectangle(
|
|
// &mut points_frame,
|
|
// core::Rect::new(
|
|
// 0,
|
|
// 0,
|
|
// 1024,
|
|
// 1024
|
|
// ),
|
|
// core::Scalar::new(255.0, 255.0, 255.0, 0.0),
|
|
// opencv::imgproc::FILLED,
|
|
// opencv::imgproc::LINE_8,
|
|
// 0
|
|
// )?;
|
|
|
|
let xx: Vec<Detection> = match line_reader.next() {
|
|
Some(x) => {
|
|
let x = x.unwrap();
|
|
|
|
if let Some(idx) = x.find(':') {
|
|
let (_, vector) = x.split_at(idx + 1);
|
|
serde_json::from_str(vector).unwrap()
|
|
} else {
|
|
Vec::new()
|
|
}
|
|
},
|
|
None => break,
|
|
};
|
|
|
|
if frame_idx % frame_skip == 0 {
|
|
// println!("frame #{}", frame_idx);
|
|
let mapping = scene.map_detections(ts_sec, &xx);
|
|
scene.update(mapping);
|
|
}
|
|
|
|
if forwarding > 0 {
|
|
forwarding -= 1;
|
|
continue;
|
|
}
|
|
|
|
for t in &scene.tracks {
|
|
if t.id > 0 {
|
|
draw_track(frame, t, core::Scalar::new(0.0, 255.0, 0.0, 0.0), ts_sec, details, rectangles, tracks)?;
|
|
} else if ghosts {
|
|
draw_track(frame, t, core::Scalar::new(128.0, 128.0, 128.0, 0.0), ts_sec, false, false, false)?;
|
|
}
|
|
}
|
|
|
|
/*
|
|
for t in &scene.peds {
|
|
if t.id > 0 {
|
|
draw_track(frame, t, core::Scalar::new(0.0, 0.0, 255.0, 0.0), ts_sec, false, false, false)?;
|
|
}
|
|
}
|
|
*/
|
|
|
|
for det in &xx {
|
|
draw_pred(frame, det).unwrap();
|
|
}
|
|
|
|
opencv::imgproc::put_text(
|
|
frame,
|
|
&format!("#{}", frame_idx),
|
|
core::Point::new(10 as i32, 30 as i32),
|
|
opencv::imgproc::FONT_HERSHEY_SIMPLEX,
|
|
0.9,
|
|
core::Scalar::new(255.0, 255.0, 0.0, 255.0),
|
|
1,
|
|
opencv::imgproc::LINE_AA,
|
|
false
|
|
)?;
|
|
|
|
highgui::imshow(window, frame)?;
|
|
// highgui::imshow(points, &points_frame)?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|