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#[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(¢er_point, &scale);
168 commands.insert_one(entity, bounding_box);
169 }
170 }
171}