semantic-scene 0.1.0

//! Semantic scene graph model.

use std::fmt;

mod category;
mod level;
mod object;
mod region;

use std::{collections::HashMap, path::Path};

pub use category::{Category, CategoryMapping, RegionCategory};
pub use level::SemanticLevel;
pub use object::SemanticObject;
pub use region::SemanticRegion;

use crate::{Aabb, LoadError, Mp3dLoader, Mp3dOptions, SemanticSceneLoader};

/// Scene element kind used by `SemanticScene::count`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ElementKind {
    /// Image count from the descriptor header.
    Images,
    /// Panorama count from the descriptor header.
    Panoramas,
    /// Vertex count from the descriptor header.
    Vertices,
    /// Surface count from the descriptor header.
    Surfaces,
    /// Segment count from the descriptor header.
    Segments,
    /// Object count from the descriptor header.
    Objects,
    /// Category count from the descriptor header.
    Categories,
    /// Region count from the descriptor header.
    Regions,
    /// Portal count from the descriptor header.
    Portals,
    /// Level count from the descriptor header.
    Levels,
}

impl ElementKind {
    pub(crate) const fn as_str(self) -> &'static str {
        match self {
            Self::Images => "images",
            Self::Panoramas => "panoramas",
            Self::Vertices => "vertices",
            Self::Surfaces => "surfaces",
            Self::Segments => "segments",
            Self::Objects => "objects",
            Self::Categories => "categories",
            Self::Regions => "regions",
            Self::Portals => "portals",
            Self::Levels => "levels",
        }
    }
}

/// Scene containing semantically annotated levels, regions, and objects.
///
/// This mirrors Habitat-Sim's `SemanticScene` at the API-shape level while
/// using Rust ownership and indices instead of shared-pointer graphs.
#[derive(Debug, Default, Clone)]
pub struct SemanticScene {
    pub(crate) name: String,
    pub(crate) label: String,
    pub(crate) counts: HashMap<&'static str, usize>,
    pub(crate) aabb: Option<Aabb>,
    pub(crate) categories: Vec<Category>,
    pub(crate) levels: Vec<SemanticLevel>,
    pub(crate) regions: Vec<SemanticRegion>,
    pub(crate) objects: Vec<SemanticObject>,
    pub(crate) semantic_index_map: HashMap<i32, usize>,
}

impl SemanticScene {
    pub(crate) fn new() -> Self {
        Self::default()
    }

    /// Returns the scene name from the descriptor header.
    #[must_use]
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Returns the scene label from the descriptor header.
    #[must_use]
    pub fn label(&self) -> &str {
        &self.label
    }

    /// Returns the scene axis-aligned bounding box.
    ///
    /// Habitat-Sim exposes this as `aabb()`. It is optional here because a
    /// malformed or partially built scene may not have a parsed header.
    #[must_use]
    pub const fn aabb(&self) -> Option<&Aabb> {
        self.aabb.as_ref()
    }

    /// Returns the total number of a given element type from the descriptor header.
    #[must_use]
    pub fn count(&self, element: ElementKind) -> Option<usize> {
        self.counts.get(element.as_str()).copied()
    }

    /// Returns all semantic object categories in the scene.
    #[must_use]
    pub fn categories(&self) -> &[Category] {
        &self.categories
    }

    /// Returns all levels in the scene.
    #[must_use]
    pub fn levels(&self) -> &[SemanticLevel] {
        &self.levels
    }

    /// Returns all regions in the scene.
    #[must_use]
    pub fn regions(&self) -> &[SemanticRegion] {
        &self.regions
    }

    /// Returns all semantic objects in the scene.
    #[must_use]
    pub fn objects(&self) -> &[SemanticObject] {
        &self.objects
    }

    /// Returns the semantic mesh mask index to object index map.
    #[must_use]
    pub const fn semantic_index_map(&self) -> &HashMap<i32, usize> {
        &self.semantic_index_map
    }

    /// Converts a semantic mesh mask index to an object index.
    ///
    /// Habitat-Sim returns `ID_UNDEFINED` on miss. This Rust API returns
    /// `None` instead.
    #[must_use]
    pub fn semantic_index_to_object_index(&self, mask_index: i32) -> Option<usize> {
        self.semantic_index_map.get(&mask_index).copied()
    }

