SemanticScene.h

// Copyright (c) Meta Platforms, Inc. and its affiliates.
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

#ifndef ESP_SCENE_SEMANTICSCENE_H_
#define ESP_SCENE_SEMANTICSCENE_H_

#include <Corrade/Containers/StringStl.h>
#include <Corrade/Utility/Path.h>
#include <map>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>

#include "esp/core/Esp.h"
#include "esp/core/Utility.h"
#include "esp/geo/OBB.h"
#include "esp/io/Json.h"

namespace esp {
namespace metadata {
namespace attributes {
class SemanticAttributes;
}
}  // namespace metadata
namespace scene {

namespace Mn = Magnum;

//! Represents a semantic category
class SemanticCategory {
 public:
  virtual ~SemanticCategory() = default;
  //! Return index of SemanticCategory under given mapping
  virtual int index(const std::string& mapping = "") const = 0;
  //! Return name of SemanticCategory under given mapping
  virtual std::string name(const std::string& mapping = "") const = 0;

  ESP_SMART_POINTERS(SemanticCategory)
};

// forward declarations
class SemanticObject;

// semantic object only used to represent characteristics of Connected
// Components of mesh verts conditioned on semantic ID/color.
class CCSemanticObject;
class SemanticRegion;
class SemanticLevel;

//! Represents a scene with containing semantically annotated
//! levels, regions and objects
class SemanticScene {
 public:
  ~SemanticScene() { ESP_DEBUG() << "Deconstructing SemanticScene"; }
  //! return axis aligned bounding box of this House
  Mn::Range3D aabb() const { return bbox_; }

  //! return total number of given element type
  int count(const std::string& element) const {
    return elementCounts_.at(element);
  }

  //! return all SemanticCategories of objects in this House
  const std::vector<std::shared_ptr<SemanticCategory>>& categories() const {
    return categories_;
  }

  //! return all Levels in this House
  const std::vector<std::shared_ptr<SemanticLevel>>& levels() const {
    return levels_;
  }

  //! return all Regions in this House
  const std::vector<std::shared_ptr<SemanticRegion>>& regions() const {
    return regions_;
  }

  //! return all Objects in this House
  const std::vector<std::shared_ptr<SemanticObject>>& objects() const {
    return objects_;
  }

  const std::unordered_map<int, int>& getSemanticIndexMap() const {
    return segmentToObjectIndex_;
  }

  //! convert semantic mesh mask index to object index or ID_UNDEFINED if
  //! not mapped
  inline int semanticIndexToObjectIndex(int maskIndex) const {
    auto segmentToObjIdxIter = segmentToObjectIndex_.find(maskIndex);
    if (segmentToObjIdxIter != segmentToObjectIndex_.end()) {
      return segmentToObjIdxIter->second;
    } else {
      return ID_UNDEFINED;
    }
  }

  /**
   * @brief Attempt to load SemanticScene descriptor from an unknown file type.
   * @param semanticAttr The semantic attributes containing a reference to the
   * name of the semantic scene descriptor (house file) to attempt to load,
   * along with informationd describing semantic constructions (region
   * annotations).
   * @param scene reference to sceneNode to assign semantic scene to
   * @param rotation rotation to apply to semantic scene upon load.
   * @return successfully loaded
   */
  static bool loadSemanticSceneDescriptor(
      const std::shared_ptr<metadata::attributes::SemanticAttributes>&
          semanticAttr,
      SemanticScene& scene,
      const Magnum::Quaternion& rotation =
          Mn::Quaternion::rotation(geo::ESP_FRONT, geo::ESP_GRAVITY));

  /**
   * @brief Attempt to load SemanticScene from a Gibson dataset house format
   * file
   * @param filename the name of the semantic scene descriptor (house file) to
   * attempt to load
   * @param scene reference to sceneNode to assign semantic scene to
   * @param rotation rotation to apply to semantic scene upon load.
   * @return successfully loaded
   */
  static bool loadGibsonHouse(const std::string& filename,
                              SemanticScene& scene,
                              const Magnum::Quaternion& rotation =
                                  Mn::Quaternion::rotation(geo::ESP_FRONT,
                                                           geo::ESP_GRAVITY));

