LichtFeld-Studio/src/python/lfs/py_splat_data.hpp at master · MrNeRF/LichtFeld-Studio · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
/* SPDX-FileCopyrightText: 2025 LichtFeld Studio Authors
 *
 * SPDX-License-Identifier: GPL-3.0-or-later */

#pragma once

#include "core/splat_data.hpp"
#include "py_tensor.hpp"
#include <memory>
#include <nanobind/nanobind.h>
#include <optional>

namespace nb = nanobind;

namespace lfs::core {
    class Scene;
}

namespace lfs::python {

    class PySplatData {
    public:
        explicit PySplatData(core::SplatData* data) : data_(data) {
            assert(data_ != nullptr);
        }

        // A scene node's own model: edits through this object report a payload change.
        PySplatData(core::SplatData* data, core::Scene* scene) : data_(data), scene_(scene) {
            assert(data_ != nullptr);
        }

        explicit PySplatData(std::shared_ptr<core::SplatData> owner) : owner_(std::move(owner)), data_(owner_.get()) {
            assert(data_ != nullptr);
        }

        // Raw tensor access (views - no copy)
        PyTensor means_raw() const;
        PyTensor sh0_raw() const;
        PyTensor shN_raw() const;
        PyTensor scaling_raw() const;
        PyTensor rotation_raw() const;
        PyTensor opacity_raw() const;

        // Computed getters (may involve GPU computation)
        PyTensor get_means() const;
        PyTensor get_opacity() const;
        PyTensor get_rotation() const;
        PyTensor get_scaling() const;
        PyTensor get_shs() const;

        // RGB color accessors (handle SH0 encoding internally)
        PyTensor get_colors_rgb() const;
        void set_colors_rgb(const PyTensor& colors);

        // Metadata
        int active_sh_degree() const { return data_->get_active_sh_degree(); }
        int max_sh_degree() const { return data_->get_max_sh_degree(); }
        float scene_scale() const { return data_->get_scene_scale(); }
        size_t num_points() const { return data_->size(); }

        // Soft deletion
        PyTensor deleted() const;
        bool has_deleted_mask() const { return data_->has_deleted_mask(); }
        size_t visible_count() const { return data_->visible_count(); }
        PyTensor soft_delete(const PyTensor& mask);
        void undelete(const PyTensor& mask);
        void clear_deleted() {
            data_->clear_deleted();
            notify_changed();
        }
        size_t apply_deleted() {
            const size_t removed = data_->apply_deleted();
            notify_changed();
            return removed;
        }

        // SH degree management
        void increment_sh_degree() {
            data_->increment_sh_degree();
            notify_changed();
        }
        void set_active_sh_degree(int degree) {
            data_->set_active_sh_degree(degree);
            notify_changed();
        }
        void set_max_sh_degree(int degree) {
            data_->set_max_sh_degree(degree);
            notify_changed();
        }

        // Capacity management
        void reserve_capacity(size_t capacity);

        // Access underlying data (for internal use)
        core::SplatData* data() { return data_; }
        const core::SplatData* data() const { return data_; }

    private:
        void notify_changed() const;

        std::shared_ptr<core::SplatData> owner_;
        core::SplatData* data_;
        core::Scene* scene_ = nullptr;
    };

    // Register PySplatData with nanobind module
    void register_splat_data(nb::module_& m);

} // namespace lfs::python