semantic_kitti_dataset.hpp

Go to the documentation of this file.

#pragma once
 
#include <cpp-toolbox/io/dataset/dataset.hpp>
#include <cpp-toolbox/io/dataset/kitti_types.hpp>
#include <cpp-toolbox/io/dataloader/kitti_odometry_dataset.hpp>
 
#include <memory>
#include <string>
#include <optional>
#include <set>
#include <map>
 
namespace toolbox::io {
 
template<typename DataType = float>
class semantic_kitti_dataset_t : 
    public dataset_t<semantic_kitti_dataset_t<DataType>, 
                     semantic_kitti_frame_t<DataType>> {
public:
    using frame_type = semantic_kitti_frame_t<DataType>;
    using base_type = dataset_t<semantic_kitti_dataset_t<DataType>, frame_type>;
    
    explicit semantic_kitti_dataset_t(const std::string& sequence_path);
    
    void enable_label_validation(bool validate) { 
        validate_labels_ = validate; 
    }
    
    void set_cache_size(std::size_t max_cached_frames) {
        base_dataset_.set_cache_size(max_cached_frames);
        label_cache_size_ = max_cached_frames;
        // Trim label cache if needed
        while (label_cache_.size() > label_cache_size_) {
            label_cache_.erase(label_cache_.begin());
        }
    }
    
    // ==================== Dataset Interface ====================
    
    [[nodiscard]] std::size_t size_impl() const { 
        return base_dataset_.size(); 
    }
    
    [[nodiscard]] std::optional<frame_type> at_impl(std::size_t index) const;
    
    // ==================== Semantic-Specific Methods ====================
    
    [[nodiscard]] std::set<uint16_t> get_unique_labels(bool scan_all = false) const;
    
    [[nodiscard]] std::map<uint16_t, std::string> get_label_names() const {
        return kitti_semantic_labels::get_label_map();
    }
    
    [[nodiscard]] std::map<uint16_t, std::size_t> 
    compute_label_statistics(std::size_t max_frames = 0) const;
    
    [[nodiscard]] std::vector<std::size_t> get_frames_with_label(
        uint16_t label,
        std::size_t min_points = 1) const;
    
    [[nodiscard]] bool has_labels() const { return has_labels_; }
    
    [[nodiscard]] std::string get_label_file(std::size_t index) const {
        if (index >= label_files_.size()) {
            throw kitti_index_out_of_range(index, label_files_.size());
        }
        return label_files_[index];
    }
    
    [[nodiscard]] const kitti_odometry_dataset_t<DataType>& 
    get_base_dataset() const { return base_dataset_; }
    
    void clear_cache() const {
        base_dataset_.clear_cache();
        label_cache_.clear();
        unique_labels_cache_.clear();
    }
 
private:
    void scan_label_files();
    
    std::vector<uint32_t> load_labels(std::size_t index) const;
    
    // ==================== Member Variables ====================
    
    mutable kitti_odometry_dataset_t<DataType> base_dataset_;
    
    std::string labels_path_;
    
    std::vector<std::string> label_files_;
    
    bool has_labels_;
    
    bool validate_labels_ = true;
    
    mutable std::map<std::size_t, std::vector<uint32_t>> label_cache_;
    
    std::size_t label_cache_size_ = 100;
    
    mutable std::set<uint16_t> unique_labels_cache_;
    mutable bool unique_labels_cached_ = false;
};
 
} // namespace toolbox::io
 
// Include implementation
#include <cpp-toolbox/io/dataloader/impl/semantic_kitti_dataset_impl.hpp>