  /**
   * @brief Attempt to load SemanticScene from a HM3D dataset house
   * format file
   * @param filename the name of the semantic scene descriptor (house file) to
   * attempt to load
   * @param scene reference to sceneNode to assign semantic scene to
   * @param rotation rotation to apply to semantic scene upon load (currently
   * not used for HM3D)
   * @return successfully loaded
   */
  static bool loadHM3DHouse(const std::string& filename,
                            SemanticScene& scene,
                            CORRADE_UNUSED const Magnum::Quaternion& rotation =
                                Mn::Quaternion::rotation(geo::ESP_FRONT,
                                                         geo::ESP_GRAVITY));

  /**
   * @brief Attempt to load SemanticScene from a Matterport3D dataset house
   * format file
   * @param filename the name of the semantic scene descriptor (house file) to
   * attempt to load
   * @param scene reference to sceneNode to assign semantic scene to
   * @param rotation rotation to apply to semantic scene upon load.
   * @return successfully loaded
   */
  static bool loadMp3dHouse(const std::string& filename,
                            SemanticScene& scene,
                            const Magnum::Quaternion& rotation =
                                Mn::Quaternion::rotation(geo::ESP_FRONT,
                                                         geo::ESP_GRAVITY));

  /**
   * @brief Attempt to load SemanticScene from a Replica dataset house format
   * file
   * @param filename the name of the semantic scene descriptor (house file) to
   * attempt to load
   * @param scene reference to sceneNode to assign semantic scene to
   * @param rotation rotation to apply to semantic scene upon load.
   * @return successfully loaded
   */
  static bool loadReplicaHouse(const std::string& filename,
                               SemanticScene& scene,
                               const Magnum::Quaternion& rotation =
                                   Mn::Quaternion::rotation(geo::ESP_FRONT,
                                                            geo::ESP_GRAVITY));

  /**
   * @brief Builds a mapping of connected component-driven bounding boxes (via
   * @ref CCSemanticObject), keyed by criteria used to decide connectivity (the
   * per-vertex attribute suche as color).  If a @p semanticScene is passed, key
   * for resultant map will be semanticID of object with specified color,
   * otherwise key is hex color value.
   * @param verts Ref to the vertex buffer holding all vertex positions in the
   * mesh.
   * @param clrsToComponents an unordered map, keyed by tag/color value encoded
   * as uint, where the value is a vector of all sets of CCs consisting of verts
   * with specified tag/"color". (see @ref findCCsByGivenColor).
   * @param semanticScene The SSD for the current semantic mesh.  Used to query
   * semantic objs. If nullptr, this function returns hex-color-keyed map,
   * otherwise returns SemanticID-keyed map.
   * @return A map keyed by a representattion of the per-vertex "color" where
   * each entry contains a vector of values for all the CCs of verts having the
   * "color" attribute specified by the key.  Each element in the vector is a
   * smart pointer to a @ref CCSemanticObject, being used to facilitate
   * collecting pertinent data.
   */
  static std::unordered_map<uint32_t,
                            std::vector<std::shared_ptr<CCSemanticObject>>>
  buildCCBasedSemanticObjs(
      const std::vector<Mn::Vector3>& verts,
      const std::unordered_map<uint32_t, std::vector<std::set<uint32_t>>>&
          clrsToComponents,
      const std::shared_ptr<SemanticScene>& semanticScene);

  /**
   * @brief Build semantic OBBs based on presence of semantic IDs on vertices.
   * @param verts Ref to the vertex buffer holding all vertex positions in the
   * mesh.
   * @param vertSemanticIDs Ref to per-vertex semantic IDs persent on source
   * mesh, both known in semantic scene descriptor, and unknown.  Known IDs are
   * expected to start at 1 and be contiguous, followed by unknown semantic IDs
   * @param ssdObjs The known semantic scene descriptor objects for the mesh
   * @param msgPrefix Debug message prefix, referencing caller.
   * @return vector of semantic object IDXs that have no vertex mapping/presence
   * in the source mesh.
   */
  static std::vector<uint32_t> buildSemanticOBBs(
      const std::vector<Mn::Vector3>& verts,
      const std::vector<uint16_t>& vertSemanticIDs,
      const std::vector<std::shared_ptr<SemanticObject>>& ssdObjs,
      const std::string& msgPrefix);

