smpl_core/common/
outputs.rs

1use crate::AppBackend;
2use burn::{
3    prelude::Backend,
4    tensor::{Float, Int, Tensor},
5};
6use gloss_geometry::geom::{self, PerVertexNormalsWeightingType};
7/// Component for shaped and un-posed mesh. This would be the output of the
8/// ``betas_to_verts`` system. This component is a generic over burn backend
9#[derive(Clone)]
10pub struct SmplOutputPoseTG<B: Backend> {
11    pub verts: Tensor<B, 2, Float>,
12    pub verts_with_expression: Tensor<B, 2, Float>,
13    pub verts_without_expression: Tensor<B, 2, Float>,
14    pub joints: Tensor<B, 2, Float>,
15}
16/// Component for a posed mesh. This would be the output of the ``apply_pose``
17/// system. This component is a generic over burn backend
18#[derive(Clone)]
19pub struct SmplOutputPosedG<B: Backend> {
20    pub joints: Tensor<B, 2, Float>,
21    pub verts: Tensor<B, 2, Float>,
22}
23/// Component for the final shaped and posed mesh. This would be the output of
24/// ``smpl_model.forward()`` This component is a generic over burn backend
25#[derive(Clone)]
26pub struct SmplOutputG<B: Backend> {
27    pub verts: Tensor<B, 2, Float>,
28    pub faces: Tensor<B, 2, Int>,
29    pub normals: Option<Tensor<B, 2, Float>>,
30    pub uvs: Option<Tensor<B, 2, Float>>,
31    pub joints: Tensor<B, 2, Float>,
32}
33impl<B: Backend> SmplOutputG<B> {
34    pub fn compute_normals(&mut self) {
35        let normals = geom::compute_per_vertex_normals_burn(&self.verts, &self.faces, &PerVertexNormalsWeightingType::Area);
36        self.normals = Some(normals);
37    }
38}
39pub type SmplOutputPoseT = SmplOutputPoseTG<AppBackend>;
40pub type SmplOutputPosed = SmplOutputPosedG<AppBackend>;
41pub type SmplOutput = SmplOutputG<AppBackend>;