    /// Loads a `Matterport3D` semantic scene from a `.house` file.
    ///
    /// # Errors
    ///
    /// Returns an error if the file cannot be read or the `.house` descriptor is invalid.
    pub fn load_mp3d_house(path: impl AsRef<Path>) -> Result<Self, LoadError> {
        Mp3dLoader::from_path(path, Mp3dOptions::default())
    }

    pub(crate) fn set_header(
        &mut self,
        name: String,
        label: String,
        counts: HashMap<&'static str, usize>,
        aabb: Aabb,
    ) {
        self.name = name;
        self.label = label;
        self.counts = counts;
        self.aabb = Some(aabb);
    }

    pub(crate) fn push_category(&mut self, category: Category) {
        self.categories.push(category);
    }

    pub(crate) fn push_level(&mut self, level: SemanticLevel) {
        self.levels.push(level);
    }

    pub(crate) fn push_region(&mut self, region: SemanticRegion) {
        self.regions.push(region);
    }

    pub(crate) fn push_object(&mut self, object: SemanticObject) {
        self.objects.push(object);
    }

    pub(crate) fn insert_segment(&mut self, segment_id: i32, object_index: usize) -> Option<usize> {
        self.semantic_index_map.insert(segment_id, object_index)
    }

    pub(crate) fn level_position_by_index(&self, index: i32) -> Option<usize> {
        self.levels.iter().position(|level| level.index() == index)
    }

    pub(crate) fn region_position_by_index(&self, index: i32) -> Option<usize> {
        self.regions
            .iter()
            .position(|region| region.index() == index)
    }

    pub(crate) fn category_position_by_index(&self, index: i32) -> Option<usize> {
        self.categories
            .iter()
            .position(|category| category.source_index() == index)
    }

    pub(crate) fn object_position_by_index(&self, index: i32) -> Option<usize> {
        self.objects
            .iter()
            .position(|object| object.index() == index)
    }

    pub(crate) fn level_mut(&mut self, index: usize) -> Option<&mut SemanticLevel> {
        self.levels.get_mut(index)
    }

    pub(crate) fn region_mut(&mut self, index: usize) -> Option<&mut SemanticRegion> {
        self.regions.get_mut(index)
    }
}

impl fmt::Display for SemanticScene {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        if formatter.alternate() {
            write!(
                formatter,
                "SemanticScene {:?}\n  label: {:?}\n  bounds: {}\n  counts: {}",
                self.name,
                self.label,
                self.display_aabb(),
                self.display_counts()
            )
        } else {
            write!(
                formatter,
                "SemanticScene {:?}: {}, bounds=[{}]",
                self.name,
                self.display_counts(),
                self.display_aabb()
            )
        }
    }
}

impl SemanticScene {
    const fn display_aabb(&self) -> DisplayAabb<'_> {
        DisplayAabb(self.aabb.as_ref())
    }

    const fn display_counts(&self) -> DisplayCounts<'_> {
        DisplayCounts(self)
    }
}

struct DisplayAabb<'a>(Option<&'a Aabb>);

impl fmt::Display for DisplayAabb<'_> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.0 {
            Some(aabb) => aabb.fmt(formatter),
            None => formatter.write_str("none"),
        }
    }
}

struct DisplayCounts<'a>(&'a SemanticScene);

impl fmt::Display for DisplayCounts<'_> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            formatter,
            "levels={}, regions={}, objects={}, categories={}, segments={}",
            self.0
                .count(ElementKind::Levels)
                .unwrap_or(self.0.levels.len()),
            self.0
                .count(ElementKind::Regions)
                .unwrap_or(self.0.regions.len()),
            self.0
                .count(ElementKind::Objects)
                .unwrap_or(self.0.objects.len()),
            self.0
                .count(ElementKind::Categories)
                .unwrap_or(self.0.categories.len()),
            self.0
                .count(ElementKind::Segments)
                .unwrap_or(self.0.semantic_index_map.len()),
        )
    }
}