  /**
   * @brief Build semantic object OBBs based on the accumulated
   * per-semantic-color CCs, and some criteria specified in @p semanticScene .
   * @param verts Ref to the vertex buffer holding all vertex positions in the
   * mesh.
   * @param clrsToComponents an unordered map, keyed by tag/color value encoded
   * as uint, where the value is a vector of all sets of CCs consisting of verts
   * with specified tag/"color". (see @ref findCCsByGivenColor).
   * @param semanticScene The SSD for the current semantic mesh.  Used to query
   * semantic objs. If nullptr, this function returns hex-color-keyed map,
   * otherwise returns SemanticID-keyed map.
   * @param maxVolFraction Fraction of maximum volume bbox CC to include in bbox
   * calc.
   * @param msgPrefix Debug message prefix, referencing caller.
   * @return vector of semantic object IDXs that have no vertex mapping/presence
   * in the source mesh.
   */
  static std::vector<uint32_t> buildSemanticOBBsFromCCs(
      const std::vector<Mn::Vector3>& verts,
      const std::unordered_map<uint32_t, std::vector<std::set<uint32_t>>>&
          clrsToComponents,
      const std::shared_ptr<SemanticScene>& semanticScene,
      float maxVolFraction,
      const std::string& msgPrefix);

  /**
   * @brief Whether the source file assigns colors to verts for Semantic
   * Mesh.
   */
  bool hasVertColorsDefined() const { return hasVertColors_; }

  /**
   * @brief return a read-only reference to the semantic color map, where value
   * is annotation color and index corresponds to id for that color. (HM3D only
   * currently)
   */
  const std::vector<Mn::Vector3ub>& getSemanticColorMap() const {
    return semanticColorMapBeingUsed_;
  }

  /**
   * @brief return a read-only reference to a mapping of semantic
   * color-as-integer to id and region id.(HM3D only currently)
   */
  const std::unordered_map<uint32_t, std::pair<int, int>>&
  getSemanticColorToIdAndRegionMap() const {
    return semanticColorToIdAndRegion_;
  }

  /**
   * @brief whether or not we should build bounding boxes around vertex
   * annotations on semantic asset load. Currently used for HM3D.
   */

  bool buildBBoxFromVertColors() const { return needBBoxFromVertColors_; }

  /**
   * @brief What fraction of largest connected component bbox to use
   * for bbox generation.  Will always use single largest, along with this
   * fraction of remaining. If set to 0.0, use all CCs (i.e. will bypass CC calc
   * and just use naive vert position mechanism for bbox), if set to 1.0, only
   * use single CC with largest volume.  Only used if @ref
   * needBBoxFromVertColors_ is true.
   */
  float CCFractionToUseForBBox() const { return ccLargestVolToUseForBBox_; }

  /**
   * @brief Compute all SemanticRegions which contain the point and return a
   * list of indices for regions in this SemanticScene.
   * @param point The query point.
   * @return A list of indices for regions which contain the point.
   */
  std::vector<int> getRegionsForPoint(const Mn::Vector3& point) const;

  /**
   * @brief Compute all the SemanticRegions that contain the passed point, and
   * return a vector of indices and weights for each region, where the weights
   * are inverted area of the region (smaller regions weighted higher)
   * @param point The query point.
   * @return std::vector<std::pair<int, float>> A sorted list of tuples
   * containing region index and inverse area of that region
   */
  std::vector<std::pair<int, double>> getWeightedRegionsForPoint(
      const Mn::Vector3& point) const;

  /**
   * @brief Compute SemanticRegion containment for a set of points. It is
   * assumed the set of points belong to the same construct (i.e. points from an
   * individual object's mesh)
   * @param points A set of points to test for semantic containment.
   * @return std::vector<std::pair<int, float>> A sorted list of tuples
   * containing region index and percentage of input points contained in that
   * region.
   */
  std::vector<std::pair<int, double>> getRegionsForPoints(
      const std::vector<Mn::Vector3>& points) const;

