warden_vision/dataset/
mod.rs

1//! Dataset loaders and exporters: three object-detection annotation
2//! formats ([`coco`], [`yolo`], [`pascal_voc`]) built around the shared
3//! [`DetectionDataset`] shape, plus a folder-per-class classification
4//! layout ([`classification`]) built around [`ClassificationDataset`].
5//!
6//! Every detection loader/exporter reads and writes the same
7//! [`DetectionDataset`], so a dataset loaded from one format can be
8//! re-exported to another with no extra glue.
9//!
10//! Requires the `datasets` feature.
11
12pub mod classification;
13pub mod coco;
14pub mod pascal_voc;
15pub mod yolo;
16
17use std::path::{Path, PathBuf};
18
19use crate::core::Detections;
20
21/// Whether `path`'s extension looks like a common raster image format.
22/// Shared by every loader in this module that walks a directory of
23/// images.
24pub(crate) fn is_image_file(path: &Path) -> bool {
25    matches!(
26        path.extension()
27            .and_then(|e| e.to_str())
28            .map(|e| e.to_ascii_lowercase())
29            .as_deref(),
30        Some("jpg") | Some("jpeg") | Some("png") | Some("bmp") | Some("gif") | Some("tiff")
31    )
32}
33
34/// One image's worth of annotations within a [`DetectionDataset`].
35#[derive(Debug, Clone, PartialEq)]
36pub struct DatasetImage {
37    /// Path to the image file. Loaders resolve this relative to whatever
38    /// images directory they were given; exporters use it only to derive
39    /// the output file's base name (they don't copy image pixels).
40    pub path: PathBuf,
41    /// Image width in pixels, as recorded in (or, for YOLO, read from) the
42    /// source annotation.
43    pub width: u32,
44    /// Image height in pixels.
45    pub height: u32,
46    /// This image's detections, in absolute pixel coordinates.
47    pub detections: Detections,
48}
49
50/// An object detection dataset: a shared class list plus one
51/// [`DatasetImage`] per annotated image.
52///
53/// `Detection::class_id` on every image's detections indexes into
54/// `classes`.
55#[derive(Debug, Clone, Default, PartialEq)]
56pub struct DetectionDataset {
57    /// Class names, indexed by `Detection::class_id`.
58    pub classes: Vec<String>,
59    /// One entry per annotated image.
60    pub images: Vec<DatasetImage>,
61}
62
63impl DetectionDataset {
64    /// Creates a new dataset from an explicit class list and image set.
65    pub fn new(classes: Vec<String>, images: Vec<DatasetImage>) -> Self {
66        Self { classes, images }
67    }
68
69    /// An empty dataset with no classes and no images.
70    pub fn empty() -> Self {
71        Self::default()
72    }
73
74    /// Number of images in the dataset.
75    pub fn len(&self) -> usize {
76        self.images.len()
77    }
78
79    /// Whether the dataset has no images.
80    pub fn is_empty(&self) -> bool {
81        self.images.is_empty()
82    }
83
84    /// Total number of annotated object instances across every image.
85    pub fn num_annotations(&self) -> usize {
86        self.images.iter().map(|image| image.detections.len()).sum()
87    }
88}
89
90/// One image's label within a [`ClassificationDataset`].
91#[derive(Debug, Clone, PartialEq)]
92pub struct ClassificationImage {
93    /// Path to the image file.
94    pub path: PathBuf,
95    /// Index into [`ClassificationDataset::classes`].
96    pub class_id: usize,
97}
98
99/// A single-label image classification dataset: a shared class list plus
100/// one [`ClassificationImage`] per image.
101#[derive(Debug, Clone, Default, PartialEq)]
102pub struct ClassificationDataset {
103    /// Class names, indexed by [`ClassificationImage::class_id`].
104    pub classes: Vec<String>,
105    /// One entry per labeled image.
106    pub images: Vec<ClassificationImage>,
107}
108
109impl ClassificationDataset {
110    /// Creates a new dataset from an explicit class list and image set.
111    pub fn new(classes: Vec<String>, images: Vec<ClassificationImage>) -> Self {
112        Self { classes, images }
113    }
114
115    /// An empty dataset with no classes and no images.
116    pub fn empty() -> Self {
117        Self::default()
118    }
119
120    /// Number of images in the dataset.
121    pub fn len(&self) -> usize {
122        self.images.len()
123    }
124
125    /// Whether the dataset has no images.
126    pub fn is_empty(&self) -> bool {
127        self.images.is_empty()
128    }
129}
warden_vision/dataset/mod.rs

warden_vision/dataset/
mod.rs
