qtrack/examples/drawer.rs

651 lines
18 KiB
Rust
Raw Permalink Normal View History

2022-03-17 16:26:48 +04:00
use anyhow::Error;
use grant_object_detector::{Detection, YoloDetector, YoloDetectorConfig};
use tracker::deep_sort::{sort, DeepSortConfig, DeepSort, Track};
use std::collections::HashMap;
use serde::{Serialize, Deserialize};
use ndarray::prelude::*;
use opencv::{
dnn,
core::{self, Mat, Scalar, Vector},
highgui,
prelude::*,
videoio,
};
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: &ObjectDetection) -> opencv::Result<()> {
let rect = core::Rect::new(
det.bbox[0] as i32,
det.bbox[1] as i32,
(det.bbox[2] - det.bbox[0]) as i32,
(det.bbox[3] - det.bbox[1]) as i32
);
let color = match det.class {
0 => core::Scalar::new(0.0, 0.0, 255.0, 0.0), // red
1 => core::Scalar::new(0.0, 255.0, 0.0, 0.0), // green
2 => core::Scalar::new(255.0, 0.0, 0.0, 0.0), // blue
3 => core::Scalar::new(0.0, 255.0, 255.0, 0.0), // yellow
5 => core::Scalar::new(255.0, 0.0, 255.0, 0.0), // magenta
7 => core::Scalar::new(255.0, 255.0, 0.0, 0.0), // cyan
_ => core::Scalar::new(0.0, 0.0, 0.0, 0.0), // black
};
// Draw a bounding box.
opencv::imgproc::rectangle(
frame,
rect,
color,
1,
opencv::imgproc::LINE_8,
0
)?;
Ok(())
}
fn draw_track(frame: &mut Mat, track_id: i32, track: &ObjectTrack, info: &ObjectInfo, highlighted: bool) -> opencv::Result<()> {
let color = if highlighted {
core::Scalar::new(0.0, 0.0, 255.0, 128.0)
} else {
core::Scalar::new(0.0, 128.0, 255.0, 128.0)
};
let rect = opencv::core::Rect::new(
track.bbox[0] as i32,
track.bbox[1] as i32,
(track.bbox[2] - track.bbox[0]) as i32,
(track.bbox[3] - track.bbox[1]) as i32
);
// Draw a bounding box.
opencv::imgproc::rectangle(
frame,
rect,
color,
1,
opencv::imgproc::LINE_8,
0
)?;
let label = format!("{}", track_id);
let kind = format!("{}", NAMES[info.class as usize]);
let mut base_line = 0;
let label_size = opencv::imgproc::get_text_size(&label, opencv::imgproc::FONT_HERSHEY_SIMPLEX, 0.4, 1, &mut base_line)?;
let kind_size = opencv::imgproc::get_text_size(&kind, opencv::imgproc::FONT_HERSHEY_SIMPLEX, 0.4, 1, &mut base_line)?;
let label_rect = core::Rect::new(
rect.x,
rect.y,
rect.width,
label_size.height + 4
);
let kind_rect = core::Rect::new(
rect.x + rect.width - kind_size.width,
rect.y,
kind_size.width,
kind_size.height + 4
);
opencv::imgproc::rectangle(frame, label_rect, color, opencv::imgproc::FILLED, opencv::imgproc::LINE_8, 0)?;
// opencv::imgproc::rectangle(frame, kind_rect, core::Scalar::new(0.0, 128.0, 255.0, 255.0), opencv::imgproc::FILLED, opencv::imgproc::LINE_8, 0)?;
opencv::imgproc::put_text(
frame,
&label,
core::Point::new(rect.x, rect.y + label_size.height),
opencv::imgproc::FONT_HERSHEY_SIMPLEX,
0.4,
core::Scalar::new(0.0, 0.0, 0.0, 0.0),
1,
opencv::imgproc::LINE_AA,
false
)?;
opencv::imgproc::put_text(
frame,
&kind,
core::Point::new(kind_rect.x, kind_rect.y + kind_size.height),
opencv::imgproc::FONT_HERSHEY_SIMPLEX,
0.4,
core::Scalar::new(0.0, 0.0, 0.0, 0.0),
1,
opencv::imgproc::LINE_AA,
false
)?;
Ok(())
}
pub enum MarkerResult {
Absent,
Weak,
Strong,
}
pub struct Incedent {
id: i64,
is_confirmed: bool,
objects: Vec<i32>,
door_marker: MarkerResult,
damage_marker: MarkerResult,
people_marker: MarkerResult,
speed_marker: MarkerResult,
alarm_marker: MarkerResult,
time_marker: MarkerResult,
}
#[derive(Serialize, Deserialize)]
pub struct ObjectInfo {
class: i32,
}
#[derive(Serialize, Deserialize)]
pub struct ObjectTrack {
direction: [f32; 2],
bbox: [f32; 4],
}
#[derive(Serialize, Deserialize)]
pub struct ObjectDetection {
confidence: f32,
class: i32,
bbox: [f32; 4],
}
#[derive(Copy, Clone, Serialize, Deserialize)]
pub enum MarkerValue {
NotDetected,
Weak,
Strong,
Extra,
}
impl Default for MarkerValue {
fn default() -> Self {
MarkerValue::NotDetected
}
}
#[derive(Clone, Serialize, Deserialize)]
pub struct IncidentInfo {
id: (i32, i32),
doors: MarkerValue,
alarm: MarkerValue,
damage: MarkerValue,
people: MarkerValue,
speed: MarkerValue,
time: MarkerValue,
}
impl IncidentInfo {
pub fn new(a: i32, b: i32) -> Self {
Self {
id: (a, b),
doors: Default::default(),
alarm: Default::default(),
damage: Default::default(),
people: Default::default(),
speed: Default::default(),
time: Default::default(),
}
}
}
#[derive(Copy, Clone, Serialize, Deserialize)]
pub enum IncidentUpdateKind {
Doors,
Alarm,
Damage,
People,
Speed,
Time,
}
#[derive(Serialize, Deserialize)]
pub struct IncidentUpdate {
id: (i32, i32),
kind: IncidentUpdateKind,
value: MarkerValue,
}
impl IncidentUpdate {
pub fn new(id: (i32, i32), kind: IncidentUpdateKind, value: MarkerValue) -> Self {
Self {
id,
value,
kind,
}
}
}
#[derive(Serialize, Deserialize)]
pub struct IncidentConfirm {
id: (i32, i32),
}
impl IncidentConfirm {
pub fn new(a: i32, b: i32) -> Self {
Self {
id: (a, b),
}
}
}
#[derive(Serialize, Deserialize)]
pub struct IncidentDelete {
id: (i32, i32),
}
impl IncidentDelete {
pub fn new(a: i32, b: i32) -> Self {
Self {
id: (a, b),
}
}
}
#[derive(Serialize, Deserialize)]
pub enum ObjectEvent {
Intersection(i32, i32),
IncidentCreate(IncidentInfo),
IncidentUpdate(IncidentUpdate),
IncidentConfirm(IncidentConfirm),
IncidentDelete(IncidentDelete),
}
#[derive(Default, Serialize, Deserialize)]
pub struct Frame {
tracks: HashMap<i32, ObjectTrack>,
detections: Vec<ObjectDetection>,
events: Vec<ObjectEvent>,
}
#[derive(Serialize, Deserialize)]
pub struct Map {
frames: HashMap<i64, Frame>,
objects: HashMap<i32, ObjectInfo>,
classes: HashMap<i32, String>,
}
impl Map {
pub fn new() -> Self {
Self {
frames: HashMap::new(),
objects: HashMap::new(),
classes: HashMap::new(),
}
}
pub fn object_info(&self, track_id: i32) -> Option<&ObjectInfo> {
self.objects.get(&track_id)
}
pub fn events(&self, frame_id: i64) -> &[ObjectEvent] {
if let Some(frame) = self.frames.get(&frame_id) {
frame.events.as_slice()
} else {
&[]
}
}
pub fn detections(&self, frame_id: i64) -> &[ObjectDetection] {
if let Some(frame) = self.frames.get(&frame_id) {
frame.detections.as_slice()
} else {
&[]
}
}
pub fn tracks(&self, frame_id: i64) -> Vec<(i32, &ObjectTrack)> {
if let Some(frame) = self.frames.get(&frame_id) {
frame.tracks.iter().map(|(&k, v)| (k, v)).collect()
} else {
vec![]
}
}
pub fn add_event(&mut self, frame_id: i64, event: ObjectEvent) {
let frame = self.frames.entry(frame_id).or_insert_with(Default::default);
frame.events.push(event);
}
pub fn add_detection(&mut self, frame_id: i64, det: &Detection) {
let frame = self.frames.entry(frame_id).or_insert_with(Default::default);
self.classes
.entry(det.class)
.or_insert_with(||NAMES[det.class as usize].to_string());
frame.detections.push(ObjectDetection {
confidence: det.confidence,
class: det.class,
bbox: [det.x, det.y, det.w, det.h],
});
}
pub fn add_tarck(&mut self, frame_id: i64, track: &sort::Track) {
let frame = self.frames.entry(frame_id).or_insert_with(Default::default);
let object = self.objects
.entry(track.track_id)
.or_insert_with(||ObjectInfo { class: track.class() });
object.class = track.class();
self.classes
.entry(object.class)
.or_insert_with(||NAMES[object.class as usize].to_string());
frame.tracks.insert(track.track_id, ObjectTrack {
direction: [0.0f32, 0.0f32],
bbox: track.bbox().as_ltrb().into(),
});
}
}
fn main() -> Result<(), anyhow::Error> {
let mut writer = VideoWriter::new("out.avi");
let map_file = std::fs::File::open("out.json")?;
let map: Map = serde_json::from_reader(map_file)?;
let mut frame_idx = 0i64;
let mut cam = videoio::VideoCapture::from_file("/home/andrey/test_video/1-3.mp4", videoio::CAP_ANY)?; // 0 is the default camera
// cam.set(videoio::CAP_PROP_POS_FRAMES, 150.0);
let opened = videoio::VideoCapture::is_opened(&cam)?;
if !opened {
panic!("Unable to open default camera!");
}
let mut frames = [core::Mat::default()];
let mut context: Vec<(bool, IncidentInfo)> = Vec::new();
let mut fq = [0f64; 64];
let mut counter = 0;
loop {
frame_idx += 1;
if cam.read(&mut frames[0]).is_err() {
break;
}
if frames[0].rows() == 0 || frames[0].cols() == 0 {
break;
}
println!("frame #{}:", frame_idx);
let frame = &mut frames[0];
let dets = map.detections(frame_idx);
for d in dets {
if d.class == 0 {
let _ = draw_pred(frame, d);
}
}
let tracks = map.tracks(frame_idx);
for &(track_id, track) in &tracks {
let info = map.object_info(track_id).unwrap();
let highlighted = {
let mut res = false;
for (_, c) in &context {
if c.id.0 == track_id || c.id.1 == track_id {
res = true;
break;
}
}
res
};
let _ = draw_track(frame, track_id, track, info, highlighted);
}
if let Some(&(_, track)) = tracks.iter().find(|&&(track_id, _)| track_id == 18) {
let [left, top, width, height] = track.bbox;
let left = (left as i32).max(0).min(frame.cols());
let top = (top as i32).max(0).min(frame.rows());
let mut width = (width as i32).max(0);
let mut height = (height as i32).max(0);
if left + width > frame.cols() {
width = frame.cols() - left;
}
if top + height > frame.rows() {
height = frame.rows() - top;
}
let rect = core::Rect::new(
left,
top,
width,
height,
);
let roi = Mat::roi(&frame, rect).unwrap().clone();
let tmp = roi.try_into_typed::<core::Vec3b>().unwrap();
let x = tmp.data_typed().unwrap();
let mut fqt = [0f64; 64];
let mut avg = 0f64;
for i in x {
avg += (i.0[0] + i.0[1] + i.0[2]) as f64;
}
avg /= (x.len() * 3) as f64;
for j in 0..64 {
if counter > 0 && counter % (j + 1) == 0 {
fqt[j as usize] += avg;
}
}
for j in 0..64 {
fqt[j as usize] /= (64 - j) as f64 / 64.0;
}
if counter % 300 == 0 {
println!("{:.3?}", fq);
for j in 0..64 {
fq[j as usize] = fqt[j as usize];
}
} else {
for j in 0..64 {
fq[j as usize] += fqt[j as usize];
}
}
counter += 1;
}
let events = map.events(frame_idx);
for event in events {
match event {
ObjectEvent::IncidentCreate(info) => context.push((false, info.clone())),
ObjectEvent::IncidentUpdate(info) => {
for i in &mut context {
if i.1.id == info.id {
match info.kind {
IncidentUpdateKind::Doors => i.1.doors = info.value,
IncidentUpdateKind::Alarm => i.1.alarm = info.value,
IncidentUpdateKind::Damage => i.1.damage = info.value,
IncidentUpdateKind::People => i.1.people = info.value,
IncidentUpdateKind::Speed => i.1.speed = info.value,
IncidentUpdateKind::Time => i.1.time = info.value,
}
}
}
},
ObjectEvent::IncidentConfirm(info) => {
for i in &mut context {
if i.1.id == info.id {
i.0 = true;
}
}
},
_ => (),
}
}
for (idx, ctx) in context.iter().enumerate() {
let color = core::Scalar::new(128.0, 128.0, 255.0, 255.0);
let base_rect = core::Rect::new(
940,
40 + (idx * 200) as i32,
300,
160,
);
let mut smaller_rect = base_rect.clone();
smaller_rect.height = 20;
opencv::imgproc::rectangle(frame, smaller_rect, color, opencv::imgproc::FILLED, opencv::imgproc::LINE_8, 0)?;
opencv::imgproc::rectangle(
frame,
base_rect,
color,
1,
opencv::imgproc::LINE_8,
0
)?;
opencv::imgproc::put_text(
frame,
&format!("{}incident_{}-{}", if ctx.0 {"[confirmed] "} else {""}, ctx.1.id.0, ctx.1.id.1),
core::Point::new(950, 54 + (idx * 200) 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
).unwrap();
draw_cb(frame, 948, 78 + (idx * 200) as i32, ctx.1.doors, "Opened Doors", color_by_marker(ctx.1.doors));
draw_cb(frame, 948, 98 + (idx * 200) as i32, ctx.1.speed, "Car Direction Change", color_by_marker(ctx.1.speed));
draw_cb(frame, 948, 118 + (idx * 200) as i32, ctx.1.alarm, "Car Alarm", color_by_marker(ctx.1.alarm));
draw_cb(frame, 948, 138 + (idx * 200) as i32, ctx.1.damage, "Car Damage", color_by_marker(ctx.1.damage));
draw_cb(frame, 948, 158 + (idx * 200) as i32, ctx.1.people, "People Around", color_by_marker(ctx.1.people));
draw_cb(frame, 948, 178 + (idx * 200) as i32, ctx.1.time, "Time Past", color_by_marker(ctx.1.time));
}
writer.feed(frame);
}
writer.release();
Ok(())
}
fn color_by_marker(val: MarkerValue) -> core::Scalar {
match val {
MarkerValue::NotDetected => core::Scalar::new(0.0, 0.0, 255.0, 255.0),
MarkerValue::Weak => core::Scalar::new(0.0, 255.0, 255.0, 255.0),
MarkerValue::Strong => core::Scalar::new(0.0, 255.0, 0.0, 255.0),
MarkerValue::Extra => core::Scalar::new(255.0, 255.0, 0.0, 255.0)
}
}
fn draw_cb(frame: &mut core::Mat, x: i32, y: i32, val: MarkerValue, title: &str, color: core::Scalar) {
let p = match val {
MarkerValue::NotDetected => " No",
MarkerValue::Weak => "Weak",
MarkerValue::Strong => "Good",
MarkerValue::Extra => "Exel"
};
opencv::imgproc::put_text(
frame,
&format!("{}", p),
core::Point::new(x, y),
opencv::imgproc::FONT_HERSHEY_SIMPLEX,
0.5,
color,
2,
opencv::imgproc::LINE_8,
false
).unwrap();
opencv::imgproc::put_text(
frame,
&format!("{}", title),
core::Point::new(x + 60, y),
opencv::imgproc::FONT_HERSHEY_SIMPLEX,
0.5,
core::Scalar::new(255.0, 255.0, 0.0, 255.0),
1,
opencv::imgproc::LINE_8,
false
).unwrap();
}