 protected:
  /**
   * @brief Verify a requested file exists.
   * @param filename the file to attempt to load
   * @param srcFunc calling function name to be displayed in failure message
   * @return whether found or not
   */
  static bool checkFileExists(const std::string& filename,
                              const std::string& srcFunc) {
    if (!Cr::Utility::Path::exists(filename)) {
      ESP_WARNING(Mn::Debug::Flag::NoSpace)
          << "::" << srcFunc << ": File `" << filename
          << "` does not exist.  Aborting load.";
      return false;
    }
    return true;
  }  // checkFileExists

  /**
   * @brief Build the HM3D semantic data from the passed file stream.  File
   * being streamed is expected to be appropriate format.
   * @param ifs The opened file stream describing the HM3D semantic annotations.
   * @param scene reference to sceneNode to assign semantic scene to
   * @param rotation rotation to apply to semantic scene upon load (currently
   * not used for HM3D)
   * @return successfully built. Currently only returns true, but retaining
   * return value for future support.
   */
  static bool buildHM3DHouse(std::ifstream& ifs,
                             SemanticScene& scene,
                             CORRADE_UNUSED const Magnum::Quaternion& rotation =
                                 Mn::Quaternion::rotation(geo::ESP_FRONT,
                                                          geo::ESP_GRAVITY));
  /**
   * @brief Build the mp3 semantic data from the passed file stream. File being
   * streamed is expected to be appropriate format.
   * @param ifs The opened file stream describing the Mp3d semantic annotations.
   * @param scene reference to sceneNode to assign semantic scene to
   * @param rotation rotation to apply to semantic scene upon load.
   * @return successfully built. Currently only returns true, but retaining
   * return value for future support.
   */
  static bool buildMp3dHouse(std::ifstream& ifs,
                             SemanticScene& scene,
                             const Magnum::Quaternion& rotation =
                                 Mn::Quaternion::rotation(geo::ESP_FRONT,
                                                          geo::ESP_GRAVITY));

  /**
   * @brief Build SemanticScene from a Gibson dataset house JSON. JSON is
   * expected to have been verified already.
   * @param jsonDoc the JSON document describing the semantic annotations.
   * @param scene reference to sceneNode to assign semantic scene to
   * @param rotation rotation to apply to semantic scene upon load.
   * @return successfully built. Currently only returns true, but retaining
   * return value for future support.
   */
  static bool buildGibsonHouse(const io::JsonDocument& jsonDoc,
                               SemanticScene& scene,
                               const Magnum::Quaternion& rotation =
                                   Mn::Quaternion::rotation(geo::ESP_FRONT,
                                                            geo::ESP_GRAVITY));

  /**
   * @brief Build SemanticScene from a Replica dataset house JSON. JSON is
   * expected to have been verified already.
   * @param jsonDoc the JSON document describing the semantic annotations.
   * @param scene reference to sceneNode to assign semantic scene to
   * @param objectsExist whether objects cell exists in json. This cell will not
   * exist in ReplicaCAD semantic lexicon.
   * @param rotation rotation to apply to semantic scene upon load.
   * @return successfully built. Currently only returns true, but retaining
   * return value for future support.
   */
  static bool buildReplicaHouse(const io::JsonDocument& jsonDoc,
                                SemanticScene& scene,
                                bool objectsExist,
                                const Magnum::Quaternion& rotation =
                                    Mn::Quaternion::rotation(geo::ESP_FRONT,
                                                             geo::ESP_GRAVITY));

  // Currently only supported by HM3D semantic files.
  bool hasVertColors_ = false;

  /**
   * @Brief whether or not we should build bounding boxes around vertex
   * annotations on semantic asset load. Currently used for HM3D.
   */
  bool needBBoxFromVertColors_ = false;

  /**
   * @brief Fraction of vert count for largest CC to use for BBox synth.  0->use
   * all (bypass CC calc), 1->use only largest CC.
   */
  float ccLargestVolToUseForBBox_ = 0.0f;

