Skip to main content

semantic_scene/mp3d/
mod.rs

1//! MP3D dataset support.
2use std::{
3    collections::{HashMap, HashSet},
4    io::BufRead,
5};
6
7use crate::{
8    Dataset, LoadError, Rotation3, SemanticLevel, SemanticObject, SemanticRegion, SemanticScene,
9    scene::ElementKind,
10};
11
12mod category;
13pub(crate) mod parser;
14pub(crate) mod raw;
15
16pub use category::{Mp3dCategoryMapping, Mp3dObjectCategory, Mp3dRegionCategory};
17use raw::Mp3dRecord;
18
19/// Loader for Habitat-Sim `Matterport3D` `.house` semantic descriptors.
20pub struct Mp3d;
21
22type Mp3dScene = SemanticScene<Mp3dObjectCategory, Mp3dRegionCategory>;
23
24/// Options for loading `Matterport3D` `.house` descriptors.
25#[derive(Debug, Clone)]
26pub struct Mp3dOptions {
27    /// Rotation to apply while loading.
28    ///
29    /// Defaults to [`Rotation3::HABITAT_MP3D`] to match Habitat-Sim's MP3D
30    /// loader. Use [`Rotation3::IDENTITY`] to keep raw MP3D coordinates.
31    pub rotation: Rotation3,
32}
33
34impl Default for Mp3dOptions {
35    fn default() -> Self {
36        Self {
37            rotation: Rotation3::HABITAT_MP3D,
38        }
39    }
40}
41
42impl Dataset for Mp3d {
43    type Options = Mp3dOptions;
44    type Error = LoadError;
45    type ObjectCategory = Mp3dObjectCategory;
46    type RegionCategory = Mp3dRegionCategory;
47
48    fn from_reader<R: BufRead>(
49        mut reader: R,
50        options: Self::Options,
51    ) -> Result<Mp3dScene, Self::Error> {
52        let mut header = String::new();
53        let bytes = reader.read_line(&mut header)?;
54        if bytes == 0 {
55            return Err(LoadError::BadHeader {
56                found: "<empty file>".to_string(),
57            });
58        }
59        let header = header.trim_end_matches(['\r', '\n']);
60        if header != "ASCII 1.1" {
61            return Err(LoadError::BadHeader {
62                found: header.to_string(),
63            });
64        }
65
66        let mut scene_header = None;
67        let mut levels = Vec::<raw::LevelRecord>::new();
68        let mut regions = Vec::<(raw::RegionRecord, usize)>::new();
69        let mut categories = HashMap::<i32, Mp3dObjectCategory>::new();
70        let mut objects = Vec::<(raw::ObjectRecord, usize)>::new();
71        let mut segments = HashMap::<i32, (i32, usize)>::new();
72
73        for (offset, line) in reader.lines().enumerate() {
74            let line_number = offset + 2;
75            let line = line?;
76            if line.trim().is_empty() {
77                continue;
78            }
79
80            let record = parser::parse_record(&line).map_err(|source| LoadError::ParseLine {
81                line_number,
82                line: line.clone(),
83                source,
84            })?;
85            match record {
86                Mp3dRecord::House(record) => {
87                    scene_header = Some(record);
88                }
89                Mp3dRecord::Level(record) => levels.push(record),
90                Mp3dRecord::Region(record) => regions.push((record, line_number)),
91                Mp3dRecord::Category(record) => {
92                    categories.insert(
93                        record.index,
94                        Mp3dObjectCategory::new(
95                            record.index,
96                            record.raw_index,
97                            record.raw_name,
98                            record.mpcat40_index,
99                            record.mpcat40_name,
100                        ),
101                    );
102                }
103                Mp3dRecord::Object(record) => objects.push((record, line_number)),
104                Mp3dRecord::Segment(record) => {
105                    if segments
106                        .insert(record.segment_id, (record.object_index, line_number))
107                        .is_some()
108                    {
109                        return Err(LoadError::DuplicateSegmentId {
110                            line_number,
111                            segment_id: record.segment_id,
112                        });
113                    }
114                }
115                Mp3dRecord::Ignored => {}
116            }
117        }
118
119        build_scene(
120            scene_header,
121            levels,
122            regions,
123            &categories,
124            objects,
125            &segments,
126            options.rotation,
127        )
128    }
129}
130
131fn build_scene(
132    scene_header: Option<raw::HouseRecord>,
133    levels: Vec<raw::LevelRecord>,
134    regions: Vec<(raw::RegionRecord, usize)>,
135    categories: &HashMap<i32, Mp3dObjectCategory>,
136    objects: Vec<(raw::ObjectRecord, usize)>,
137    segments: &HashMap<i32, (i32, usize)>,
138    rotation: Rotation3,
139) -> Result<Mp3dScene, LoadError> {
140    let scene_header = scene_header.ok_or_else(|| LoadError::BadHeader {
141        found: "<missing house record>".to_string(),
142    })?;
143    validate_segments(&objects, segments)?;
144
145    let region_indices = source_indices(&regions);
146    let objects_by_region =
147        build_objects_by_region(objects, categories, &region_indices, rotation)?;
148
149    let level_indices = levels
150        .iter()
151        .map(|record| record.index)
152        .collect::<HashSet<_>>();
153    let regions_by_level =
154        build_regions_by_level(regions, objects_by_region, &level_indices, rotation)?;
155    let levels = build_levels(levels, regions_by_level, rotation);
156
157    let counts = house_counts(&scene_header);
158    Ok(SemanticScene::from_parts(
159        scene_header.name,
160        scene_header.label,
161        counts,
162        Some(scene_header.aabb.rotated(rotation)),
163        levels,
164    ))
165}
166
167fn validate_segments(
168    objects: &[(raw::ObjectRecord, usize)],
169    segments: &HashMap<i32, (i32, usize)>,
170) -> Result<(), LoadError> {
171    let object_indices = objects
172        .iter()
173        .map(|(record, _)| record.index)
174        .collect::<HashSet<_>>();
175    for (object_index, line_number) in segments.values().copied() {
176        if !object_indices.contains(&object_index) {
177            return Err(LoadError::MissingParent {
178                line_number,
179                kind: "object",
180                index: object_index,
181            });
182        }
183    }
184    Ok(())
185}
186
187fn source_indices(records: &[(raw::RegionRecord, usize)]) -> HashSet<i32> {
188    records.iter().map(|(record, _)| record.index).collect()
189}
190
191fn build_objects_by_region(
192    objects: Vec<(raw::ObjectRecord, usize)>,
193    categories: &HashMap<i32, Mp3dObjectCategory>,
194    region_indices: &HashSet<i32>,
195    rotation: Rotation3,
196) -> Result<HashMap<i32, Vec<SemanticObject<Mp3dObjectCategory>>>, LoadError> {
197    let mut objects_by_region = HashMap::<i32, Vec<SemanticObject<Mp3dObjectCategory>>>::new();
198    for (record, line_number) in objects {
199        assert!(
200            record.region_index >= 0,
201            "orphan MP3D object {} on line {}",
202            record.index,
203            line_number
204        );
205        if !region_indices.contains(&record.region_index) {
206            return Err(LoadError::MissingParent {
207                line_number,
208                kind: "region",
209                index: record.region_index,
210            });
211        }
212        let category = if record.category_index < 0 {
213            None
214        } else {
215            Some(categories.get(&record.category_index).cloned().ok_or(
216                LoadError::MissingCategory {
217                    line_number,
218                    index: record.category_index,
219                },
220            )?)
221        };
222        objects_by_region
223            .entry(record.region_index)
224            .or_default()
225            .push(SemanticObject::new(
226                record.index,
227                category,
228                record.obb.rotated(rotation),
229            ));
230    }
231    Ok(objects_by_region)
232}
233
234fn build_regions_by_level(
235    regions: Vec<(raw::RegionRecord, usize)>,
236    mut objects_by_region: HashMap<i32, Vec<SemanticObject<Mp3dObjectCategory>>>,
237    level_indices: &HashSet<i32>,
238    rotation: Rotation3,
239) -> Result<HashMap<i32, Vec<SemanticRegion<Mp3dObjectCategory, Mp3dRegionCategory>>>, LoadError> {
240    let mut regions_by_level =
241        HashMap::<i32, Vec<SemanticRegion<Mp3dObjectCategory, Mp3dRegionCategory>>>::new();
242    for (record, line_number) in regions {
243        assert!(
244            record.level_index >= 0,
245            "orphan MP3D region {} on line {}",
246            record.index,
247            line_number
248        );
249        if !level_indices.contains(&record.level_index) {
250            return Err(LoadError::MissingParent {
251                line_number,
252                kind: "level",
253                index: record.level_index,
254            });
255        }
256        let objects = objects_by_region.remove(&record.index).unwrap_or_default();
257        regions_by_level
258            .entry(record.level_index)
259            .or_default()
260            .push(SemanticRegion::new(
261                record.index,
262                Mp3dRegionCategory::new(record.category_code),
263                rotation.transform_vector(record.position),
264                record.aabb.rotated(rotation),
265                objects,
266            ));
267    }
268    Ok(regions_by_level)
269}
270
271fn build_levels(
272    levels: Vec<raw::LevelRecord>,
273    mut regions_by_level: HashMap<i32, Vec<SemanticRegion<Mp3dObjectCategory, Mp3dRegionCategory>>>,
274    rotation: Rotation3,
275) -> Vec<SemanticLevel<Mp3dObjectCategory, Mp3dRegionCategory>> {
276    levels
277        .into_iter()
278        .map(|record| {
279            let regions = regions_by_level.remove(&record.index).unwrap_or_default();
280            SemanticLevel::new(
281                record.index,
282                record.label,
283                rotation.transform_vector(record.position),
284                record.aabb.rotated(rotation),
285                regions,
286            )
287        })
288        .collect()
289}
290
291fn house_counts(record: &raw::HouseRecord) -> HashMap<&'static str, usize> {
292    HashMap::from([
293        (ElementKind::Images.as_str(), record.images),
294        (ElementKind::Panoramas.as_str(), record.panoramas),
295        (ElementKind::Vertices.as_str(), record.vertices),
296        (ElementKind::Surfaces.as_str(), record.surfaces),
297        (ElementKind::Segments.as_str(), record.segments),
298        (ElementKind::Objects.as_str(), record.objects),
299        (ElementKind::Categories.as_str(), record.categories),
300        (ElementKind::Regions.as_str(), record.regions),
301        (ElementKind::Portals.as_str(), record.portals),
302        (ElementKind::Levels.as_str(), record.levels),
303    ])
304}