smpl_gloss_integration/
conversions.rs

1use burn::tensor::backend::Backend;
2use burn::tensor::{Int, Tensor};
3use gloss_burn_multibackend::backend::MultiBackend;
4use gloss_hecs::{CommandBuffer, Entity};
5use gloss_renderer::components::FacesGPU;
6use gloss_renderer::components::NormalsGPU;
7use gloss_renderer::components::TangentsGPU;
8use gloss_renderer::components::{BoundingBox, ColorsGPU, UVsGPU, VertsGPU};
9use gloss_renderer::{
10    components::{Faces, ModelMatrix, Normals, Tangents, UVs, Verts, VisMesh, VisPoints},
11    scene::Scene,
12};
13use gloss_utils::{
14    bshare::ToNdArray,
15    tensor::{DynamicTensorFloat2D, DynamicTensorInt2D},
16};
17use log::error;
18use nalgebra as na;
19use smpl_core::common::{pose::Pose, smpl_model::SmplModel};
20use wgpu_burn_interop::interop::tensor_float2wgpu_buffer;
21use wgpu_burn_interop::interop::tensor_int2wgpu_buffer;
22/// Insert vertices and vertex attributes for the entity based on changes made
23/// to it, on a generic Burn Backend. We currently support - ``Candle``,
24/// ``NdArray``, and ``Wgpu``
25#[allow(clippy::too_many_arguments)]
26#[allow(clippy::similar_names)]
27#[allow(clippy::too_many_lines)]
28pub fn update_entity_on_backend<B: Backend>(
29    entity: Entity,
30    scene: &Scene,
31    commands: &mut CommandBuffer,
32    with_uv: bool,
33    new_verts: &DynamicTensorFloat2D,
34    new_normals: &DynamicTensorFloat2D,
35    new_tangents: Option<DynamicTensorFloat2D>,
36    uv: DynamicTensorFloat2D,
37    faces: DynamicTensorInt2D,
38    _smpl_model: &dyn SmplModel<B>,
39) {
40    if with_uv && !scene.world.has::<UVs>(entity).unwrap() {
41        commands.insert_one(entity, UVs(uv));
42    }
43    if with_uv {
44        if let Some(tangents) = new_tangents {
45            commands.insert_one(entity, Tangents(tangents.clone()));
46        }
47    }
48    if !scene.world.has::<Faces>(entity).unwrap() {
49        commands.insert_one(entity, Faces(faces));
50    }
51    commands.insert_one(entity, Normals(new_normals.clone()));
52    commands.insert_one(entity, Verts(new_verts.clone()));
53    if !scene.world.has::<VisMesh>(entity).unwrap() {
54        commands.insert_one(
55            entity,
56            VisMesh {
57                added_automatically: true,
58                ..Default::default()
59            },
60        );
61    }
62    if !scene.world.has::<VisPoints>(entity).unwrap() {
63        commands.insert_one(
64            entity,
65            VisPoints {
66                added_automatically: true,
67                ..Default::default()
68            },
69        );
70    }
71    if !scene.world.has::<ModelMatrix>(entity).unwrap() {
72        commands.insert_one(entity, ModelMatrix::default());
73    }
74}
75#[allow(clippy::too_many_arguments)]
76#[allow(clippy::cast_possible_truncation)]
77pub fn update_entity_on_backend_wgpu(
78    entity: Entity,
79    scene: &Scene,
80    gpu: &easy_wgpu::gpu::Gpu,
81    commands: &mut CommandBuffer,
82    with_uv: bool,
83    verts: &Tensor<MultiBackend, 2>,
84    normals: &Tensor<MultiBackend, 2>,
85    tangents: Option<Tensor<MultiBackend, 2>>,
86    uv: &Tensor<MultiBackend, 2>,
87    faces: &Tensor<MultiBackend, 2, Int>,
88) {
89    if !with_uv {
90        error!("UVs are required for WGPU backend. Currently only the case of with_uv is supported");
91        return;
92    }
93    let verts_buf = tensor_float2wgpu_buffer(verts.clone(), wgpu::BufferUsages::VERTEX, &gpu.device().clone(), &gpu.queue().clone());
94    let uv_buf = tensor_float2wgpu_buffer(uv.clone(), wgpu::BufferUsages::VERTEX, &gpu.device().clone(), &gpu.queue().clone());
95    let normals_buf = tensor_float2wgpu_buffer(normals.clone(), wgpu::BufferUsages::VERTEX, &gpu.device().clone(), &gpu.queue().clone());
96    let tangents_buf = tangents
97        .clone()
98        .map(|x| tensor_float2wgpu_buffer(x.clone(), wgpu::BufferUsages::VERTEX, &gpu.device().clone(), &gpu.queue().clone()));
99    let tangents_buf = tangents_buf.unwrap();
100    let faces_buf = tensor_int2wgpu_buffer(faces.clone(), wgpu::BufferUsages::INDEX, &gpu.device().clone(), &gpu.queue().clone());
101    commands.insert_one(
102        entity,
103        VertsGPU {
104            buf: verts_buf.clone(),
105            nr_vertices: verts.shape().dims[0] as u32,
106        },
107    );
108    commands.insert_one(
109        entity,
110        ColorsGPU {
111            buf: verts_buf,
112            nr_vertices: verts.shape().dims[0] as u32,
113        },
114    );
115    commands.insert_one(
116        entity,
117        UVsGPU {
118            buf: uv_buf,
119            nr_vertices: uv.shape().dims[0] as u32,
120        },
121    );
122    commands.insert_one(
123        entity,
124        NormalsGPU {
125            buf: normals_buf,
126            nr_vertices: normals.shape().dims[0] as u32,
127        },
128    );
129    commands.insert_one(
130        entity,
131        TangentsGPU {
132            buf: tangents_buf,
133            nr_vertices: tangents.unwrap().shape().dims[0] as u32,
134        },
135    );
136    commands.insert_one(
137        entity,
138        FacesGPU {
139            buf: faces_buf,
140            nr_triangles: faces.shape().dims[0] as u32,
141        },
142    );
143    if !scene.world.has::<VisPoints>(entity).unwrap() {
144        commands.insert_one(
145            entity,
146            VisPoints {
147                added_automatically: true,
148                ..Default::default()
149            },
150        );
151    }
152    if !scene.world.has::<VisMesh>(entity).unwrap() {
153        commands.insert_one(
154            entity,
155            VisMesh {
156                added_automatically: true,
157                ..Default::default()
158            },
159        );
160    }
161    if let Ok(pose) = scene.get_comp::<&Pose>(&entity) {
162        if !scene.world.has::<BoundingBox>(entity).unwrap() {
163            let mut center = pose.global_trans.clone().to_ndarray();
164            center[1] -= 0.5;
165            let scale = na::Vector3::new(1.5, 1.5, 1.5);
166            let center_point = na::Point3::<f32>::from_slice(center.as_slice().unwrap());
167            let bounding_box = BoundingBox::from_center_and_scale(&center_point, &scale);
168            commands.insert_one(entity, bounding_box);
169        }
170    }
171}