  std::string name_;
  std::string label_;
  Mn::Range3D bbox_;
  std::map<std::string, int> elementCounts_;
  std::vector<std::shared_ptr<SemanticCategory>> categories_;
  std::vector<std::shared_ptr<SemanticLevel>> levels_;
  std::vector<std::shared_ptr<SemanticRegion>> regions_;
  std::vector<std::shared_ptr<SemanticObject>> objects_;
  //! map from combined region-segment id to objectIndex for semantic mesh
  std::unordered_map<int, int> segmentToObjectIndex_;
  /**
   * @brief List of mapped vertex colors, where index corresponds to object
   * Index/semantic ID (HM3D only currently)
   */
  std::vector<Mn::Vector3ub> semanticColorMapBeingUsed_{};
  /**
   * @brief Map of integer color to segment and region ids. Used for
   * transforming provided vertex colors. (HM3D only currently)
   */
  std::unordered_map<uint32_t, std::pair<int, int>>
      semanticColorToIdAndRegion_{};

  ESP_SMART_POINTERS(SemanticScene)
};

//! Represents a level of a SemanticScene
class SemanticLevel {
 public:
  virtual ~SemanticLevel() = default;
  virtual std::string id() const { return std::to_string(index_); }
  const std::vector<std::shared_ptr<SemanticRegion>>& regions() const {
    return regions_;
  }

  const std::vector<std::shared_ptr<SemanticObject>>& objects() const {
    return objects_;
  }

  Mn::Range3D aabb() const { return bbox_; }

 protected:
  int index_{};
  std::string labelCode_;
  Mn::Vector3 position_;
  Mn::Range3D bbox_;
  std::vector<std::shared_ptr<SemanticObject>> objects_;
  std::vector<std::shared_ptr<SemanticRegion>> regions_;
  friend SemanticScene;
  ESP_SMART_POINTERS(SemanticLevel)
};

class LoopRegionCategory : public SemanticCategory {
 public:
  LoopRegionCategory(const int id, const std::string& name)
      : id_(id), name_(name) {}

  int index(const std::string& /*mapping*/) const override { return id_; }

  std::string name(const std::string& mapping) const override {
    if (mapping == "category" || mapping.empty()) {
      return name_;
    } else {
      ESP_ERROR() << "Unknown mapping type:" << mapping;
      return "UNKNOWN";
    }
  }

 protected:
  int id_;
  std::string name_;
  ESP_SMART_POINTERS(LoopRegionCategory)

};  // class LoopRegionCategory

//! Represents a region (typically room) in a level of a house
class SemanticRegion {
 public:
  virtual ~SemanticRegion() = default;
  virtual std::string id() const {
    if (!name_.empty()) {
      return name_;
    }
    if (level_ != nullptr) {
      return level_->id() + "_" + std::to_string(index_);
    } else {
      return "_" + std::to_string(index_);
    }
  }

  int getIndex() const { return index_; }
  SemanticLevel::ptr level() const { return level_; }

  const std::vector<std::shared_ptr<SemanticObject>>& objects() const {
    return objects_;
  }

  /**
   * @brief Test whether this region contains the passed point
   */
  virtual bool contains(const Mn::Vector3& point) const;

  void setBBox(const Mn::Vector3& min, const Mn::Vector3& max);

  Mn::Range3D aabb() const { return bbox_; }

  const std::vector<Mn::Vector2>& getPolyLoopPoints() const {
    return polyLoopPoints_;
  }

  /**
   * @brief Return a list of the semantic region's bounding volume edges.
   */
  const std::vector<std::vector<Mn::Vector3>>& getVisEdges() const {
    return visEdges_;
  }

  double getExtrusionHeight() const { return extrusionHeight_; }

  double getFloorHeight() const { return floorHeight_; }

  SemanticCategory::ptr category() const { return category_; }

  /**
   * @brief Returns the area of the polyloop forming the base of the region
   * extrusion
   */
  double getArea() const { return area_; }
  /**
   * @brief Returns the volume of the polyloop-based extrusion defining the
   * bounds of this region.
   */
  double getVolume() const { return area_ * extrusionHeight_; }

 protected:
  int index_{};
  int parentIndex_{};
  std::shared_ptr<SemanticCategory> category_;
  Mn::Vector3 position_;
  Mn::Range3D bbox_;

  std::string name_;

  // The area of the surface enclosed by the region
  double area_{};

