Coco format generation completed, accurate bounding boxes

Author: dec4234

Cargo.toml

@@ -18,6 +18,7 @@ imageproc = "0.25.0"
 serde = {version = "1.0.210", features = ["derive"]}
 serde_json = "1.0.68"
 strum = {version = "0.26.3", features = ["derive"]}
+chrono = "0.4.38"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 [profile.release]

backgrounds/archive/output.png

@@ -1,8 +1,10 @@
 use std::fs;
+use std::fs::DirEntry;
 use std::path::Path;
 use std::sync::Mutex;
 use std::time::Instant;
 use anyhow::{anyhow, Result};
+use chrono::{DateTime, Local};
 use image::{DynamicImage, RgbaImage};
 use log::{debug, trace, warn};
 use rand::seq::SliceRandom;
@@ -11,7 +13,7 @@ use rayon::iter::ParallelBridge;
 use rayon::iter::ParallelIterator;
 
 pub struct BackgroundLoader {
-	pub backgrounds: Vec<RgbaImage>,
+	pub backgrounds: Vec<BackgroundImage>,
 }
 
 impl BackgroundLoader {
@@ -27,8 +29,8 @@ impl BackgroundLoader {
 
 		debug!("Loading backgrounds from: {:?}", dir);
 
-		// parallelize here because loading images can be slow
-		fs::read_dir(dir)?.for_each(|entry| {
+		// TODO: parallelize here because loading images can be slow
+		fs::read_dir(dir)?.for_each(|entry | {
 			let entry = entry.unwrap();
 			let path = entry.path();
 
@@ -39,32 +41,42 @@ impl BackgroundLoader {
 			let path_name = path.display().to_string();
 
 			if let Ok(img) = image::open(path) {
-				v.push(img);
+				let datetime: DateTime<Local> = entry.metadata().unwrap().created().unwrap().into();
+				
+				let back = BackgroundImage {
+					image: img.to_rgba8(),
+					filename: path_name.clone(),
+					date_captured: datetime.to_string(),
+					id: v.len() as u32,
+				};
+				
+				v.push(back);
 			} else {
 				warn!("Failed to load image: {}", path_name);
 			}
 
 			debug!("Loaded image in {}ms: {}", start.elapsed().as_millis(), path_name);
 		});
-
-		let mut r = vec![];
-		
-		for img in v.iter() {
-			r.push(img.to_rgba8());
-		}
 
 		let x = Ok(Self { // don't ask why we need to make a variable here, it just works
-			backgrounds: r
+			backgrounds: v
 		});
 
 		x
 	}
 
-	pub fn random(&self) -> Option<&RgbaImage> {
+	pub fn random(&self) -> Option<&BackgroundImage> {
 		self.backgrounds.choose(&mut thread_rng())
 	}
 }
 
+pub struct BackgroundImage {
+	pub image: RgbaImage,
+	pub filename: String,
+	pub date_captured: String,
+	pub id: u32,
+}
+
 #[test]
 #[ignore]
 fn test_bg_loader() {

src/backgrounds/mod.rs

@@ -1,5 +1,6 @@
 // https://docs.aws.amazon.com/rekognition/latest/customlabels-dg/md-coco-overview.html
 
+use std::fs::OpenOptions;
 use std::path::{Path, PathBuf};
 use serde::{Deserialize, Serialize};
 
@@ -36,8 +37,9 @@ impl CocoGenerator {
 	}
 
 	pub fn save(&self) {
-		let json = serde_json::to_string_pretty(&self.file).unwrap();
-		std::fs::write(&self.file_path, json).unwrap();
+		let file = OpenOptions::new().write(true).create(true).open(&self.file_path).unwrap();
+		
+		serde_json::to_writer_pretty(file, &self.file).unwrap();
 	}
 
 	/// Add a background image in, then return the image id
@@ -143,5 +145,5 @@ impl CocoCategory {
 }
 
 pub trait CocoCategoryInfo {
-	fn new(&self) -> Vec<CocoCategory>;
+	fn categories() -> Vec<CocoCategory>;
 }

src/generator/coco/mod.rs

@@ -8,9 +8,11 @@ use imageproc::drawing::draw_text_mut;
 use log::{debug, trace};
 use rand::{Rng, thread_rng};
 use simple_logger::SimpleLogger;
+use strum::Display;
+use thiserror::Error;
 use crate::backgrounds::BackgroundLoader;
-use crate::generator::coco::CocoGenerator;
-use crate::objects::ObjectManager;
+use crate::generator::coco::{CocoCategoryInfo, CocoGenerator};
+use crate::objects::{ObjectClass, ObjectManager};
 
 pub mod coco;
 
@@ -33,31 +35,48 @@ impl TargetGenerator {
 			backgrounds_path: background_path.as_ref().to_path_buf(),
 			object_manager,
 			background_loader: BackgroundLoader::new(background_path)?,
-			coco_generator: CocoGenerator::new(annotations_path)
+			coco_generator: CocoGenerator::new(annotations_path, ObjectClass::categories())
 		})
 	}
 
-	pub fn generate_target(&self, altitude: f32, fov: f32, iteration: u32, pixels_per_meter: f32) -> Result<RgbaImage> {
+	pub fn generate_target(&mut self, pixels_per_meter: f32) -> Result<RgbaImage> {
 		trace!("Beginning to generate a target...");
 
-		let mut background = self.background_loader.random().unwrap().clone();
-		let (w, h) = (background.width(), background.height());
+		let background = self.background_loader.random().unwrap();
+		let mut image = background.image.clone();
+		let (w, h) = (image.width(), image.height());
 		let set = self.object_manager.generate_set(1)?;
 
+		// add background image to coco here
+		self.coco_generator.add_image(w, h, background.filename.clone(), background.date_captured.clone());
+		
 		for obj in set {
 			let clone = &obj.dynamic_image.clone();
 			let (obj_w, obj_h) = (obj.dynamic_image.width(), obj.dynamic_image.height());
 			let (x, y) = (thread_rng().gen_range(0..w - obj_w), thread_rng().gen_range(0..h - obj_h));
 			trace!("Placing object at {}, {}", x, y);
 
-			let (obj_w, obj_h) = new_sizes(obj_w, obj_h, pixels_per_meter, obj.object_width_meters);
+			let (obj_w, obj_h) = new_sizes(obj_w, obj_h, pixels_per_meter, obj.object_width_meters)?;
 			debug!("Resizing object to {}x{}", obj_w, obj_h);
 
 			// overlay respects transparent pixels unlike copy_from
-			image::imageops::overlay(&mut background, &clone.resize(obj_w, obj_h, FilterType::Gaussian), x as i64, y as i64);
+			image::imageops::overlay(&mut image, &clone.resize(obj_w, obj_h, FilterType::Gaussian), x as i64, y as i64);
+			
+			// TODO: remove, both are top left
+			//imageproc::drawing::draw_filled_circle_mut(&mut image, (x as i32, y as i32), 4, Rgba([255, 0, 0, 255]));
+			//imageproc::drawing::draw_filled_circle_mut(&mut image, (0i32, 0i32), 8, Rgba([255, 0, 255, 255]));
+			imageproc::drawing::draw_hollow_rect_mut(&mut image, imageproc::rect::Rect::at(x as i32, y as i32).of_size(obj_w, obj_h), Rgba([0, 255, 0, 255]));
+			
+			// add annotation to coco here
+			self.coco_generator.add_annotation(0, obj.object_class as u32, 0, vec![], (obj_w * obj_h) as f64, coco::BoundingBox {
+				x,
+				y,
+				width: obj_w,
+				height: obj_h,
+			});
 		}
 
-		Ok(background)
+		Ok(image)
 	}
 
 	pub fn generate_targets<A: AsRef<Path>>(&self, amount: u32, range_to: RangeTo<u32>, path: A) -> Result<()> {
@@ -69,6 +88,10 @@ impl TargetGenerator {
 
 		Ok(())
 	}
+	
+	pub fn close(&self) {
+		self.coco_generator.save();
+	}
 }
 
 const STANDARD_PPM: f32 = 35.0;
@@ -84,47 +107,39 @@ fn resize_ratio(object_real_size: f32, pixels_per_meter: f32) -> f32 {
 /// 1. Calculate the aspect ratio
 /// 2. Calculate the width of the object in pixels that we expect based on the real width and the Pixels Per Meter value
 /// 3. Calculate the height from this new width using the previously calculated aspect ratio
-fn new_sizes(object_width: u32, object_height: u32, pixels_per_meter: f32, real_width: f32) -> (u32, u32) {
+fn new_sizes(object_width: u32, object_height: u32, pixels_per_meter: f32, real_width: f32) -> Result<(u32, u32), GenerationError> {
 	let (w, h) = (object_width as f32, object_height as f32);
 	let aspect_ratio = w / h;
 	let new_width = resize_ratio(real_width, pixels_per_meter) as u32;
 	let new_height = (new_width as f32 / aspect_ratio) as u32;
 
-	(new_width, new_height)
-}
-
-fn degrees_to_radians(degrees: f32) -> f32 {
-	degrees * std::f32::consts::PI / 180.0
-}
-
-// Calculate the fov in radians based on the image width, height and focal length
-fn calculate_fov(image_width: u32, image_height: u32, focal_length: f32) -> f32 {
-	2.0 * (0.5 * image_width as f32 / focal_length).atan()
-}
-
-/// Calculate the width of the ground in meters based on camera position and field of view
-fn calculate_ground_width(altitude: f32, fov: f32) -> f32 {
-	2.0 * altitude * (fov.to_radians() / 2.0).tan()
-}
-
-fn meters_per_pixel(image_width: u32, ground_width: f32) -> f32 {
-	ground_width / image_width as f32
+	if new_height == 0 || new_width == 0 {
+		return Err(GenerationError::SizeError);
+	}
+	
+	Ok((new_width, new_height))
 }
 
-/// This gives you the expected size of an object in pixels, use this to calculate a ratio between
-/// the expected and actual values. Then use that to scale up/down the object to the expected size
-fn expected_object_size_pixels(real_size: f32, meters_per_pixel: f32) -> f32 {
-	real_size / meters_per_pixel
+#[derive(Debug, Error)]
+pub enum GenerationError {
+	#[error("Serde decoding or encoding error")]
+	SerdeError(#[from] serde_json::Error),
+	#[error("IO error occurred while generating target")]
+	IOError(#[from] std::io::Error),
+	#[error("Calculated new sizes provided an invalid size")]
+	SizeError
 }
 
 #[test]
 #[ignore]
 pub fn test_generate_target() {
 	SimpleLogger::new().init().unwrap();
 
-	let tg = TargetGenerator::new("output", "backgrounds", "objects", "output").unwrap();
-	let b = tg.generate_target(22.8, degrees_to_radians(35.0), 1, STANDARD_PPM).unwrap();
+	let mut tg = TargetGenerator::new("output", "backgrounds", "objects", "output/annotations.json").unwrap();
+	let b = tg.generate_target(STANDARD_PPM).unwrap();
 
 	b.save("output_1.png".to_string()).unwrap();
 	debug!("Saved generated target to output_1.png");
+	
+	tg.close();
 }

src/generator/mod.rs

@@ -90,7 +90,7 @@ impl ObjectManager {
 
 #[derive(Debug, Clone)]
 pub struct Object {
-	object_class: ObjectClass,
+	pub(crate) object_class: ObjectClass,
 	id: u16,
 	pub(crate) dynamic_image: DynamicImage,
 	pub(crate) object_width_meters: f32,
@@ -138,16 +138,23 @@ impl ObjectClass {
 			ObjectClass::TIRE => "tire",
 		}
 	}
+	
+	/// Gets the numerical id of the object class for COCO format category id.
+	/// NOTE: object ids should be sequential and non-duplicative
+	pub fn id(&self) -> u32 {
+		match self {
+			ObjectClass::BICYCLE => 0,
+			ObjectClass::TIRE => 1,
+		}
+	}
 }
 
 impl CocoCategoryInfo for ObjectClass {
-	fn new(&self) -> Vec<CocoCategory> {
+	fn categories() -> Vec<CocoCategory> {
 		let mut v = vec![];
-		let mut id = 0;
 
 		for object_class in ObjectClass::iter() {
-			v.push(CocoCategory::new(id, object_class.prefix().to_string()));
-			id += 1;
+			v.push(CocoCategory::new(object_class.id(), object_class.prefix().to_string()));
 		}
 
 		v
@@ -175,6 +182,7 @@ pub struct PlacedObject {
 	bounding_box: BoundingBox
 }
 
+// Used to generate the starting object mapping file
 #[ignore]
 #[test]
 fn generate_objects_json_file() {
@@ -199,4 +207,14 @@ fn generate_objects_json_file() {
 	let json = serde_json::to_string_pretty(&object_details_file).unwrap();
 
 	std::fs::write("objects/objects.json", json).unwrap();
+}
+
+#[test]
+fn ensure_sequential_no_duplicate_ids() {
+	let mut id = 0;
+	
+	for object_class in ObjectClass::iter() {
+		assert_eq!(object_class.id(), id, "Object class id is not sequential for {:?}", object_class);
+		id += 1;
+	}
 }