semantic-scene 0.1.0

//! `Matterport3D` `.house` loader.
use std::{collections::HashMap, io::BufRead};

use crate::{
    Category, ElementKind, LoadError, RegionCategory, Rotation3, SemanticLevel, SemanticObject,
    SemanticRegion, SemanticScene, loader::SemanticSceneLoader,
};

mod category;
pub(crate) mod parser;
pub(crate) mod raw;

use raw::Mp3dRecord;

/// Loader for Habitat-Sim `Matterport3D` `.house` semantic descriptors.
pub struct Mp3dLoader;

/// Options for loading `Matterport3D` `.house` descriptors.
#[derive(Debug, Clone, Default)]
pub struct Mp3dOptions {
    /// Rotation to apply while loading.
    ///
    /// Only [`Rotation3::Identity`] is supported in this milestone. Passing a
    /// non-identity rotation returns [`LoadError::UnsupportedOption`].
    pub rotation: Rotation3,
}

impl SemanticSceneLoader for Mp3dLoader {
    type Options = Mp3dOptions;
    type Error = LoadError;

    fn from_reader<R: BufRead>(
        mut reader: R,
        options: Self::Options,
    ) -> Result<SemanticScene, Self::Error> {
        if options.rotation != Rotation3::Identity {
            return Err(LoadError::UnsupportedOption("non-identity MP3D rotation"));
        }

        let mut header = String::new();
        let bytes = reader.read_line(&mut header)?;
        if bytes == 0 {
            return Err(LoadError::BadHeader {
                found: "<empty file>".to_string(),
            });
        }
        let header = header.trim_end_matches(['\r', '\n']);
        if header != "ASCII 1.1" {
            return Err(LoadError::BadHeader {
                found: header.to_string(),
            });
        }

        let mut scene = SemanticScene::new();
        for (offset, line) in reader.lines().enumerate() {
            let line_number = offset + 2;
            let line = line?;
            if line.trim().is_empty() {
                continue;
            }

            let record = parser::parse_record(&line).map_err(|source| LoadError::ParseLine {
                line_number,
                line: line.clone(),
                source,
            })?;
            apply_record(&mut scene, record, line_number)?;
        }

        Ok(scene)
    }
}

fn apply_record(
    scene: &mut SemanticScene,
    record: Mp3dRecord,
    line_number: usize,
) -> Result<(), LoadError> {
    match record {
        Mp3dRecord::House(record) => {
            let counts = house_counts(&record);
            scene.set_header(record.name, record.label, counts, record.aabb);
        }
        Mp3dRecord::Level(record) => {
            scene.push_level(SemanticLevel::new(
                record.index,
                record.label,
                record.position,
                record.aabb,
            ));
        }
        Mp3dRecord::Region(record) => apply_region(scene, &record, line_number)?,
        Mp3dRecord::Category(record) => {
            scene.push_category(Category::new(
                record.index,
                record.raw_index,
                record.raw_name,
                record.mpcat40_index,
                record.mpcat40_name,
            ));
        }
        Mp3dRecord::Object(record) => apply_object(scene, &record, line_number)?,
        Mp3dRecord::Segment(record) => {
            let object_index = scene.object_position_by_index(record.object_index).ok_or(
                LoadError::MissingParent {
                    line_number,
                    kind: "object",
                    index: record.object_index,
                },
            )?;
            if scene
                .insert_segment(record.segment_id, object_index)
                .is_some()
            {
                return Err(LoadError::DuplicateSegmentId {
                    line_number,
                    segment_id: record.segment_id,
                });
            }
        }
        Mp3dRecord::Ignored => {}
    }
    Ok(())
}

fn house_counts(record: &raw::HouseRecord) -> HashMap<&'static str, usize> {
    HashMap::from([
        (ElementKind::Images.as_str(), record.images),
        (ElementKind::Panoramas.as_str(), record.panoramas),
        (ElementKind::Vertices.as_str(), record.vertices),
        (ElementKind::Surfaces.as_str(), record.surfaces),
        (ElementKind::Segments.as_str(), record.segments),
        (ElementKind::Objects.as_str(), record.objects),
        (ElementKind::Categories.as_str(), record.categories),
        (ElementKind::Regions.as_str(), record.regions),
        (ElementKind::Portals.as_str(), record.portals),
        (ElementKind::Levels.as_str(), record.levels),
    ])
}

fn apply_region(
    scene: &mut SemanticScene,
    record: &raw::RegionRecord,
    line_number: usize,
) -> Result<(), LoadError> {
    let level_index =
        if record.level_index < 0 {
            None
        } else {
            Some(scene.level_position_by_index(record.level_index).ok_or(
                LoadError::MissingParent {
                    line_number,
                    kind: "level",
                    index: record.level_index,
                },
            )?)
        };
    let region_index = scene.regions().len();
    scene.push_region(SemanticRegion::new(
        record.index,
        level_index,
        RegionCategory::new(record.category_code),
        record.position,
        record.aabb,
    ));
    if let Some(level_index) = level_index
        && let Some(level) = scene.level_mut(level_index)
    {
        level.add_region(region_index);
    }
    Ok(())
}

fn apply_object(
    scene: &mut SemanticScene,
    record: &raw::ObjectRecord,
    line_number: usize,
) -> Result<(), LoadError> {
    let region_index = if record.region_index < 0 {
        None
    } else {
        Some(scene.region_position_by_index(record.region_index).ok_or(
            LoadError::MissingParent {
                line_number,
                kind: "region",
                index: record.region_index,
            },
        )?)
    };
    let category_index = if record.category_index < 0 {
        None
    } else {
        Some(
            scene
                .category_position_by_index(record.category_index)
                .ok_or(LoadError::MissingCategory {
                    line_number,
                    index: record.category_index,
                })?,
        )
    };

    let object_index = scene.objects().len();
    scene.push_object(SemanticObject::new(
        record.index,
        region_index,
        category_index,
        record.obb,
    ));
    if let Some(region_index) = region_index {
        let level_index = scene.regions()[region_index].level_index();
        if let Some(region) = scene.region_mut(region_index) {
            region.add_object(object_index);
        }
        if let Some(level_index) = level_index
            && let Some(level) = scene.level_mut(level_index)
        {
            level.add_object(object_index);
        }
    }
    Ok(())
}