gffx 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- gffx-0.1.0/LICENSE +21 -0
- gffx-0.1.0/PKG-INFO +84 -0
- gffx-0.1.0/README.md +65 -0
- gffx-0.1.0/pyproject.toml +18 -0
- gffx-0.1.0/src/gffx/__init__.py +5 -0
- gffx-0.1.0/src/gffx/__main__.py +172 -0
- gffx-0.1.0/src/gffx/context/__init__.py +14 -0
- gffx-0.1.0/src/gffx/linalg/__init__.py +158 -0
- gffx-0.1.0/src/gffx/obj/__init__.py +367 -0
- gffx-0.1.0/src/gffx/ray/__init__.py +284 -0
- gffx-0.1.0/src/gffx/ray/camera.py +66 -0
gffx-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 Danzel Serrano
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
gffx-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,84 @@
|
|
|
1
|
+
Metadata-Version: 2.3
|
|
2
|
+
Name: gffx
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: A minimal graphics library in PyTorch
|
|
5
|
+
License: MIT
|
|
6
|
+
Author: Danzel Serrano
|
|
7
|
+
Author-email: danzel.serrano@gmail.com
|
|
8
|
+
Requires-Python: >=3.8,<4.0
|
|
9
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
10
|
+
Classifier: Programming Language :: Python :: 3
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.8
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
17
|
+
Description-Content-Type: text/markdown
|
|
18
|
+
|
|
19
|
+
### gffx - a minimal library for (differentiable) graphics
|
|
20
|
+
|
|
21
|
+
```Python
|
|
22
|
+
import gffx
|
|
23
|
+
import torch
|
|
24
|
+
import matplotlib.pyplot as plt
|
|
25
|
+
|
|
26
|
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
27
|
+
|
|
28
|
+
# Setup Camera
|
|
29
|
+
camera = gffx.ray.Camera(
|
|
30
|
+
width = 512,
|
|
31
|
+
height = 512,
|
|
32
|
+
pos = [0, 0, 3],
|
|
33
|
+
dir = [0, 0, -1],
|
|
34
|
+
|
|
35
|
+
device = device
|
|
36
|
+
)
|
|
37
|
+
|
|
38
|
+
# Setup list of meshes
|
|
39
|
+
vertices = torch.load('path/to/vertices.pt').to(device)
|
|
40
|
+
faces = torch.load('path/to/vertices.pt').to(device)
|
|
41
|
+
|
|
42
|
+
object_list = [
|
|
43
|
+
gffx.obj.mesh_from_vertices_and_faces(
|
|
44
|
+
vertices = vertices,
|
|
45
|
+
faces = faces,
|
|
46
|
+
init_translation = [0, 0.0, 2],
|
|
47
|
+
init_rotation = [0, 0, 0],
|
|
48
|
+
init_scale = [1, 1, 1],
|
|
49
|
+
|
|
50
|
+
ambient_color = [0.5, 0.5, 0.5],
|
|
51
|
+
diffuse_color = [0.5, 0.5, 0.5],
|
|
52
|
+
specular_color = [0.5, 0.5, 0.5],
|
|
53
|
+
specular_coefficient = 1,
|
|
54
|
+
|
|
55
|
+
device = device
|
|
56
|
+
)
|
|
57
|
+
]
|
|
58
|
+
|
|
59
|
+
# Ray Trace Render
|
|
60
|
+
images = gffx.ray.mesh_render(
|
|
61
|
+
meshes = object_list,
|
|
62
|
+
camera = camera,
|
|
63
|
+
light_intensity = 1.0,
|
|
64
|
+
ambient_intensity = 0.2,
|
|
65
|
+
light_pos = [5, 5, 5],
|
|
66
|
+
background_color = [0, 0, 0],
|
|
67
|
+
ray_chunk_size = 4096,
|
|
68
|
+
|
|
69
|
+
device = device
|
|
70
|
+
)
|
|
71
|
+
|
|
72
|
+
plt.imshow((images[0].cpu()).permute(1, 0, 2))
|
|
73
|
+
plt.gca().invert_yaxis()
|
|
74
|
+
plt.show()
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
### [WIP] CLI
|
|
78
|
+
```
|
|
79
|
+
python -m gffx --vertices '/path/to/vertices.pt' --faces 'path/to/faces.pt'
|
|
80
|
+
```
|
|
81
|
+
|
|
82
|
+
### Motivation
|
|
83
|
+
1. Installing PyTorch3D everytime in colab is a hassle. This library's first aim is to render a mesh using native operations in PyTorch.
|
|
84
|
+
|
gffx-0.1.0/README.md
ADDED
|
@@ -0,0 +1,65 @@
|
|
|
1
|
+
### gffx - a minimal library for (differentiable) graphics
|
|
2
|
+
|
|
3
|
+
```Python
|
|
4
|
+
import gffx
|
|
5
|
+
import torch
|
|
6
|
+
import matplotlib.pyplot as plt
|
|
7
|
+
|
|
8
|
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
9
|
+
|
|
10
|
+
# Setup Camera
|
|
11
|
+
camera = gffx.ray.Camera(
|
|
12
|
+
width = 512,
|
|
13
|
+
height = 512,
|
|
14
|
+
pos = [0, 0, 3],
|
|
15
|
+
dir = [0, 0, -1],
|
|
16
|
+
|
|
17
|
+
device = device
|
|
18
|
+
)
|
|
19
|
+
|
|
20
|
+
# Setup list of meshes
|
|
21
|
+
vertices = torch.load('path/to/vertices.pt').to(device)
|
|
22
|
+
faces = torch.load('path/to/vertices.pt').to(device)
|
|
23
|
+
|
|
24
|
+
object_list = [
|
|
25
|
+
gffx.obj.mesh_from_vertices_and_faces(
|
|
26
|
+
vertices = vertices,
|
|
27
|
+
faces = faces,
|
|
28
|
+
init_translation = [0, 0.0, 2],
|
|
29
|
+
init_rotation = [0, 0, 0],
|
|
30
|
+
init_scale = [1, 1, 1],
|
|
31
|
+
|
|
32
|
+
ambient_color = [0.5, 0.5, 0.5],
|
|
33
|
+
diffuse_color = [0.5, 0.5, 0.5],
|
|
34
|
+
specular_color = [0.5, 0.5, 0.5],
|
|
35
|
+
specular_coefficient = 1,
|
|
36
|
+
|
|
37
|
+
device = device
|
|
38
|
+
)
|
|
39
|
+
]
|
|
40
|
+
|
|
41
|
+
# Ray Trace Render
|
|
42
|
+
images = gffx.ray.mesh_render(
|
|
43
|
+
meshes = object_list,
|
|
44
|
+
camera = camera,
|
|
45
|
+
light_intensity = 1.0,
|
|
46
|
+
ambient_intensity = 0.2,
|
|
47
|
+
light_pos = [5, 5, 5],
|
|
48
|
+
background_color = [0, 0, 0],
|
|
49
|
+
ray_chunk_size = 4096,
|
|
50
|
+
|
|
51
|
+
device = device
|
|
52
|
+
)
|
|
53
|
+
|
|
54
|
+
plt.imshow((images[0].cpu()).permute(1, 0, 2))
|
|
55
|
+
plt.gca().invert_yaxis()
|
|
56
|
+
plt.show()
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
### [WIP] CLI
|
|
60
|
+
```
|
|
61
|
+
python -m gffx --vertices '/path/to/vertices.pt' --faces 'path/to/faces.pt'
|
|
62
|
+
```
|
|
63
|
+
|
|
64
|
+
### Motivation
|
|
65
|
+
1. Installing PyTorch3D everytime in colab is a hassle. This library's first aim is to render a mesh using native operations in PyTorch.
|
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["poetry-core>=1.0.0"]
|
|
3
|
+
build-backend = "poetry.core.masonry.api"
|
|
4
|
+
|
|
5
|
+
[tool.poetry]
|
|
6
|
+
name = "gffx"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "A minimal graphics library in PyTorch"
|
|
9
|
+
authors = ["Danzel Serrano <danzel.serrano@gmail.com>"]
|
|
10
|
+
license = "MIT" # Change as needed
|
|
11
|
+
readme = "README.md"
|
|
12
|
+
|
|
13
|
+
[tool.poetry.dependencies]
|
|
14
|
+
python = "^3.8"
|
|
15
|
+
|
|
16
|
+
[tool.poetry.group.dev.dependencies]
|
|
17
|
+
pytest = "^6.2.5"
|
|
18
|
+
coverage = "^6.2.0"
|
|
@@ -0,0 +1,172 @@
|
|
|
1
|
+
import gffx
|
|
2
|
+
import torch
|
|
3
|
+
import argparse
|
|
4
|
+
from pathlib import Path
|
|
5
|
+
import matplotlib.pyplot as plt
|
|
6
|
+
|
|
7
|
+
################################################################
|
|
8
|
+
# ARGUMENT PARSER
|
|
9
|
+
################################################################
|
|
10
|
+
parser = argparse.ArgumentParser(description='gffx')
|
|
11
|
+
parser.add_argument(
|
|
12
|
+
'--vertices',
|
|
13
|
+
type=Path,
|
|
14
|
+
help='Path to vertices PyTorch file (*.pt)',
|
|
15
|
+
required=True
|
|
16
|
+
)
|
|
17
|
+
parser.add_argument(
|
|
18
|
+
'--faces',
|
|
19
|
+
type=Path,
|
|
20
|
+
help='Path to faces PyTorch file (*.pt)',
|
|
21
|
+
required=True
|
|
22
|
+
)
|
|
23
|
+
parser.add_argument(
|
|
24
|
+
'--center_and_scale',
|
|
25
|
+
action = 'store_true',
|
|
26
|
+
help = 'Center and scale object',
|
|
27
|
+
default = False
|
|
28
|
+
)
|
|
29
|
+
|
|
30
|
+
parser.add_argument(
|
|
31
|
+
'--camera_width',
|
|
32
|
+
type=int,
|
|
33
|
+
help='Camera width',
|
|
34
|
+
default=512
|
|
35
|
+
)
|
|
36
|
+
parser.add_argument(
|
|
37
|
+
'--camera_height',
|
|
38
|
+
type=int,
|
|
39
|
+
help='Camera height',
|
|
40
|
+
default=512
|
|
41
|
+
)
|
|
42
|
+
parser.add_argument(
|
|
43
|
+
'--camera_pos',
|
|
44
|
+
type=list,
|
|
45
|
+
help='Camera position',
|
|
46
|
+
default=[0, 0, 1.2]
|
|
47
|
+
)
|
|
48
|
+
parser.add_argument(
|
|
49
|
+
'--camera_dir',
|
|
50
|
+
type=list,
|
|
51
|
+
help='Camera direction',
|
|
52
|
+
default=[0, 0, -1]
|
|
53
|
+
)
|
|
54
|
+
parser.add_argument(
|
|
55
|
+
'--camera_screen_distance',
|
|
56
|
+
type = float,
|
|
57
|
+
help = 'Camera screen distance',
|
|
58
|
+
default = 1.0
|
|
59
|
+
)
|
|
60
|
+
|
|
61
|
+
parser.add_argument(
|
|
62
|
+
'--ambient_color',
|
|
63
|
+
type=list,
|
|
64
|
+
help='Ambient color',
|
|
65
|
+
default=[0.5, 0.5, 0.5]
|
|
66
|
+
)
|
|
67
|
+
parser.add_argument(
|
|
68
|
+
'--diffuse_color',
|
|
69
|
+
type=list,
|
|
70
|
+
help='Diffuse color',
|
|
71
|
+
default=[0.5, 0.5, 0.5]
|
|
72
|
+
)
|
|
73
|
+
parser.add_argument(
|
|
74
|
+
'--specular_color',
|
|
75
|
+
type=list,
|
|
76
|
+
help='Specular color',
|
|
77
|
+
default=[0.5, 0.5, 0.5]
|
|
78
|
+
)
|
|
79
|
+
parser.add_argument(
|
|
80
|
+
'--specular_coefficient',
|
|
81
|
+
type=float,
|
|
82
|
+
help='Specular coefficient',
|
|
83
|
+
default=1
|
|
84
|
+
)
|
|
85
|
+
parser.add_argument(
|
|
86
|
+
'--light_intensity',
|
|
87
|
+
type=float,
|
|
88
|
+
help='Light intensity',
|
|
89
|
+
default=1
|
|
90
|
+
)
|
|
91
|
+
parser.add_argument(
|
|
92
|
+
'--ambient_intensity',
|
|
93
|
+
type=float,
|
|
94
|
+
help='Ambient intensity',
|
|
95
|
+
default=0.2
|
|
96
|
+
)
|
|
97
|
+
parser.add_argument(
|
|
98
|
+
'--light_pos',
|
|
99
|
+
type=list,
|
|
100
|
+
help='Light position',
|
|
101
|
+
default=[5, 5, 5]
|
|
102
|
+
)
|
|
103
|
+
parser.add_argument(
|
|
104
|
+
'--background_color',
|
|
105
|
+
type=list,
|
|
106
|
+
help='Background color',
|
|
107
|
+
default=[0, 0, 0]
|
|
108
|
+
)
|
|
109
|
+
parser.add_argument(
|
|
110
|
+
'--ray_chunk_size',
|
|
111
|
+
type=int,
|
|
112
|
+
help='Ray chunk size',
|
|
113
|
+
default=1024
|
|
114
|
+
)
|
|
115
|
+
|
|
116
|
+
args = parser.parse_args()
|
|
117
|
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
118
|
+
|
|
119
|
+
# Setup Camera
|
|
120
|
+
camera = gffx.ray.Camera(
|
|
121
|
+
width = args.camera_width,
|
|
122
|
+
height = args.camera_height,
|
|
123
|
+
pos = args.camera_pos,
|
|
124
|
+
dir = args.camera_dir,
|
|
125
|
+
screen_distance = args.camera_screen_distance,
|
|
126
|
+
|
|
127
|
+
device = device
|
|
128
|
+
)
|
|
129
|
+
|
|
130
|
+
# Setup list of meshes
|
|
131
|
+
vertices = torch.load(args.vertices)[0].to(device)
|
|
132
|
+
faces = torch.load(args.faces)[0].to(device)
|
|
133
|
+
|
|
134
|
+
if args.center_and_scale:
|
|
135
|
+
vertices_mean = torch.mean(vertices, dim=0)
|
|
136
|
+
vertices_centered = vertices - vertices_mean
|
|
137
|
+
vertices_scale = torch.max(torch.linalg.norm(vertices_centered, dim=-1))
|
|
138
|
+
vertices = vertices / vertices_scale
|
|
139
|
+
|
|
140
|
+
object_list = [
|
|
141
|
+
gffx.obj.mesh_from_vertices_and_faces(
|
|
142
|
+
vertices = vertices,
|
|
143
|
+
faces = faces,
|
|
144
|
+
init_translation = [0, 0.0, 0],
|
|
145
|
+
init_rotation = [0, 0, 0],
|
|
146
|
+
init_scale = [1, 1, 1],
|
|
147
|
+
|
|
148
|
+
ambient_color = args.ambient_color,
|
|
149
|
+
diffuse_color = args.diffuse_color,
|
|
150
|
+
specular_color = args.specular_color,
|
|
151
|
+
specular_coefficient = args.specular_coefficient,
|
|
152
|
+
|
|
153
|
+
device = device
|
|
154
|
+
)
|
|
155
|
+
]
|
|
156
|
+
|
|
157
|
+
# Ray Trace Render
|
|
158
|
+
images = gffx.ray.mesh_render(
|
|
159
|
+
meshes = object_list,
|
|
160
|
+
camera = camera,
|
|
161
|
+
light_intensity = args.light_intensity,
|
|
162
|
+
ambient_intensity = args.ambient_intensity,
|
|
163
|
+
light_pos = args.light_pos,
|
|
164
|
+
background_color = args.background_color,
|
|
165
|
+
ray_chunk_size = args.ray_chunk_size,
|
|
166
|
+
|
|
167
|
+
device = device
|
|
168
|
+
)
|
|
169
|
+
|
|
170
|
+
plt.imshow((images[0].cpu()).permute(1, 0, 2))
|
|
171
|
+
plt.gca().invert_yaxis()
|
|
172
|
+
plt.show()
|
|
@@ -0,0 +1,158 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
|
|
3
|
+
def cramer(
|
|
4
|
+
A : torch.Tensor,
|
|
5
|
+
b : torch.Tensor
|
|
6
|
+
) -> torch.Tensor:
|
|
7
|
+
"""
|
|
8
|
+
Solve Ax = b via Cramer's Rule
|
|
9
|
+
|
|
10
|
+
Parameters
|
|
11
|
+
----------
|
|
12
|
+
A : torch.Tensor
|
|
13
|
+
Coefficient matrices -> dim(..., N, N)
|
|
14
|
+
b : torch.Tensor
|
|
15
|
+
Right-hand side vectors -> dim(..., N, 1)
|
|
16
|
+
"""
|
|
17
|
+
# Get shape for reshape later
|
|
18
|
+
A_shape = A.shape
|
|
19
|
+
|
|
20
|
+
# Flatten A and b
|
|
21
|
+
A = A.flatten(start_dim=0, end_dim=-3)
|
|
22
|
+
b = b.flatten(start_dim=0, end_dim=-3)
|
|
23
|
+
device = A.device
|
|
24
|
+
|
|
25
|
+
# Asserts
|
|
26
|
+
B, N, _ = A.shape
|
|
27
|
+
assert N == A.shape[2], "Matrices A must be square"
|
|
28
|
+
assert B == b.shape[0], "Matrices A and vectors b must have the same batch size"
|
|
29
|
+
assert N == b.shape[1], "Matrices A and vectors b must have the same number of rows"
|
|
30
|
+
assert b.shape[2] == 1, "Vectors b must be a column vector"
|
|
31
|
+
|
|
32
|
+
# Compute determinant of A
|
|
33
|
+
M = torch.linalg.det(A) # dim(B, 1)
|
|
34
|
+
|
|
35
|
+
# Compute determinants of A_i
|
|
36
|
+
A_dets = torch.zeros((B, N), device=device)
|
|
37
|
+
for i in range(N):
|
|
38
|
+
A_i = A.clone()
|
|
39
|
+
A_i[:, :, i] = b.squeeze(-1)
|
|
40
|
+
A_dets[:, i] = torch.linalg.det(A_i)
|
|
41
|
+
|
|
42
|
+
# Compute solution
|
|
43
|
+
x = A_dets / M.unsqueeze(-1)
|
|
44
|
+
|
|
45
|
+
# Reshape x
|
|
46
|
+
x = torch.reshape(x, [*A_shape[:-2], N, 1]) # dim(..., N, 1)
|
|
47
|
+
|
|
48
|
+
return x
|
|
49
|
+
|
|
50
|
+
def euler_angles_to_matrix(
|
|
51
|
+
rot_vec : torch.Tensor
|
|
52
|
+
):
|
|
53
|
+
"""
|
|
54
|
+
Convert Euler angles to rotation matrix
|
|
55
|
+
|
|
56
|
+
Parameters
|
|
57
|
+
----------
|
|
58
|
+
rot_vec : torch.Tensor
|
|
59
|
+
Rotation vector -> dim(B, 3)
|
|
60
|
+
"""
|
|
61
|
+
B = rot_vec.shape[0]
|
|
62
|
+
device = rot_vec.device
|
|
63
|
+
|
|
64
|
+
# Rodrigues' rotation formula
|
|
65
|
+
theta = torch.linalg.norm(rot_vec, dim=-1, keepdim=True)[..., None]
|
|
66
|
+
K = torch.zeros((B, 3, 3), device=device)
|
|
67
|
+
K[:,0,1] = -rot_vec[:,2]
|
|
68
|
+
K[:,0,2] = rot_vec[:,1]
|
|
69
|
+
K[:,1,0] = rot_vec[:,2]
|
|
70
|
+
K[:,1,2] = -rot_vec[:,0]
|
|
71
|
+
K[:,2,0] = -rot_vec[:,1]
|
|
72
|
+
K[:,2,1] = rot_vec[:,0]
|
|
73
|
+
|
|
74
|
+
R = (
|
|
75
|
+
torch.eye(3, device=device)[None,:].expand(B,-1,-1)
|
|
76
|
+
+ torch.sin(theta) * K
|
|
77
|
+
+ (1 - torch.cos(theta)) * (K @ K)
|
|
78
|
+
)
|
|
79
|
+
R = torch.cat([R, torch.zeros((B, 3, 1), device=device)], dim=-1)
|
|
80
|
+
R = torch.cat([R, torch.zeros((B, 1, 4), device=device)], dim=-2)
|
|
81
|
+
R[:,3,3] = 1
|
|
82
|
+
|
|
83
|
+
return R
|
|
84
|
+
|
|
85
|
+
def scale_matrix(
|
|
86
|
+
scale_vec : torch.Tensor
|
|
87
|
+
):
|
|
88
|
+
"""
|
|
89
|
+
Convert scale vector to scale matrix
|
|
90
|
+
|
|
91
|
+
Parameters
|
|
92
|
+
----------
|
|
93
|
+
scale_vec : torch.Tensor
|
|
94
|
+
Scale vector -> dim(B, 3)
|
|
95
|
+
"""
|
|
96
|
+
B = scale_vec.shape[0]
|
|
97
|
+
device = scale_vec.device
|
|
98
|
+
|
|
99
|
+
# Scale matrix
|
|
100
|
+
return torch.cat([
|
|
101
|
+
torch.cat([
|
|
102
|
+
torch.eye(3, device=device)[None,:].expand(B,-1,-1) * scale_vec[..., None],
|
|
103
|
+
torch.zeros((B, 3, 1), device=device)
|
|
104
|
+
], dim=-1),
|
|
105
|
+
torch.cat([
|
|
106
|
+
torch.zeros((B, 1, 3), device=device),
|
|
107
|
+
torch.ones((B, 1, 1), device=device)
|
|
108
|
+
], dim=-1)
|
|
109
|
+
], dim=-2)
|
|
110
|
+
|
|
111
|
+
def translation_matrix(
|
|
112
|
+
translation_vec : torch.Tensor
|
|
113
|
+
):
|
|
114
|
+
"""
|
|
115
|
+
Convert translation vector to translation matrix
|
|
116
|
+
|
|
117
|
+
Parameters
|
|
118
|
+
----------
|
|
119
|
+
translation_vec : torch.Tensor
|
|
120
|
+
Translation vector -> dim(B, 3)
|
|
121
|
+
"""
|
|
122
|
+
B = translation_vec.shape[0]
|
|
123
|
+
device = translation_vec.device
|
|
124
|
+
|
|
125
|
+
# Translation matrix
|
|
126
|
+
T = torch.eye(4, device=device)[None,:].expand(B,-1,-1).clone()
|
|
127
|
+
T[:,0:3,3] = translation_vec
|
|
128
|
+
|
|
129
|
+
return T
|
|
130
|
+
|
|
131
|
+
def transformation_matrix(
|
|
132
|
+
translation_vec : torch.Tensor,
|
|
133
|
+
rotation_vec : torch.Tensor,
|
|
134
|
+
scale_vec : torch.Tensor
|
|
135
|
+
):
|
|
136
|
+
"""
|
|
137
|
+
Convert translation, rotation, and scale vectors to transformation matrix
|
|
138
|
+
|
|
139
|
+
Parameters
|
|
140
|
+
----------
|
|
141
|
+
translation_vec : torch.Tensor
|
|
142
|
+
Translation vector -> dim(B, 3)
|
|
143
|
+
rotation_vec : torch.Tensor
|
|
144
|
+
Rotation vector -> dim(B, 3)
|
|
145
|
+
scale_vec : torch.Tensor
|
|
146
|
+
Scale vector -> dim(B, 3)
|
|
147
|
+
"""
|
|
148
|
+
|
|
149
|
+
# Rodrigues' rotation formula
|
|
150
|
+
R = euler_angles_to_matrix(rotation_vec)
|
|
151
|
+
|
|
152
|
+
# Scale Matrix
|
|
153
|
+
S = scale_matrix(scale_vec)
|
|
154
|
+
|
|
155
|
+
# Translation Matrix
|
|
156
|
+
T = translation_matrix(translation_vec)
|
|
157
|
+
|
|
158
|
+
return T @ R @ S
|
|
@@ -0,0 +1,367 @@
|
|
|
1
|
+
import math
|
|
2
|
+
import gffx
|
|
3
|
+
import torch
|
|
4
|
+
from typing import Optional
|
|
5
|
+
|
|
6
|
+
class Transform:
|
|
7
|
+
def __init__(
|
|
8
|
+
self,
|
|
9
|
+
translation : Optional[list | torch.Tensor] = None,
|
|
10
|
+
rotation : Optional[list | torch.Tensor] = None,
|
|
11
|
+
scale : Optional[list | torch.Tensor] = None,
|
|
12
|
+
device : Optional[torch.device] = None
|
|
13
|
+
):
|
|
14
|
+
if device is None:
|
|
15
|
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
16
|
+
|
|
17
|
+
# Defaults
|
|
18
|
+
if not isinstance(translation, torch.Tensor) and not translation:
|
|
19
|
+
translation = [0., 0., 0.]
|
|
20
|
+
if isinstance(translation, list):
|
|
21
|
+
translation = torch.tensor(translation, device=device)[None,:]
|
|
22
|
+
|
|
23
|
+
if not isinstance(rotation, torch.Tensor) and not rotation:
|
|
24
|
+
rotation = [0., 0., 0.]
|
|
25
|
+
if isinstance(rotation, list):
|
|
26
|
+
rotation = torch.tensor(rotation, device=device)[None,:]
|
|
27
|
+
|
|
28
|
+
if not isinstance(scale, torch.Tensor) and not scale:
|
|
29
|
+
scale = [1., 1., 1.]
|
|
30
|
+
if isinstance(scale, list):
|
|
31
|
+
scale = torch.tensor(scale, device=device)[None,:]
|
|
32
|
+
|
|
33
|
+
#
|
|
34
|
+
self.translation = translation.float()
|
|
35
|
+
self.rotation = rotation.float()
|
|
36
|
+
self.scale = scale.float()
|
|
37
|
+
|
|
38
|
+
class MeshObject:
|
|
39
|
+
def __init__(
|
|
40
|
+
self,
|
|
41
|
+
vertices : torch.Tensor,
|
|
42
|
+
faces : torch.Tensor,
|
|
43
|
+
normals : Optional[torch.Tensor] = None,
|
|
44
|
+
init_transform : Optional[Transform] = None,
|
|
45
|
+
init_translation : Optional[torch.Tensor] = None,
|
|
46
|
+
init_rotation : Optional[torch.Tensor] = None,
|
|
47
|
+
init_scale : Optional[torch.Tensor] = None,
|
|
48
|
+
|
|
49
|
+
ambient_color : list | torch.Tensor = [0.5, 0.5, 0.5],
|
|
50
|
+
diffuse_color : list | torch.Tensor = [0.5, 0.5, 0.5],
|
|
51
|
+
specular_color : list | torch.Tensor = [0.5, 0.5, 0.5],
|
|
52
|
+
specular_coefficient : float = 1.0
|
|
53
|
+
):
|
|
54
|
+
# Defaults
|
|
55
|
+
if init_transform:
|
|
56
|
+
assert init_translation is None, "If init_transform is not None, init_translation must be None"
|
|
57
|
+
assert init_rotation is None, "If init_transform is not None, init_rotation must be None"
|
|
58
|
+
assert init_scale is None, "If init_transform is not None, init_scale must be None"
|
|
59
|
+
self.transform = init_transform
|
|
60
|
+
else:
|
|
61
|
+
self.transform = Transform(
|
|
62
|
+
translation = init_translation,
|
|
63
|
+
rotation = init_rotation,
|
|
64
|
+
scale = init_scale
|
|
65
|
+
)
|
|
66
|
+
|
|
67
|
+
self.vertices = vertices
|
|
68
|
+
self.faces = faces
|
|
69
|
+
self.normals = normals
|
|
70
|
+
self.device = vertices.device
|
|
71
|
+
|
|
72
|
+
# [4NOW] Flat colors
|
|
73
|
+
self.ambient_color = ambient_color
|
|
74
|
+
if isinstance(self.ambient_color, list):
|
|
75
|
+
self.ambient_color = torch.tensor(self.ambient_color, device=self.device, dtype=torch.float32)
|
|
76
|
+
|
|
77
|
+
self.diffuse_color = diffuse_color
|
|
78
|
+
if isinstance(self.diffuse_color, list):
|
|
79
|
+
self.diffuse_color = torch.tensor(self.diffuse_color, device=self.device, dtype=torch.float32)
|
|
80
|
+
|
|
81
|
+
self.specular_color = specular_color
|
|
82
|
+
if isinstance(self.specular_color, list):
|
|
83
|
+
self.specular_color = torch.tensor(self.specular_color, device=self.device, dtype=torch.float32)
|
|
84
|
+
|
|
85
|
+
self.specular_coefficient = specular_coefficient
|
|
86
|
+
|
|
87
|
+
def get_transformed(self):
|
|
88
|
+
"""
|
|
89
|
+
Apply transformation to vertices
|
|
90
|
+
|
|
91
|
+
Returns:
|
|
92
|
+
-------
|
|
93
|
+
torch.Tensor: Transformed vertices
|
|
94
|
+
|
|
95
|
+
if normals are not None, also returns transformed normals
|
|
96
|
+
"""
|
|
97
|
+
M = gffx.linalg.transformation_matrix(
|
|
98
|
+
translation_vec = self.transform.translation[None,...],
|
|
99
|
+
rotation_vec = self.transform.rotation[None,...],
|
|
100
|
+
scale_vec = self.transform.scale[None,...]
|
|
101
|
+
) # dim(1, 4, 4)
|
|
102
|
+
vertices_h = torch.cat([self.vertices, torch.ones((self.vertices.shape[0], 1), device=self.device)], dim=-1)
|
|
103
|
+
vertices_h = vertices_h @ M.transpose(-1, -2)
|
|
104
|
+
|
|
105
|
+
# [Optional] Transform normals
|
|
106
|
+
if self.normals is not None:
|
|
107
|
+
M = M[:, 0:3, 0:3].inverse()
|
|
108
|
+
normals = self.normals @ M
|
|
109
|
+
return vertices_h[0,..., 0:3], normals[0]
|
|
110
|
+
|
|
111
|
+
return vertices_h[0,..., 0:3], None
|
|
112
|
+
|
|
113
|
+
def compute_normals(vertices: torch.Tensor, faces: torch.Tensor) -> torch.Tensor:
|
|
114
|
+
"""
|
|
115
|
+
Compute per-vertex normals for a batch of meshes.
|
|
116
|
+
|
|
117
|
+
Args:
|
|
118
|
+
vertices (torch.Tensor): Tensor of shape (B, V, 3) with vertex positions.
|
|
119
|
+
faces (torch.Tensor): Long tensor of shape (B, F, 3) with indices into vertices.
|
|
120
|
+
|
|
121
|
+
Returns:
|
|
122
|
+
torch.Tensor: Tensor of shape (B, V, 3) with normalized vertex normals.
|
|
123
|
+
"""
|
|
124
|
+
B, V, _ = vertices.shape
|
|
125
|
+
_, F, _ = faces.shape
|
|
126
|
+
device = vertices.device
|
|
127
|
+
|
|
128
|
+
# Initialize normals accumulator
|
|
129
|
+
vertex_normals = torch.zeros_like(vertices)
|
|
130
|
+
|
|
131
|
+
# Gather vertex positions for each corner of every face
|
|
132
|
+
v0 = torch.gather(vertices, 1, faces[:, :, 0].unsqueeze(-1).expand(-1, -1, 3))
|
|
133
|
+
v1 = torch.gather(vertices, 1, faces[:, :, 1].unsqueeze(-1).expand(-1, -1, 3))
|
|
134
|
+
v2 = torch.gather(vertices, 1, faces[:, :, 2].unsqueeze(-1).expand(-1, -1, 3))
|
|
135
|
+
|
|
136
|
+
# Compute face normals (unnormalized) via cross product
|
|
137
|
+
face_normals = torch.cross(v1 - v0, v2 - v0, dim=2) # dim(B,F,3)
|
|
138
|
+
face_normals = face_normals / (face_normals.norm(dim=2, keepdim=True) + 1e-8)
|
|
139
|
+
|
|
140
|
+
# Accumulate face normals to each vertex using scatter_add_
|
|
141
|
+
for i in range(3):
|
|
142
|
+
idx = faces[:, :, i].unsqueeze(-1).expand(-1, -1, 3)
|
|
143
|
+
vertex_normals.scatter_add_(1, idx, face_normals)
|
|
144
|
+
|
|
145
|
+
# Normalize the accumulated vertex normals
|
|
146
|
+
norm = vertex_normals.norm(dim=2, keepdim=True).clamp(min=1e-8)
|
|
147
|
+
vertex_normals = vertex_normals / norm
|
|
148
|
+
|
|
149
|
+
return vertex_normals
|
|
150
|
+
|
|
151
|
+
def mesh_from_vertices_and_faces(
|
|
152
|
+
vertices : torch.Tensor,
|
|
153
|
+
faces : torch.Tensor,
|
|
154
|
+
init_transform : Optional[Transform] = None,
|
|
155
|
+
init_translation : Optional[list | torch.Tensor] = None,
|
|
156
|
+
init_rotation : Optional[list | torch.Tensor] = None,
|
|
157
|
+
init_scale : Optional[list | torch.Tensor] = None,
|
|
158
|
+
|
|
159
|
+
ambient_color : list | torch.Tensor = [0.5, 0.5, 0.5],
|
|
160
|
+
diffuse_color : list | torch.Tensor = [0.5, 0.5, 0.5],
|
|
161
|
+
specular_color : list | torch.Tensor = [0.5, 0.5, 0.5],
|
|
162
|
+
specular_coefficient : float = 1.0,
|
|
163
|
+
|
|
164
|
+
device : Optional[torch.device] = None
|
|
165
|
+
):
|
|
166
|
+
#
|
|
167
|
+
if isinstance(init_translation, list):
|
|
168
|
+
init_translation = torch.tensor(init_translation, device=device)
|
|
169
|
+
if isinstance(init_rotation, list):
|
|
170
|
+
init_rotation = torch.tensor(init_rotation, device=device)
|
|
171
|
+
if isinstance(init_scale, list):
|
|
172
|
+
init_scale = torch.tensor(init_scale, device=device)
|
|
173
|
+
|
|
174
|
+
# Compute Normals
|
|
175
|
+
normals = compute_normals(vertices[None,:], faces[None,:])[0]
|
|
176
|
+
|
|
177
|
+
return MeshObject(
|
|
178
|
+
vertices = vertices,
|
|
179
|
+
faces = faces,
|
|
180
|
+
normals = normals,
|
|
181
|
+
init_transform = init_transform,
|
|
182
|
+
init_translation = init_translation,
|
|
183
|
+
init_rotation = init_rotation,
|
|
184
|
+
init_scale = init_scale,
|
|
185
|
+
|
|
186
|
+
ambient_color = ambient_color,
|
|
187
|
+
diffuse_color = diffuse_color,
|
|
188
|
+
specular_color = specular_color,
|
|
189
|
+
specular_coefficient = specular_coefficient
|
|
190
|
+
)
|
|
191
|
+
|
|
192
|
+
def generate_cube_mesh(
|
|
193
|
+
init_transform : Optional[Transform] = None,
|
|
194
|
+
init_translation : Optional[list | torch.Tensor] = None,
|
|
195
|
+
init_rotation : Optional[list | torch.Tensor] = None,
|
|
196
|
+
init_scale : Optional[list | torch.Tensor] = None,
|
|
197
|
+
|
|
198
|
+
ambient_color : list | torch.Tensor = [0.5, 0.5, 0.5],
|
|
199
|
+
diffuse_color : list | torch.Tensor = [0.5, 0.5, 0.5],
|
|
200
|
+
specular_color : list | torch.Tensor = [0.5, 0.5, 0.5],
|
|
201
|
+
|
|
202
|
+
device : Optional[torch.device] = None
|
|
203
|
+
):
|
|
204
|
+
#
|
|
205
|
+
if isinstance(init_translation, list):
|
|
206
|
+
init_translation = torch.tensor(init_translation, device=device)
|
|
207
|
+
if isinstance(init_rotation, list):
|
|
208
|
+
init_rotation = torch.tensor(init_rotation, device=device)
|
|
209
|
+
if isinstance(init_scale, list):
|
|
210
|
+
init_scale = torch.tensor(init_scale, device=device)
|
|
211
|
+
|
|
212
|
+
# Define vertices for a cube centered at the origin with side length 2
|
|
213
|
+
vertices = torch.tensor([
|
|
214
|
+
[-1, -1, -1],
|
|
215
|
+
[-1, -1, 1],
|
|
216
|
+
[-1, 1, -1],
|
|
217
|
+
[-1, 1, 1],
|
|
218
|
+
[ 1, -1, -1],
|
|
219
|
+
[ 1, -1, 1],
|
|
220
|
+
[ 1, 1, -1],
|
|
221
|
+
[ 1, 1, 1]
|
|
222
|
+
], device=device, dtype=torch.float32)
|
|
223
|
+
|
|
224
|
+
# Define faces as two triangles per cube face (12 triangles total)
|
|
225
|
+
faces = torch.tensor([
|
|
226
|
+
[0, 1, 2], [1, 3, 2], # Left face
|
|
227
|
+
[4, 6, 5], [5, 6, 7], # Right face
|
|
228
|
+
[0, 4, 1], [1, 4, 5], # Bottom face
|
|
229
|
+
[2, 3, 6], [3, 7, 6], # Top face
|
|
230
|
+
[0, 2, 4], [2, 6, 4], # Back face
|
|
231
|
+
[1, 5, 3], [3, 5, 7] # Front face
|
|
232
|
+
], device=device, dtype=torch.int64)
|
|
233
|
+
|
|
234
|
+
return MeshObject(
|
|
235
|
+
vertices = vertices,
|
|
236
|
+
faces = faces,
|
|
237
|
+
normals = vertices / torch.linalg.norm(vertices, dim=-1, keepdim=True),
|
|
238
|
+
init_transform = init_transform,
|
|
239
|
+
init_translation = init_translation,
|
|
240
|
+
init_rotation = init_rotation,
|
|
241
|
+
init_scale = init_scale,
|
|
242
|
+
|
|
243
|
+
ambient_color = ambient_color,
|
|
244
|
+
diffuse_color = diffuse_color,
|
|
245
|
+
specular_color = specular_color
|
|
246
|
+
)
|
|
247
|
+
|
|
248
|
+
|
|
249
|
+
def generate_icosphere_mesh(
|
|
250
|
+
num_subdivisions : int = 1,
|
|
251
|
+
init_transform : Optional[Transform] = None,
|
|
252
|
+
init_translation : Optional[list | torch.Tensor] = None,
|
|
253
|
+
init_rotation : Optional[list | torch.Tensor] = None,
|
|
254
|
+
init_scale : Optional[list | torch.Tensor] = None,
|
|
255
|
+
|
|
256
|
+
ambient_color : list | torch.Tensor = [0.5, 0.5, 0.5],
|
|
257
|
+
diffuse_color : list | torch.Tensor = [0.5, 0.5, 0.5],
|
|
258
|
+
specular_color : list | torch.Tensor = [0.5, 0.5, 0.5],
|
|
259
|
+
|
|
260
|
+
device : Optional[torch.device] = None
|
|
261
|
+
):
|
|
262
|
+
#
|
|
263
|
+
if isinstance(init_translation, list):
|
|
264
|
+
init_translation = torch.tensor(init_translation, device=device)
|
|
265
|
+
if isinstance(init_rotation, list):
|
|
266
|
+
init_rotation = torch.tensor(init_rotation, device=device)
|
|
267
|
+
if isinstance(init_scale, list):
|
|
268
|
+
init_scale = torch.tensor(init_scale, device=device)
|
|
269
|
+
|
|
270
|
+
t = (1.0 + math.sqrt(5.0)) / 2.0
|
|
271
|
+
|
|
272
|
+
verts = [
|
|
273
|
+
[-1, t, 0],
|
|
274
|
+
[ 1, t, 0],
|
|
275
|
+
[-1, -t, 0],
|
|
276
|
+
[ 1, -t, 0],
|
|
277
|
+
[ 0, -1, t],
|
|
278
|
+
[ 0, 1, t],
|
|
279
|
+
[ 0, -1, -t],
|
|
280
|
+
[ 0, 1, -t],
|
|
281
|
+
[ t, 0, -1],
|
|
282
|
+
[ t, 0, 1],
|
|
283
|
+
[-t, 0, -1],
|
|
284
|
+
[-t, 0, 1],
|
|
285
|
+
]
|
|
286
|
+
|
|
287
|
+
# Normalize vertices to unit length
|
|
288
|
+
vertices = []
|
|
289
|
+
for v in verts:
|
|
290
|
+
norm = math.sqrt(v[0]**2 + v[1]**2 + v[2]**2)
|
|
291
|
+
vertices.append([v[0] / norm, v[1] / norm, v[2] / norm])
|
|
292
|
+
vertices = torch.tensor(vertices, device=device, dtype=torch.float32)
|
|
293
|
+
|
|
294
|
+
faces = torch.tensor([
|
|
295
|
+
[0, 11, 5],
|
|
296
|
+
[0, 5, 1],
|
|
297
|
+
[0, 1, 7],
|
|
298
|
+
[0, 7, 10],
|
|
299
|
+
[0, 10, 11],
|
|
300
|
+
[1, 5, 9],
|
|
301
|
+
[5, 11, 4],
|
|
302
|
+
[11, 10, 2],
|
|
303
|
+
[10, 7, 6],
|
|
304
|
+
[7, 1, 8],
|
|
305
|
+
[3, 9, 4],
|
|
306
|
+
[3, 4, 2],
|
|
307
|
+
[3, 2, 6],
|
|
308
|
+
[3, 6, 8],
|
|
309
|
+
[3, 8, 9],
|
|
310
|
+
[4, 9, 5],
|
|
311
|
+
[2, 4, 11],
|
|
312
|
+
[6, 2, 10],
|
|
313
|
+
[8, 6, 7],
|
|
314
|
+
[9, 8, 1],
|
|
315
|
+
], device=device, dtype=torch.int64)
|
|
316
|
+
|
|
317
|
+
# Subdivide faces
|
|
318
|
+
for _ in range(num_subdivisions):
|
|
319
|
+
new_faces = []
|
|
320
|
+
midpoints = {}
|
|
321
|
+
for face in faces:
|
|
322
|
+
a, b, c = face.tolist()
|
|
323
|
+
ab = tuple(sorted((a, b)))
|
|
324
|
+
ac = tuple(sorted((a, c)))
|
|
325
|
+
bc = tuple(sorted((b, c)))
|
|
326
|
+
|
|
327
|
+
if ab not in midpoints:
|
|
328
|
+
midpoints[ab] = len(vertices)
|
|
329
|
+
vertices = torch.cat([vertices, (vertices[a:a+1] + vertices[b:b+1]) / 2.0], dim=0)
|
|
330
|
+
if ac not in midpoints:
|
|
331
|
+
midpoints[ac] = len(vertices)
|
|
332
|
+
vertices = torch.cat([vertices, (vertices[a:a+1] + vertices[c:c+1]) / 2.0], dim=0)
|
|
333
|
+
if bc not in midpoints:
|
|
334
|
+
midpoints[bc] = len(vertices)
|
|
335
|
+
vertices = torch.cat([vertices, (vertices[b:b+1] + vertices[c:c+1]) / 2.0], dim=0)
|
|
336
|
+
|
|
337
|
+
d = midpoints[ab]
|
|
338
|
+
e = midpoints[ac]
|
|
339
|
+
f = midpoints[bc]
|
|
340
|
+
|
|
341
|
+
new_faces.extend([
|
|
342
|
+
[a, d, e],
|
|
343
|
+
[b, d, f],
|
|
344
|
+
[c, e, f],
|
|
345
|
+
[d, e, f]
|
|
346
|
+
])
|
|
347
|
+
faces = torch.tensor(new_faces, device=device, dtype=torch.int64)
|
|
348
|
+
|
|
349
|
+
# Project vertices to unit length
|
|
350
|
+
vertices = torch.nn.functional.normalize(vertices, dim=-1)
|
|
351
|
+
vertices = torch.cat([vertices, torch.ones((vertices.shape[0], 1), device=device)], dim=-1)
|
|
352
|
+
vertices = vertices @ torch.diag(torch.tensor([1, 1, 1, 0], device=device)).to(vertices.dtype)
|
|
353
|
+
vertices = vertices[..., 0:3]
|
|
354
|
+
|
|
355
|
+
return MeshObject(
|
|
356
|
+
vertices = vertices,
|
|
357
|
+
faces = faces,
|
|
358
|
+
normals = vertices,
|
|
359
|
+
init_transform = init_transform,
|
|
360
|
+
init_translation = init_translation,
|
|
361
|
+
init_rotation = init_rotation,
|
|
362
|
+
init_scale = init_scale,
|
|
363
|
+
|
|
364
|
+
ambient_color = ambient_color,
|
|
365
|
+
diffuse_color = diffuse_color,
|
|
366
|
+
specular_color = specular_color
|
|
367
|
+
)
|
|
@@ -0,0 +1,284 @@
|
|
|
1
|
+
import gffx
|
|
2
|
+
import torch
|
|
3
|
+
from tqdm import tqdm
|
|
4
|
+
|
|
5
|
+
from .camera import *
|
|
6
|
+
|
|
7
|
+
def ray_sphere_intersection(ray_origins, ray_directions, sphere_pos, sphere_radius):
|
|
8
|
+
# Check discriminant (ray-sphere intersection)
|
|
9
|
+
e_min_c = ray_origins - sphere_pos[None, None, None, :] # dim(B, H, W, 3)
|
|
10
|
+
d_dot_e_min_c = torch.sum(ray_directions * e_min_c, dim=-1, keepdim=True) # dim(B, H, W, 1)
|
|
11
|
+
d_dot_d = torch.sum(ray_directions * ray_directions, dim=-1, keepdim=True) # dim(B, H, W, 1)
|
|
12
|
+
discriminant = d_dot_e_min_c ** 2 - d_dot_d * (torch.sum(e_min_c * e_min_c, dim=-1, keepdim=True) - sphere_radius ** 2)
|
|
13
|
+
|
|
14
|
+
# Compute t
|
|
15
|
+
hit_mask = (discriminant >= 0).float()
|
|
16
|
+
t = torch.sqrt(discriminant * hit_mask) / d_dot_d
|
|
17
|
+
t_minus = torch.clip(-d_dot_e_min_c - t, 0)
|
|
18
|
+
t_plus = torch.clip(-d_dot_e_min_c + t, 0)
|
|
19
|
+
t = torch.maximum(t_minus, t_plus) * hit_mask
|
|
20
|
+
|
|
21
|
+
return t
|
|
22
|
+
|
|
23
|
+
def ray_triangle_intersection(
|
|
24
|
+
ray_origins,
|
|
25
|
+
ray_directions,
|
|
26
|
+
triangle_vertices,
|
|
27
|
+
t0 = 0,
|
|
28
|
+
t1 = 100
|
|
29
|
+
):
|
|
30
|
+
"""
|
|
31
|
+
Compute ray-triangle intersection using Cramer's Rule
|
|
32
|
+
|
|
33
|
+
Terms
|
|
34
|
+
-----
|
|
35
|
+
B = batch size
|
|
36
|
+
N = number of triangles
|
|
37
|
+
|
|
38
|
+
Parameters
|
|
39
|
+
----------
|
|
40
|
+
ray_origins : torch.Tensor
|
|
41
|
+
Ray origins -> dim(B, 3)
|
|
42
|
+
ray_directions : torch.Tensor
|
|
43
|
+
Ray directions -> dim(B, 3)
|
|
44
|
+
triangle_vertices : torch.Tensor
|
|
45
|
+
Triangle vertices -> dim(N, 3, 3)
|
|
46
|
+
t0 : float
|
|
47
|
+
Minimum t value (default: 0)
|
|
48
|
+
t1 : float
|
|
49
|
+
Maximum t value (default: 100)
|
|
50
|
+
|
|
51
|
+
Returns
|
|
52
|
+
-------
|
|
53
|
+
beta : torch.Tensor
|
|
54
|
+
beta values -> dim(B, N)
|
|
55
|
+
gamma : torch.Tensor
|
|
56
|
+
gamma values -> dim(B, N)
|
|
57
|
+
t : torch.Tensor
|
|
58
|
+
t values -> dim(B, N)
|
|
59
|
+
intersect : torch.Tensor
|
|
60
|
+
Intersection mask -> dim(B, N)
|
|
61
|
+
|
|
62
|
+
Notes
|
|
63
|
+
-----
|
|
64
|
+
| x_a - x_b; x_a - x_c; x_d | | beta | | x_a - x_e |
|
|
65
|
+
| y_a - y_b; y_a - y_c; y_d | | gamma | = | y_a - y_e |
|
|
66
|
+
| z_a - z_b; z_a - z_c; z_d | | t | | z_a - z_e |
|
|
67
|
+
|
|
68
|
+
A x = b
|
|
69
|
+
|
|
70
|
+
Solve via Cramer's Rule:
|
|
71
|
+
beta = det(A_1) / det(A)
|
|
72
|
+
gamma = det(A_2) / det(A)
|
|
73
|
+
t = det(A_3) / det(A)
|
|
74
|
+
Where A_i is A with column i replaced by b
|
|
75
|
+
"""
|
|
76
|
+
B = ray_origins.shape[0]
|
|
77
|
+
N = triangle_vertices.shape[0]
|
|
78
|
+
device = triangle_vertices.device
|
|
79
|
+
|
|
80
|
+
# Setup
|
|
81
|
+
A = torch.zeros((B, N, 3, 3), device=device)
|
|
82
|
+
b = torch.zeros((B, N, 3, 1), device=device)
|
|
83
|
+
triangle_vertices = triangle_vertices[None, :, :, :]
|
|
84
|
+
ray_directions = ray_directions[:, None, :]
|
|
85
|
+
ray_origins = ray_origins[:, None, :]
|
|
86
|
+
|
|
87
|
+
#
|
|
88
|
+
A[...,0,0] = triangle_vertices[...,0,0] - triangle_vertices[...,1,0]
|
|
89
|
+
A[...,0,1] = triangle_vertices[...,0,0] - triangle_vertices[...,2,0]
|
|
90
|
+
A[...,0,2] = ray_directions[...,0]
|
|
91
|
+
A[...,1,0] = triangle_vertices[...,0,1] - triangle_vertices[...,1,1]
|
|
92
|
+
A[...,1,1] = triangle_vertices[...,0,1] - triangle_vertices[...,2,1]
|
|
93
|
+
A[...,1,2] = ray_directions[...,1]
|
|
94
|
+
A[...,2,0] = triangle_vertices[...,0,2] - triangle_vertices[...,1,2]
|
|
95
|
+
A[...,2,1] = triangle_vertices[...,0,2] - triangle_vertices[...,2,2]
|
|
96
|
+
A[...,2,2] = ray_directions[...,2]
|
|
97
|
+
|
|
98
|
+
b[...,0,0] = triangle_vertices[...,0,0] - ray_origins[...,0]
|
|
99
|
+
b[...,1,0] = triangle_vertices[...,0,1] - ray_origins[...,1]
|
|
100
|
+
b[...,2,0] = triangle_vertices[...,0,2] - ray_origins[...,2]
|
|
101
|
+
|
|
102
|
+
# Cramer's Rule
|
|
103
|
+
x = gffx.linalg.cramer(A, b)
|
|
104
|
+
beta = x[...,0,0]
|
|
105
|
+
gamma = x[...,1,0]
|
|
106
|
+
t = x[...,2,0]
|
|
107
|
+
|
|
108
|
+
# Intersection Test
|
|
109
|
+
intersect = ((t > t0) & (t < t1)) & ((gamma >= 0) & (gamma <= 1)) & ((beta > 0) & (beta + gamma < 1))
|
|
110
|
+
|
|
111
|
+
return beta, gamma, t, intersect
|
|
112
|
+
|
|
113
|
+
# [4NOW] 1 camera
|
|
114
|
+
def mesh_render(
|
|
115
|
+
meshes : list | gffx.obj.MeshObject,
|
|
116
|
+
camera : Camera,
|
|
117
|
+
light_intensity : float = 1,
|
|
118
|
+
ambient_intensity : float = 0.2,
|
|
119
|
+
light_pos : list | torch.Tensor = [5, 5, 5],
|
|
120
|
+
background_color : list | torch.Tensor = [0, 0, 0],
|
|
121
|
+
ray_chunk_size : int = 4096,
|
|
122
|
+
|
|
123
|
+
verbose : bool = False,
|
|
124
|
+
|
|
125
|
+
device : Optional[torch.device] = None
|
|
126
|
+
):
|
|
127
|
+
# [4NOW] 1 camera
|
|
128
|
+
B = 1
|
|
129
|
+
|
|
130
|
+
if isinstance(meshes, gffx.obj.MeshObject):
|
|
131
|
+
meshes = [meshes]
|
|
132
|
+
|
|
133
|
+
# [4NOW] Flat colors
|
|
134
|
+
if isinstance(background_color, list):
|
|
135
|
+
background_color = torch.tensor(background_color, device=device, dtype=torch.float32)
|
|
136
|
+
|
|
137
|
+
diffuse_colors = [obj.diffuse_color for obj in meshes] + [background_color]
|
|
138
|
+
diffuse_colors = torch.stack(diffuse_colors, dim=0).to(device)
|
|
139
|
+
|
|
140
|
+
specular_coefficients = [obj.specular_coefficient for obj in meshes] + [1]
|
|
141
|
+
specular_coefficients = torch.tensor(specular_coefficients, device=device, dtype=torch.float32)
|
|
142
|
+
|
|
143
|
+
specular_colors = [obj.specular_color for obj in meshes] + [background_color]
|
|
144
|
+
specular_colors = torch.stack(specular_colors, dim=0).to(device)
|
|
145
|
+
|
|
146
|
+
ambient_colors = [obj.ambient_color for obj in meshes] + [background_color]
|
|
147
|
+
ambient_colors = torch.stack(ambient_colors, dim=0).to(device)
|
|
148
|
+
|
|
149
|
+
# Light setup
|
|
150
|
+
if isinstance(light_pos, list):
|
|
151
|
+
light_pos = torch.tensor(light_pos, device=device, dtype=torch.float32)
|
|
152
|
+
|
|
153
|
+
#
|
|
154
|
+
object_hit = torch.full((B * camera.width * camera.height,), -1, device=device, dtype=torch.int64)
|
|
155
|
+
face_hit = torch.full((B * camera.width * camera.height,), -1, device=device, dtype=torch.int64)
|
|
156
|
+
t_val = torch.full((B * camera.width * camera.height,), float('inf'), device=device)
|
|
157
|
+
normals = torch.zeros((B * camera.width * camera.height, 3), device=device)
|
|
158
|
+
hit_pos = torch.zeros((B * camera.width * camera.height, 3), device=device)
|
|
159
|
+
|
|
160
|
+
#
|
|
161
|
+
if verbose:
|
|
162
|
+
print('Ray Chunk Size:', ray_chunk_size)
|
|
163
|
+
ray_origins = camera.ray_origins.view(B * camera.width * camera.height, 3) # dim(B * W * H, 3)
|
|
164
|
+
ray_directions = camera.ray_directions.view(B * camera.width * camera.height, 3) # dim(B * W * H, 3)
|
|
165
|
+
|
|
166
|
+
#
|
|
167
|
+
for obj_idx, obj in enumerate(meshes):
|
|
168
|
+
transformed_vertices, transformed_normals = obj.get_transformed()
|
|
169
|
+
|
|
170
|
+
face_tri_vertices = transformed_vertices[obj.faces] # dim(F, 3, 3)
|
|
171
|
+
face_tri_normals = transformed_normals[obj.faces] # dim(F, 3, 3)
|
|
172
|
+
|
|
173
|
+
# Chunking
|
|
174
|
+
pbar = tqdm(
|
|
175
|
+
iterable = range(0, B * camera.width * camera.height, ray_chunk_size),
|
|
176
|
+
desc = f'Object {obj_idx}',
|
|
177
|
+
)
|
|
178
|
+
for i in pbar:
|
|
179
|
+
C = ray_chunk_size if i + ray_chunk_size < B * camera.width * camera.height else (B * camera.width * camera.height - i)
|
|
180
|
+
|
|
181
|
+
#
|
|
182
|
+
beta, gamma, t, intersect = gffx.ray.ray_triangle_intersection(
|
|
183
|
+
ray_origins = ray_origins[i:i+C], # dim(C, 3)
|
|
184
|
+
ray_directions = ray_directions[i:i+C], # dim(C, 3)
|
|
185
|
+
triangle_vertices = face_tri_vertices, # dim(F, 3, 3)
|
|
186
|
+
t0 = 0,
|
|
187
|
+
t1 = 100
|
|
188
|
+
)
|
|
189
|
+
|
|
190
|
+
# Choose the smallest t which intersects
|
|
191
|
+
t_valid = torch.where(
|
|
192
|
+
intersect,
|
|
193
|
+
torch.where(t < 0, float('inf'), t),
|
|
194
|
+
float('inf')
|
|
195
|
+
)
|
|
196
|
+
face_idx_valid_min = t_valid.argmin(dim=-1, keepdim=True)
|
|
197
|
+
beta_valid_min = torch.gather(
|
|
198
|
+
input = beta,
|
|
199
|
+
dim = -1,
|
|
200
|
+
index = face_idx_valid_min
|
|
201
|
+
)[:,0]
|
|
202
|
+
gamma_valid_min = torch.gather(
|
|
203
|
+
input = gamma,
|
|
204
|
+
dim = -1,
|
|
205
|
+
index = face_idx_valid_min
|
|
206
|
+
)[:,0]
|
|
207
|
+
|
|
208
|
+
t_valid_min = torch.gather(
|
|
209
|
+
input = t_valid,
|
|
210
|
+
dim = -1,
|
|
211
|
+
index = face_idx_valid_min
|
|
212
|
+
)[:,0]
|
|
213
|
+
intersect_valid_min = torch.gather(
|
|
214
|
+
input = intersect,
|
|
215
|
+
dim = -1,
|
|
216
|
+
index = face_idx_valid_min
|
|
217
|
+
)[:,0]
|
|
218
|
+
|
|
219
|
+
#
|
|
220
|
+
object_hit[i:i+C] = torch.where(
|
|
221
|
+
((t_valid_min < t_val[i:i+C]) & intersect_valid_min),
|
|
222
|
+
obj_idx,
|
|
223
|
+
object_hit[i:i+C]
|
|
224
|
+
)
|
|
225
|
+
|
|
226
|
+
#
|
|
227
|
+
face_hit[i:i+C] = torch.where(
|
|
228
|
+
((t_valid_min < t_val[i:i+C]) & intersect_valid_min),
|
|
229
|
+
face_idx_valid_min[:,0],
|
|
230
|
+
face_hit[i:i+C]
|
|
231
|
+
)
|
|
232
|
+
|
|
233
|
+
# Interpolate normals
|
|
234
|
+
interpolated_normals = (1 - beta - gamma)[...,None] * face_tri_normals[None, :, 0] + beta[...,None] * face_tri_normals[None, :, 1] + gamma[...,None] * face_tri_normals[None, :, 2]
|
|
235
|
+
|
|
236
|
+
interpolated_normals = torch.gather(
|
|
237
|
+
input = interpolated_normals, # dim(C, F, 3)
|
|
238
|
+
dim = 1,
|
|
239
|
+
index = face_idx_valid_min[..., None].expand(-1, -1, 3)
|
|
240
|
+
)[:,0,:]
|
|
241
|
+
normals[i:i+C] = torch.where(
|
|
242
|
+
((t_valid_min < t_val[i:i+C]) & intersect_valid_min)[:,None],
|
|
243
|
+
interpolated_normals,
|
|
244
|
+
normals[i:i+C]
|
|
245
|
+
)
|
|
246
|
+
|
|
247
|
+
# Interpolate hit positions
|
|
248
|
+
hit_pos[i:i+C] = torch.where(
|
|
249
|
+
((t_valid_min < t_val[i:i+C]) & intersect_valid_min)[:,None],
|
|
250
|
+
ray_origins[i:i+C] + t_valid_min[:,None] * ray_directions[i:i+C],
|
|
251
|
+
hit_pos[i:i+C]
|
|
252
|
+
)
|
|
253
|
+
|
|
254
|
+
#
|
|
255
|
+
t_val[i:i+C] = torch.where(
|
|
256
|
+
((t_valid_min < t_val[i:i+C]) & intersect_valid_min),
|
|
257
|
+
t_valid_min,
|
|
258
|
+
t_val[i:i+C]
|
|
259
|
+
)
|
|
260
|
+
|
|
261
|
+
# Reshape
|
|
262
|
+
object_hit = object_hit.view(B, camera.width, camera.height)
|
|
263
|
+
face_hit = face_hit.view(B, camera.width, camera.height)
|
|
264
|
+
normals = normals.view(B, camera.width, camera.height, 3)
|
|
265
|
+
hit_pos = hit_pos.view(B, camera.width, camera.height, 3)
|
|
266
|
+
|
|
267
|
+
# Compute Phong Shading
|
|
268
|
+
light_pos = light_pos[None, None, None, :]
|
|
269
|
+
light_dir = light_pos - hit_pos
|
|
270
|
+
light_dir /= torch.linalg.norm(light_dir, dim=-1, keepdim=True)
|
|
271
|
+
|
|
272
|
+
view_dir = camera.pos[None, None, None, :] - hit_pos
|
|
273
|
+
view_dir /= torch.linalg.norm(view_dir, dim=-1, keepdim=True)
|
|
274
|
+
bisector_vec = light_dir + view_dir
|
|
275
|
+
bisector_vec /= torch.linalg.norm(bisector_vec, dim=-1, keepdim=True)
|
|
276
|
+
|
|
277
|
+
diffuse_weight = torch.clamp(torch.sum(light_dir * normals, dim=-1, keepdim=True), min=0)
|
|
278
|
+
specular_weight = torch.clamp(torch.sum(bisector_vec * normals, dim=-1, keepdim=True), min=0) ** specular_coefficients[object_hit][...,None]
|
|
279
|
+
|
|
280
|
+
L = diffuse_colors[object_hit] * light_intensity * diffuse_weight
|
|
281
|
+
L += specular_colors[object_hit] * light_intensity * specular_weight
|
|
282
|
+
L += ambient_colors[object_hit] * ambient_intensity
|
|
283
|
+
|
|
284
|
+
return L
|
|
@@ -0,0 +1,66 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
from typing import Optional
|
|
3
|
+
|
|
4
|
+
class Camera:
|
|
5
|
+
def __init__(
|
|
6
|
+
self,
|
|
7
|
+
width : int = 512,
|
|
8
|
+
height : int = 512,
|
|
9
|
+
pos : list | torch.Tensor = [0, 0, 5],
|
|
10
|
+
dir : list | torch.Tensor = [0, 0, -1],
|
|
11
|
+
world_up : list | torch.Tensor = [0, 1, 0],
|
|
12
|
+
ray_origin_scale_du : float = 0,
|
|
13
|
+
ray_origin_scale_dv : float = 0,
|
|
14
|
+
screen_distance : float = 1.0,
|
|
15
|
+
device : Optional[torch.device] = None
|
|
16
|
+
):
|
|
17
|
+
self.width = width
|
|
18
|
+
self.height = height
|
|
19
|
+
self.ray_origin_scale_du = ray_origin_scale_du
|
|
20
|
+
self.ray_origin_scale_dv = ray_origin_scale_dv
|
|
21
|
+
self.screen_distance = screen_distance
|
|
22
|
+
|
|
23
|
+
if not device:
|
|
24
|
+
self.device = device
|
|
25
|
+
|
|
26
|
+
self.pos = pos
|
|
27
|
+
if isinstance(self.pos, list):
|
|
28
|
+
self.pos = torch.tensor(self.pos, device=device, dtype=torch.float32)
|
|
29
|
+
|
|
30
|
+
self.dir = dir
|
|
31
|
+
if isinstance(self.dir, list):
|
|
32
|
+
self.dir = torch.tensor(self.dir, device=device, dtype=torch.float32)
|
|
33
|
+
|
|
34
|
+
self.world_up = world_up
|
|
35
|
+
if isinstance(self.world_up, list):
|
|
36
|
+
self.world_up = torch.tensor(self.world_up, device=device, dtype=torch.float32)
|
|
37
|
+
|
|
38
|
+
# Camera basis
|
|
39
|
+
self.u = torch.cross(self.dir, self.world_up, dim=-1)
|
|
40
|
+
self.u /= torch.linalg.norm(self.u)
|
|
41
|
+
self.v = torch.cross(self.u, self.dir, dim=-1)
|
|
42
|
+
self.w = -self.dir
|
|
43
|
+
|
|
44
|
+
# Screen
|
|
45
|
+
self.aspect_ratio = self.width / self.height
|
|
46
|
+
self.left = -self.aspect_ratio
|
|
47
|
+
self.right = self.aspect_ratio
|
|
48
|
+
self.bottom = -1
|
|
49
|
+
self.top = 1
|
|
50
|
+
|
|
51
|
+
self.screen_grids = torch.stack(torch.meshgrid(
|
|
52
|
+
(self.right - self.left) * ((torch.arange(0, self.width, device=device) + 0.5) / self.width) + self.left,
|
|
53
|
+
(self.top - self.bottom) * ((torch.arange(0, self.height, device=device) + 0.5) / self.height) + self.bottom,
|
|
54
|
+
indexing='ij'
|
|
55
|
+
), dim=-1).float().expand(1, -1, -1, -1)
|
|
56
|
+
|
|
57
|
+
self.ray_origins = self.ray_origin_scale_du * self.screen_grids[..., 0:1] * self.u[None, None, None, :]
|
|
58
|
+
self.ray_origins += self.ray_origin_scale_dv * self.screen_grids[..., 1:2] * self.v[None, None, None, :]
|
|
59
|
+
self.ray_origins += self.pos[None,None,None,:]
|
|
60
|
+
|
|
61
|
+
# Compute viewing rays
|
|
62
|
+
self.ray_directions = self.screen_grids[..., 0:1] * self.u[None, None, None, :]
|
|
63
|
+
self.ray_directions += self.screen_grids[..., 1:2] * self.v[None, None, None, :]
|
|
64
|
+
self.ray_directions += self.pos[None,None,None,:] - (self.w[None,None,None,:] * self.screen_distance)
|
|
65
|
+
self.ray_directions -= self.ray_origins
|
|
66
|
+
self.ray_directions /= torch.linalg.norm(self.ray_directions, dim=-1, keepdim=True)
|