pineapple_core/im/
polygons.rs1use std::fs::File;
5use std::io::{BufWriter, Read};
6use std::path::Path;
7
8use serde::Serialize;
9use serde_json::Value;
10
11use crate::constant::POLYGON_JSON_VALID_KEYS;
12use crate::cv::points::{dedup_points, order_points, resample_points};
13use crate::error::PineappleError;
14use crate::im::boxes::BoundingBoxes;
15use crate::mp::form;
16
17#[derive(Debug, Clone)]
46pub struct Polygons {
47 data: Vec<Vec<[f32; 2]>>,
48 deduped: bool,
49 ordered: bool,
50}
51
52impl Polygons {
53 pub fn new(data: Vec<Vec<[f32; 2]>>) -> Result<Self, PineappleError> {
72 let n = data.len();
73
74 let data: Vec<Vec<[f32; 2]>> = data
75 .into_iter()
76 .filter(|polygon| polygon.len() > 2)
77 .collect();
78
79 if data.len() != n {
80 return Err(PineappleError::PolygonsSizeError);
81 }
82
83 Ok(Self {
84 data,
85 deduped: false,
86 ordered: false,
87 })
88 }
89}
90
91impl Polygons {
94 pub fn open<P: AsRef<Path>>(path: P) -> Result<Polygons, PineappleError> {
107 let extension = path
108 .as_ref()
109 .extension()
110 .and_then(|s| s.to_str())
111 .map(|s| s.to_lowercase());
112
113 if let Some(ext) = extension && ext == "json" {
114 return read_polygons_json(path);
115 }
116
117 Err(PineappleError::PolygonsReadError)
118 }
119
120 pub fn save<P: AsRef<Path>>(&self, path: P) -> Result<(), PineappleError> {
134 let extension = path
135 .as_ref()
136 .extension()
137 .and_then(|s| s.to_str())
138 .map(|s| s.to_lowercase());
139
140 if let Some(ext) = extension && ext == "json" {
141 return write_polygons_json(path, &self.data);
142 }
143
144 Err(PineappleError::PolygonsWriteError)
145 }
146}
147
148impl Polygons {
152 pub fn len(&self) -> usize {
154 self.data.len()
155 }
156
157 pub fn is_empty(&self) -> bool {
159 self.data.len() == 0
160 }
161}
162
163impl Polygons {
168 pub fn as_points(&self) -> &Vec<Vec<[f32; 2]>> {
170 &self.data
171 }
172
173 pub fn to_points(self) -> Vec<Vec<[f32; 2]>> {
175 self.data
176 }
177
178 pub fn to_bounding_boxes(&self) -> Result<BoundingBoxes, PineappleError> {
180 BoundingBoxes::new(
181 self.data
182 .iter()
183 .map(|polygon| {
184 let &[fx, fy] = &polygon[0];
185
186 let mut min_x = fx;
187 let mut min_y = fy;
188 let mut max_x = fx;
189 let mut max_y = fy;
190
191 for &[x, y] in polygon {
192 min_x = min_x.min(x);
193 min_y = min_y.min(y);
194 max_x = max_x.max(x);
195 max_y = max_y.max(y);
196 }
197
198 [min_x, min_y, max_x, max_y]
199 })
200 .collect::<Vec<[f32; 4]>>(),
201 )
202 }
203}
204
205impl Polygons {
210 pub fn dedup_points(&mut self) {
212 if !self.deduped {
213 self.data.iter_mut().for_each(dedup_points);
214 self.deduped = true;
215 }
216 }
217
218 pub fn order_points(&mut self) {
220 if !self.ordered {
221 self.data
222 .iter_mut()
223 .for_each(|polygon| order_points(polygon));
224
225 self.ordered = true;
226 }
227 }
228
229 pub fn resample_points(&mut self, n: usize) {
231 self.dedup_points();
232 self.order_points();
233 self.data
234 .iter_mut()
235 .for_each(|polygon| resample_points(polygon, n));
236 }
237
238 pub fn remove(&mut self, indices: &[usize]) {
240 if indices.is_empty() {
241 return;
242 }
243
244 let mut data: Vec<Vec<[f32; 2]>> = Vec::with_capacity(self.len() - indices.len());
245 let mut indices_iter = indices.iter().peekable();
246 let mut next_remove = indices_iter.next().copied();
247
248 for (idx, polygon) in self.data.iter().enumerate() {
249 if Some(idx) == next_remove {
250 next_remove = indices_iter.next().copied();
251 } else {
252 data.push(polygon.to_vec());
253 }
254 }
255
256 self.data = data;
257 }
258
259 pub fn descriptors(&mut self) -> Vec<[f32; 23]> {
261 if !self.deduped {
262 self.dedup_points();
263 self.deduped = true;
264 }
265
266 if !self.ordered {
267 self.order_points();
268 self.ordered = true;
269 }
270
271 self.data
272 .iter()
273 .map(|points| form::descriptors(points))
274 .collect()
275 }
276}
277
278pub fn read_polygons_json<P: AsRef<Path>>(path: P) -> Result<Polygons, PineappleError> {
282 let mut contents = String::new();
283
284 File::open(path)
285 .map_err(|err| PineappleError::NoFileError(err.to_string()))?
286 .read_to_string(&mut contents)
287 .map_err(|err| PineappleError::NoFileError(err.to_string()))?;
288
289 let data: Value = serde_json::from_str(&contents).map_err(|_| PineappleError::PolygonsReadError)?;
290
291 fn to_f32(value: &Value) -> Result<f32, PineappleError> {
292 if let Some(n) = value.as_f64() {
293 Ok(n as f32)
294 } else if let Some(n) = value.as_u64() {
295 Ok(n as f32)
296 } else if let Some(n) = value.as_i64() {
297 Ok(n as f32)
298 } else {
299 Err(PineappleError::PolygonsReadError)
300 }
301 }
302
303 for key in &POLYGON_JSON_VALID_KEYS {
304 if let Some(polygons) = data.get(key).and_then(|v| v.as_array()) {
305 let polygons: Result<Vec<Vec<[f32; 2]>>, _> = polygons
306 .iter()
307 .filter_map(Value::as_array)
308 .map(|polygon| {
309 polygon
310 .iter()
311 .filter_map(Value::as_array)
312 .map(|p| {
313 if p.len() == 2 {
314 let x = to_f32(&p[0])?;
315 let y = to_f32(&p[1])?;
316 Ok([x, y])
317 } else {
318 Err(PineappleError::PolygonsReadError)
319 }
320 })
321 .collect::<Result<Vec<[f32; 2]>, _>>()
322 })
323 .collect();
324
325 if let Ok(polygons) = polygons {
326 return Polygons::new(polygons);
327 }
328 }
329 }
330
331 Err(PineappleError::PolygonsReadError)
332}
333
334pub fn write_polygons_json<P, T>(path: P, polygons: &[Vec<[T; 2]>]) -> Result<(), PineappleError>
336where
337 P: AsRef<Path>,
338 T: Serialize,
339{
340 let file = File::create(path).map_err(|_| PineappleError::PolygonsWriteError)?;
341 let writer = BufWriter::new(file);
342
343 serde_json::to_writer(writer, &serde_json::json!({ "polygons": polygons }))
344 .map_err(|_| PineappleError::PolygonsWriteError)?;
345
346 Ok(())
347}
348
349#[cfg(test)]
350mod test {
351
352 use super::*;
353
354 const TEST_DATA_JSON: &str = "../data/tests/test_polygons.json";
355
356 #[test]
357 pub fn test_open_json_success() {
358 let polygons = Polygons::open(TEST_DATA_JSON);
359 polygons.clone().unwrap();
360 assert!(polygons.is_ok());
361 }
362
363 #[test]
364 pub fn test_open_json_failure() {
365 let polygons = Polygons::open("does_not_exist/");
366 assert!(polygons.is_err())
367 }
368
369 #[test]
370 pub fn test_open_json_count() {
371 let polygons = Polygons::open(TEST_DATA_JSON).unwrap();
372 assert_eq!(polygons.len(), 2);
373 assert!(polygons.as_points()[0].len() > 2);
374 assert_eq!(polygons.as_points()[0][0].len(), 2);
375 }
376
377 #[test]
378 pub fn test_write_json() {
379 const OUTPUT: &str = "TEST_POLYGONS_WRITE.json";
380
381 let polygons = Polygons::open(TEST_DATA_JSON).unwrap();
382
383 polygons.save(OUTPUT).unwrap();
384
385 let polygons = Polygons::open(OUTPUT).unwrap();
386
387 assert_eq!(polygons.as_points(), polygons.as_points());
388
389 std::fs::remove_file(OUTPUT).unwrap();
390 }
391}