  // Height of extrusion for Extruded poly-loop-based volumes
  double extrusionHeight_{};
  // Floor height
  double floorHeight_{};

  // poly loop points for base extrusion
  std::vector<Mn::Vector2> polyLoopPoints_;

  // Edges for visualization of volume
  std::vector<std::vector<Mn::Vector3>> visEdges_;

  std::vector<std::shared_ptr<SemanticObject>> objects_;
  std::shared_ptr<SemanticLevel> level_;
  friend SemanticScene;
  ESP_SMART_POINTERS(SemanticRegion)
};  // class SemanticRegion

//! Represents a distinct semantically annotated object
class SemanticObject {
 public:
  virtual ~SemanticObject() = default;
  virtual std::string id() const {
    if (region_ != nullptr) {
      return region_->id() + "_" + std::to_string(index_);
    } else {
      return "_" + std::to_string(index_);
    }
  }

  /**
   * @brief Retrieve the unique semantic ID corresponding to this object
   */
  int semanticID() const { return index_; }

  SemanticRegion::ptr region() const { return region_; }

  Mn::Range3D aabb() const { return obb_.toAABB(); }

  geo::OBB obb() const { return obb_; }

  SemanticCategory::ptr category() const { return category_; }

  void setObb(
      const Mn::Vector3& center,
      const Mn::Vector3& dimensions,
      const Mn::Quaternion& rotation = Mn::Quaternion(Mn::Math::IdentityInit)) {
    obb_ = geo::OBB{center, dimensions, rotation};
  }
  void setObb(const geo::OBB& otr) { obb_ = geo::OBB{otr}; }
  Mn::Vector3ub getColor() const { return color_; }

  void setColor(Mn::Vector3ub _color) {
    color_ = _color;
    // update colorAsInt_
    colorAsInt_ = geo::getValueAsUInt(color_);
  }

  uint32_t getColorAsInt() const { return colorAsInt_; }

  void setColorAsInt(const uint32_t _colorAsInt) {
    colorAsInt_ = _colorAsInt;
    // update color_ vector
    color_ = {uint8_t((colorAsInt_ >> 16) & 0xff),
              uint8_t((colorAsInt_ >> 8) & 0xff), uint8_t(colorAsInt_ & 0xff)};
  }

 protected:
  /**
   * @brief The unique semantic ID corresponding to this object
   */
  int index_{};
  /**
   * @brief specified color for this object instance.
   */
  Mn::Vector3ub color_{};
  /**
   * @brief  specified color as unsigned int
   */
  uint32_t colorAsInt_{};
  /**
   * @brief References the parent region for this object
   */
  int parentIndex_{};
  std::shared_ptr<SemanticCategory> category_;
  geo::OBB obb_;
  std::shared_ptr<SemanticRegion> region_;
  friend SemanticScene;
  ESP_SMART_POINTERS(SemanticObject)
};  // class SemanticObject

/**
 * @brief This class exists to facilitate semantic object data access for bboxes
 * derived from connected component analysis.
 */

class CCSemanticObject : public SemanticObject {
 public:
  CCSemanticObject(uint32_t _colorInt, const std::set<uint32_t>& _vertSet)
      : SemanticObject(), vertSet_(_vertSet), numSrcVerts_(vertSet_.size()) {
    // need to set index manually when mapping from color/colorInt is known
    index_ = ID_UNDEFINED;
    // sets both colorAsInt_ and updates color_ vector to match
    setColorAsInt(_colorInt);
  }

  void setIndex(int _index) { index_ = _index; }
  uint32_t getNumSrcVerts() const { return numSrcVerts_; }

  const std::set<uint32_t>& getVertSet() const { return vertSet_; }

 protected:
  // reference to vertex set used to build this CCSemantic object
  const std::set<uint32_t>& vertSet_;

  /**
   * @brief Number of verts in source mesh of CC used for this Semantic Object
   */
  uint32_t numSrcVerts_;

  ESP_SMART_POINTERS(CCSemanticObject)
};  // class CCSemanticObject

}  // namespace scene
}  // namespace esp

#endif  // ESP_SCENE_SEMANTICSCENE_H_