1use crate::{
2 error::{Error, Result},
3 serializable::{
4 Attribute, CalibratedSensor, Category, EgoPose, Instance, Log, Map, Sample,
5 SampleAnnotation, SampleData, Scene, Sensor, Token, Visibility, VisibilityToken,
6 },
7};
8use chrono::NaiveDateTime;
9use std::{collections::HashMap, path::PathBuf};
10
11#[derive(Debug, Clone)]
12pub struct DatasetInner {
13 pub version: String,
14 pub dataset_dir: PathBuf,
15 pub attribute_map: HashMap<Token, Attribute>,
16 pub calibrated_sensor_map: HashMap<Token, CalibratedSensor>,
17 pub category_map: HashMap<Token, Category>,
18 pub ego_pose_map: HashMap<Token, EgoPose>,
19 pub instance_map: HashMap<Token, InstanceInner>,
20 pub log_map: HashMap<Token, Log>,
21 pub map_map: HashMap<Token, Map>,
22 pub scene_map: HashMap<Token, SceneInner>,
23 pub sample_map: HashMap<Token, SampleInner>,
24 pub sample_annotation_map: HashMap<Token, SampleAnnotation>,
25 pub sample_data_map: HashMap<Token, SampleData>,
26 pub sensor_map: HashMap<Token, Sensor>,
27 pub visibility_map: HashMap<VisibilityToken, Visibility>,
28 pub sorted_ego_pose_tokens: Vec<Token>,
29 pub sorted_sample_tokens: Vec<Token>,
30 pub sorted_sample_data_tokens: Vec<Token>,
31 pub sorted_scene_tokens: Vec<Token>,
32}
33
34#[derive(Debug, Clone)]
35pub struct SampleInner {
36 pub token: Token,
37 pub next: Option<Token>,
38 pub prev: Option<Token>,
39 pub timestamp: NaiveDateTime,
40 pub scene_token: Token,
41 pub annotation_tokens: Vec<Token>,
42 pub sample_data_tokens: Vec<Token>,
43}
44
45impl SampleInner {
46 pub fn from(
47 sample: Sample,
48 annotation_tokens: Vec<Token>,
49 sample_data_tokens: Vec<Token>,
50 ) -> Self {
51 let Sample {
52 token,
53 next,
54 prev,
55 scene_token,
56 timestamp,
57 } = sample;
58
59 Self {
60 token,
61 next,
62 prev,
63 scene_token,
64 timestamp,
65 annotation_tokens,
66 sample_data_tokens,
67 }
68 }
69}
70
71#[derive(Debug, Clone)]
72pub struct InstanceInner {
73 pub token: Token,
74 pub category_token: Token,
75 pub annotation_tokens: Vec<Token>,
76}
77
78impl InstanceInner {
79 pub fn from(
80 instance: Instance,
81 sample_annotation_map: &HashMap<Token, SampleAnnotation>,
82 ) -> Result<Self> {
83 let Instance {
84 token,
85 nbr_annotations,
86 category_token,
87 first_annotation_token,
88 last_annotation_token,
89 } = instance;
90
91 let mut annotation_token_opt = Some(first_annotation_token);
92 let mut annotation_tokens = vec![];
93
94 while let Some(annotation_token) = annotation_token_opt {
95 let annotation = &sample_annotation_map
96 .get(&annotation_token)
97 .expect("internal error: invalid annotation_token");
98 assert_eq!(
99 annotation_token, annotation.token,
100 "internal error: annotation.token mismatch"
101 );
102 annotation_tokens.push(annotation_token);
103 annotation_token_opt = annotation.next;
104 }
105
106 if annotation_tokens.len() != nbr_annotations {
107 let msg = format!(
108 "the instance with token {} assures nbr_annotations = {}, but in fact {}",
109 token,
110 nbr_annotations,
111 annotation_tokens.len()
112 );
113 return Err(Error::CorruptedDataset(msg));
114 }
115 if annotation_tokens.last().unwrap() != &last_annotation_token {
116 let msg = format!(
117 "the instance with token {} assures last_annotation_token = {}, but in fact {}",
118 token,
119 last_annotation_token,
120 annotation_tokens.last().unwrap()
121 );
122 return Err(Error::CorruptedDataset(msg));
123 }
124
125 let ret = Self {
126 token,
127 category_token,
128 annotation_tokens,
129 };
130 Ok(ret)
131 }
132}
133
134#[derive(Debug, Clone)]
135pub struct SceneInner {
136 pub token: Token,
137 pub name: String,
138 pub description: String,
139 pub log_token: Token,
140 pub sample_tokens: Vec<Token>,
141}
142
143impl SceneInner {
144 pub fn from(scene: Scene, sample_map: &HashMap<Token, Sample>) -> Result<Self> {
145 let Scene {
146 token,
147 name,
148 description,
149 log_token,
150 nbr_samples,
151 first_sample_token,
152 last_sample_token,
153 } = scene;
154
155 let mut sample_tokens = vec![];
156 let mut sample_token_opt = Some(first_sample_token);
157
158 while let Some(sample_token) = sample_token_opt {
159 let sample = &sample_map[&sample_token];
160 assert_eq!(
161 sample.token, sample_token,
162 "internal error: sample.token mismatch"
163 );
164 sample_tokens.push(sample_token);
165 sample_token_opt = sample.next;
166 }
167
168 if sample_tokens.len() != nbr_samples {
169 let msg = format!(
170 "the sample with token {} assures nbr_samples = {}, but in fact {}",
171 token,
172 nbr_samples,
173 sample_tokens.len()
174 );
175 return Err(Error::CorruptedDataset(msg));
176 }
177 if *sample_tokens.last().unwrap() != last_sample_token {
178 let msg = format!(
179 "the sample with token {} assures last_sample_token = {}, but in fact {}",
180 token,
181 last_sample_token,
182 sample_tokens.last().unwrap()
183 );
184 return Err(Error::CorruptedDataset(msg));
185 }
186
187 let ret = Self {
188 token,
189 name,
190 description,
191 log_token,
192 sample_tokens,
193 };
194 Ok(ret)
195 }
196}