Skip to main content

panlabel/ir/
model.rs

1//! Core dataset model for the panlabel intermediate representation.
2//!
3//! This module defines the canonical format-agnostic representation of
4//! object detection datasets. All format-specific readers convert to this
5//! IR, and all writers convert from it.
6
7use serde::{Deserialize, Serialize};
8use std::collections::BTreeMap;
9
10use super::bbox::BBoxXYXY;
11use super::ids::{AnnotationId, CategoryId, ImageId, LicenseId};
12use super::space::Pixel;
13
14/// A complete object detection dataset in the panlabel IR format.
15///
16/// This is the central data structure that all format conversions work through.
17/// Think of it as the "AST" in a compiler - formats parse into this representation,
18/// and this representation renders out to target formats.
19#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
20pub struct Dataset {
21    /// Metadata about the dataset (name, version, license, etc.)
22    #[serde(default)]
23    pub info: DatasetInfo,
24
25    /// License definitions for the dataset.
26    #[serde(default, skip_serializing_if = "Vec::is_empty")]
27    pub licenses: Vec<License>,
28
29    /// All images in the dataset.
30    pub images: Vec<Image>,
31
32    /// All category definitions.
33    pub categories: Vec<Category>,
34
35    /// All annotations (bounding boxes with labels).
36    pub annotations: Vec<Annotation>,
37}
38
39/// Metadata about the dataset.
40#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
41pub struct DatasetInfo {
42    /// Optional name of the dataset.
43    #[serde(default, skip_serializing_if = "Option::is_none")]
44    pub name: Option<String>,
45
46    /// Optional version string.
47    #[serde(default, skip_serializing_if = "Option::is_none")]
48    pub version: Option<String>,
49
50    /// Optional description.
51    #[serde(default, skip_serializing_if = "Option::is_none")]
52    pub description: Option<String>,
53
54    /// Optional URL for more information.
55    #[serde(default, skip_serializing_if = "Option::is_none")]
56    pub url: Option<String>,
57
58    /// Optional year the dataset was created.
59    #[serde(default, skip_serializing_if = "Option::is_none")]
60    pub year: Option<u32>,
61
62    /// Optional contributor name or organization.
63    #[serde(default, skip_serializing_if = "Option::is_none")]
64    pub contributor: Option<String>,
65
66    /// Optional date the dataset was created (ISO 8601 or similar).
67    #[serde(default, skip_serializing_if = "Option::is_none")]
68    pub date_created: Option<String>,
69
70    /// Adapter-specific dataset attributes and provenance metadata.
71    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
72    pub attributes: BTreeMap<String, String>,
73}
74
75impl DatasetInfo {
76    /// Returns true if all fields are None (i.e., no metadata is set).
77    pub fn is_empty(&self) -> bool {
78        self.name.is_none()
79            && self.version.is_none()
80            && self.description.is_none()
81            && self.url.is_none()
82            && self.year.is_none()
83            && self.contributor.is_none()
84            && self.date_created.is_none()
85            && self.attributes.is_empty()
86    }
87}
88
89/// A license that can be associated with images in the dataset.
90#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
91pub struct License {
92    /// Unique identifier for this license.
93    pub id: LicenseId,
94
95    /// Name of the license (e.g., "CC BY 4.0").
96    pub name: String,
97
98    /// Optional URL to the license text.
99    #[serde(default, skip_serializing_if = "Option::is_none")]
100    pub url: Option<String>,
101}
102
103impl License {
104    /// Creates a new license with the given properties.
105    pub fn new(id: impl Into<LicenseId>, name: impl Into<String>) -> Self {
106        Self {
107            id: id.into(),
108            name: name.into(),
109            url: None,
110        }
111    }
112
113    /// Creates a new license with a URL.
114    pub fn with_url(
115        id: impl Into<LicenseId>,
116        name: impl Into<String>,
117        url: impl Into<String>,
118    ) -> Self {
119        Self {
120            id: id.into(),
121            name: name.into(),
122            url: Some(url.into()),
123        }
124    }
125}
126
127/// An image in the dataset.
128#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
129pub struct Image {
130    /// Unique identifier for this image.
131    pub id: ImageId,
132
133    /// Filename or path of the image.
134    pub file_name: String,
135
136    /// Width of the image in pixels.
137    pub width: u32,
138
139    /// Height of the image in pixels.
140    pub height: u32,
141
142    /// Optional license ID for this image.
143    #[serde(default, skip_serializing_if = "Option::is_none")]
144    pub license_id: Option<LicenseId>,
145
146    /// Optional date the image was captured (ISO 8601 or similar).
147    #[serde(default, skip_serializing_if = "Option::is_none")]
148    pub date_captured: Option<String>,
149
150    /// Additional image-level attributes (e.g., VOC depth metadata).
151    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
152    pub attributes: BTreeMap<String, String>,
153}
154
155impl Image {
156    /// Creates a new image with the given properties.
157    pub fn new(
158        id: impl Into<ImageId>,
159        file_name: impl Into<String>,
160        width: u32,
161        height: u32,
162    ) -> Self {
163        Self {
164            id: id.into(),
165            file_name: file_name.into(),
166            width,
167            height,
168            license_id: None,
169            date_captured: None,
170            attributes: BTreeMap::new(),
171        }
172    }
173
174    /// Sets the license ID for this image.
175    pub fn with_license(mut self, license_id: impl Into<LicenseId>) -> Self {
176        self.license_id = Some(license_id.into());
177        self
178    }
179
180    /// Sets the date captured for this image.
181    pub fn with_date_captured(mut self, date: impl Into<String>) -> Self {
182        self.date_captured = Some(date.into());
183        self
184    }
185}
186
187impl From<u64> for ImageId {
188    fn from(id: u64) -> Self {
189        ImageId::new(id)
190    }
191}
192
193/// A category (class label) in the dataset.
194#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
195pub struct Category {
196    /// Unique identifier for this category.
197    pub id: CategoryId,
198
199    /// Name of the category (e.g., "person", "car", "dog").
200    pub name: String,
201
202    /// Optional supercategory for hierarchical taxonomies.
203    #[serde(default, skip_serializing_if = "Option::is_none")]
204    pub supercategory: Option<String>,
205}
206
207impl Category {
208    /// Creates a new category with the given properties.
209    pub fn new(id: impl Into<CategoryId>, name: impl Into<String>) -> Self {
210        Self {
211            id: id.into(),
212            name: name.into(),
213            supercategory: None,
214        }
215    }
216
217    /// Creates a new category with a supercategory.
218    pub fn with_supercategory(
219        id: impl Into<CategoryId>,
220        name: impl Into<String>,
221        supercategory: impl Into<String>,
222    ) -> Self {
223        Self {
224            id: id.into(),
225            name: name.into(),
226            supercategory: Some(supercategory.into()),
227        }
228    }
229}
230
231impl From<u64> for CategoryId {
232    fn from(id: u64) -> Self {
233        CategoryId::new(id)
234    }
235}
236
237/// An annotation (bounding box with label) in the dataset.
238#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
239pub struct Annotation {
240    /// Unique identifier for this annotation.
241    pub id: AnnotationId,
242
243    /// ID of the image this annotation belongs to.
244    pub image_id: ImageId,
245
246    /// ID of the category (class) for this annotation.
247    pub category_id: CategoryId,
248
249    /// Bounding box in pixel coordinates (XYXY format).
250    pub bbox: BBoxXYXY<Pixel>,
251
252    /// Optional confidence score (e.g., from model predictions).
253    #[serde(default, skip_serializing_if = "Option::is_none")]
254    pub confidence: Option<f64>,
255
256    /// Additional attributes (e.g., "occluded", "truncated").
257    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
258    pub attributes: BTreeMap<String, String>,
259}
260
261impl Annotation {
262    /// Creates a new annotation with the minimum required fields.
263    pub fn new(
264        id: impl Into<AnnotationId>,
265        image_id: impl Into<ImageId>,
266        category_id: impl Into<CategoryId>,
267        bbox: BBoxXYXY<Pixel>,
268    ) -> Self {
269        Self {
270            id: id.into(),
271            image_id: image_id.into(),
272            category_id: category_id.into(),
273            bbox,
274            confidence: None,
275            attributes: BTreeMap::new(),
276        }
277    }
278
279    /// Adds a confidence score to the annotation.
280    pub fn with_confidence(mut self, confidence: f64) -> Self {
281        self.confidence = Some(confidence);
282        self
283    }
284
285    /// Adds an attribute to the annotation.
286    pub fn with_attribute(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
287        self.attributes.insert(key.into(), value.into());
288        self
289    }
290}
291
292impl From<u64> for AnnotationId {
293    fn from(id: u64) -> Self {
294        AnnotationId::new(id)
295    }
296}
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301
302    #[test]
303    fn test_dataset_creation() {
304        let dataset = Dataset {
305            info: DatasetInfo {
306                name: Some("Test Dataset".into()),
307                ..Default::default()
308            },
309            licenses: vec![],
310            images: vec![Image::new(1u64, "image001.jpg", 640, 480)],
311            categories: vec![Category::new(1u64, "person")],
312            annotations: vec![Annotation::new(
313                1u64,
314                1u64,
315                1u64,
316                BBoxXYXY::from_xyxy(10.0, 20.0, 100.0, 200.0),
317            )],
318        };
319
320        assert_eq!(dataset.images.len(), 1);
321        assert_eq!(dataset.categories.len(), 1);
322        assert_eq!(dataset.annotations.len(), 1);
323    }
324
325    #[test]
326    fn test_annotation_builder_pattern() {
327        let annotation =
328            Annotation::new(1u64, 1u64, 1u64, BBoxXYXY::from_xyxy(0.0, 0.0, 50.0, 50.0))
329                .with_confidence(0.95)
330                .with_attribute("occluded", "false")
331                .with_attribute("truncated", "true");
332
333        assert_eq!(annotation.confidence, Some(0.95));
334        assert_eq!(annotation.attributes.len(), 2);
335        assert_eq!(
336            annotation.attributes.get("occluded"),
337            Some(&"false".to_string())
338        );
339    }
340}