More error types

Author: dec4234

src/backgrounds/mod.rs

@@ -3,28 +3,28 @@ 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;
 use rand::{Rng, thread_rng};
 use rayon::iter::ParallelBridge;
 use rayon::iter::ParallelIterator;
+use crate::generator::error::GenerationError;
 
 pub struct BackgroundLoader {
 	pub backgrounds: Vec<BackgroundImage>,
 }
 
 impl BackgroundLoader {
-	pub fn new<Q: AsRef<Path>>(path: Q) -> Result<BackgroundLoader> {
+	pub fn new<Q: AsRef<Path>>(path: Q) -> Result<BackgroundLoader, GenerationError> {
 		let dir = path.as_ref().to_path_buf();
 		let mut v = vec![]; // mutex for multi-thread access
 
 		let start = Instant::now();
 
 		if !dir.is_dir() {
-			return Err(anyhow!("Path provided is not a directory!"));
+			return Err(GenerationError::NotADirectory);
 		}
 
 		debug!("Loading backgrounds from: {:?}", dir);

src/generator/coco/mod.rs

@@ -59,7 +59,7 @@ impl CocoGenerator {
 		self.image_id
 	}
 
-	pub fn add_annotation(&mut self, image_id: u32, category_id: u32, iscrowd: u8, segmentation: Vec<Vec<f32>>, area: f64, bbox: BoundingBox) {
+	pub fn add_annotation(&mut self, image_id: u32, category_id: u32, iscrowd: u8, segmentation: Vec<Vec<f32>>, area: f64, bbox: BoundingBox) -> u32 {
 		self.file.annotations.push(CocoAnnotation {
 			id: Some(self.annotation_id),
 			image_id,
@@ -70,6 +70,8 @@ impl CocoGenerator {
 			bbox,
 		});
 		self.annotation_id += 1;
+		
+		self.annotation_id
 	}
 }
 

src/generator/config.rs

@@ -0,0 +1,16 @@
+use image::Rgba;
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct TargetGeneratorConfig {
+	pub visualize_bboxes: bool,
+	pub maskover_color: Option<Rgba<u8>>
+}
+
+impl Default for TargetGeneratorConfig {
+	fn default() -> Self {
+		Self {
+			visualize_bboxes: false,
+			maskover_color: None
+		}
+	}
+}

src/generator/error.rs

@@ -0,0 +1,27 @@
+use thiserror::Error;
+
+#[derive(Debug, Error)]
+pub enum GenerationError {
+	#[error("Calculated new sizes provided an invalid size")]
+	SizeError,
+	#[error("No objects were provided to generate")]
+	NoObjects,
+	#[error("Path provided is not a directory")]
+	NotADirectory,
+	#[error("Missing objects.json file")]
+	MissingObjectsJSON,
+	#[error("{0}")]
+	GenericError(String),
+	#[error("Not enough objects available to generate")]
+	NotEnoughObjectsAvailable,
+
+	// conversions
+	#[error("Error parsing integer")]
+	ParseIntError(#[from] std::num::ParseIntError),
+	#[error("Error parsing image")]
+	ImageError(#[from] image::ImageError),
+	#[error("Serde decoding or encoding error")]
+	SerdeError(#[from] serde_json::Error),
+	#[error("IO error occurred while generating target")]
+	IOError(#[from] std::io::Error),
+}

src/generator/mod.rs

@@ -1,31 +1,37 @@
 use std::ops::RangeTo;
 use std::path::{Path, PathBuf};
 use std::time::Instant;
-use anyhow::Result;
 use image::{GenericImage, ImageBuffer, Rgba, RgbaImage};
 use image::imageops::FilterType;
 use imageproc::drawing::draw_text_mut;
 use log::{debug, trace};
-use rand::{Rng, thread_rng};
+use rand::{thread_rng, Rng};
 use simple_logger::SimpleLogger;
 use strum::Display;
 use thiserror::Error;
+use error::GenerationError;
+use util::STANDARD_PPM;
 use crate::backgrounds::BackgroundLoader;
 use crate::generator::coco::{CocoCategoryInfo, CocoGenerator};
-use crate::objects::{ObjectClass, ObjectManager};
+use crate::generator::config::TargetGeneratorConfig;
+use crate::objects::{ObjectClass, ObjectManager, PlacedObject};
 
 pub mod coco;
+pub mod error;
+pub(crate) mod util;
+pub mod config;
 
 pub struct TargetGenerator {
 	output: PathBuf,
 	backgrounds_path: PathBuf,
 	pub object_manager: ObjectManager,
 	background_loader: BackgroundLoader,
-	coco_generator: CocoGenerator
+	coco_generator: CocoGenerator,
+	config: TargetGeneratorConfig,
 }
 
 impl TargetGenerator {
-	pub fn new<Q: AsRef<Path>>(output: Q, background_path: Q, objects_path: Q, annotations_path: Q) -> Result<Self> {
+	pub fn new<Q: AsRef<Path>>(output: Q, background_path: Q, objects_path: Q, annotations_path: Q) -> Result<Self, GenerationError> {
 
 		let mut object_manager = ObjectManager::new(objects_path);
 		object_manager.load_objects()?;
@@ -35,51 +41,67 @@ 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, ObjectClass::categories())
+			coco_generator: CocoGenerator::new(annotations_path, ObjectClass::categories()),
+			config: TargetGeneratorConfig::default(),
 		})
 	}
 
-	pub fn generate_target(&mut self, pixels_per_meter: f32) -> Result<RgbaImage> {
+	pub fn generate_target(&mut self, pixels_per_meter: f32, number_of_objects: u16) -> Result<RgbaImage, GenerationError> {
 		trace!("Beginning to generate a target...");
+		
+		if number_of_objects == 0 {
+			return Err(GenerationError::NoObjects);
+		}
 
 		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)?;
+		let set = self.object_manager.generate_set(number_of_objects as u32)?;
+		let mut placed_objects = vec![];
 
 		// add background image to coco here
-		self.coco_generator.add_image(w, h, background.filename.clone(), background.date_captured.clone());
+		let background_id = 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) = util::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 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]));
+			if self.config.visualize_bboxes {
+				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 {
+			if let Some(color) = self.config.maskover_color {
+				imageproc::drawing::draw_filled_rect_mut(&mut image, imageproc::rect::Rect::at(x as i32, y as i32).of_size(obj_w, obj_h), color);
+			}
+			
+			let bbox = coco::BoundingBox {
 				x,
 				y,
 				width: obj_w,
 				height: obj_h,
+			};
+			
+			// add annotation to coco here
+			let object_id = self.coco_generator.add_annotation(background_id, obj.object_class as u32, 0, vec![], (obj_w * obj_h) as f64, bbox);
+			
+			placed_objects.push(PlacedObject {
+				id: object_id,
+				bounding_box: bbox,
 			});
 		}
 
 		Ok(image)
 	}
 
-	pub fn generate_targets<A: AsRef<Path>>(&self, amount: u32, range_to: RangeTo<u32>, path: A) -> Result<()> {
+	pub fn generate_targets<A: AsRef<Path>>(&self, amount: u32, range_to: RangeTo<u32>, path: A) -> Result<(), GenerationError> {
 		let start = Instant::now(); // start timer
 
 
@@ -89,45 +111,13 @@ impl TargetGenerator {
 		Ok(())
 	}
 
-	pub fn close(&self) {
-		self.coco_generator.save();
+	pub fn generate_new_location_no_collision(&self) -> (u32, u32) {
+		todo!() // TODO: use collision detection to find a new location
 	}
-}
-
-const STANDARD_PPM: f32 = 35.0;
-
-/// Use the real size of an object and the Pixels Per Meter value to calculate the size in 
-/// pixels that it should be in order to be at scale
-fn resize_ratio(object_real_size: f32, pixels_per_meter: f32) -> f32 {
-	debug!("Real size: {}, Pixels per meter: {}", object_real_size, pixels_per_meter);
-	object_real_size * pixels_per_meter
-}
-
-/// Calculate the new sizes of an object in pixels based on the requested Pixel Per Meter value
-/// 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) -> 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;
 
-	if new_height == 0 || new_width == 0 {
-		return Err(GenerationError::SizeError);
+	pub fn close(&self) {
+		self.coco_generator.save();
 	}
-	
-	Ok((new_width, new_height))
-}
-
-#[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]
@@ -136,7 +126,7 @@ pub fn test_generate_target() {
 	SimpleLogger::new().init().unwrap();
 
 	let mut tg = TargetGenerator::new("output", "backgrounds", "objects", "output/annotations.json").unwrap();
-	let b = tg.generate_target(STANDARD_PPM).unwrap();
+	let b = tg.generate_target(STANDARD_PPM, 1).unwrap();
 
 	b.save("output_1.png".to_string()).unwrap();
 	debug!("Saved generated target to output_1.png");

src/generator/util.rs

@@ -0,0 +1,28 @@
+use log::debug;
+use crate::generator::error::GenerationError;
+
+pub const STANDARD_PPM: f32 = 35.0;
+
+/// Use the real size of an object and the Pixels Per Meter value to calculate the size in 
+/// pixels that it should be in order to be at scale
+fn resize_ratio(object_real_size: f32, pixels_per_meter: f32) -> f32 {
+	debug!("Real size: {}, Pixels per meter: {}", object_real_size, pixels_per_meter);
+	object_real_size * pixels_per_meter
+}
+
+/// Calculate the new sizes of an object in pixels based on the requested Pixel Per Meter value
+/// 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
+pub fn new_sizes(object_width: u32, object_height: u32, pixels_per_meter: f32, real_width: f32) -> anyhow::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;
+	
+	if new_height == 0 || new_width == 0 {
+		return Err(GenerationError::SizeError);
+	}
+	
+	Ok((new_width, new_height))
+}