1use 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#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
20pub struct Dataset {
21 #[serde(default)]
23 pub info: DatasetInfo,
24
25 #[serde(default, skip_serializing_if = "Vec::is_empty")]
27 pub licenses: Vec<License>,
28
29 pub images: Vec<Image>,
31
32 pub categories: Vec<Category>,
34
35 pub annotations: Vec<Annotation>,
37}
38
39#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
41pub struct DatasetInfo {
42 #[serde(default, skip_serializing_if = "Option::is_none")]
44 pub name: Option<String>,
45
46 #[serde(default, skip_serializing_if = "Option::is_none")]
48 pub version: Option<String>,
49
50 #[serde(default, skip_serializing_if = "Option::is_none")]
52 pub description: Option<String>,
53
54 #[serde(default, skip_serializing_if = "Option::is_none")]
56 pub url: Option<String>,
57
58 #[serde(default, skip_serializing_if = "Option::is_none")]
60 pub year: Option<u32>,
61
62 #[serde(default, skip_serializing_if = "Option::is_none")]
64 pub contributor: Option<String>,
65
66 #[serde(default, skip_serializing_if = "Option::is_none")]
68 pub date_created: Option<String>,
69
70 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
72 pub attributes: BTreeMap<String, String>,
73}
74
75impl DatasetInfo {
76 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#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
91pub struct License {
92 pub id: LicenseId,
94
95 pub name: String,
97
98 #[serde(default, skip_serializing_if = "Option::is_none")]
100 pub url: Option<String>,
101}
102
103impl License {
104 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 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#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
129pub struct Image {
130 pub id: ImageId,
132
133 pub file_name: String,
135
136 pub width: u32,
138
139 pub height: u32,
141
142 #[serde(default, skip_serializing_if = "Option::is_none")]
144 pub license_id: Option<LicenseId>,
145
146 #[serde(default, skip_serializing_if = "Option::is_none")]
148 pub date_captured: Option<String>,
149
150 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
152 pub attributes: BTreeMap<String, String>,
153}
154
155impl Image {
156 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 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 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#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
195pub struct Category {
196 pub id: CategoryId,
198
199 pub name: String,
201
202 #[serde(default, skip_serializing_if = "Option::is_none")]
204 pub supercategory: Option<String>,
205}
206
207impl Category {
208 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 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#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
239pub struct Annotation {
240 pub id: AnnotationId,
242
243 pub image_id: ImageId,
245
246 pub category_id: CategoryId,
248
249 pub bbox: BBoxXYXY<Pixel>,
251
252 #[serde(default, skip_serializing_if = "Option::is_none")]
254 pub confidence: Option<f64>,
255
256 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
258 pub attributes: BTreeMap<String, String>,
259}
260
261impl Annotation {
262 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 pub fn with_confidence(mut self, confidence: f64) -> Self {
281 self.confidence = Some(confidence);
282 self
283 }
284
285 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}