autoforge 1.5.1__tar.gz → 1.5.2__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.
- {autoforge-1.5.1 → autoforge-1.5.2}/PKG-INFO +1 -1
- {autoforge-1.5.1 → autoforge-1.5.2}/pyproject.toml +1 -1
- {autoforge-1.5.1 → autoforge-1.5.2}/src/AutoForge.egg-info/PKG-INFO +1 -1
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Helper/OptimizerHelper.py +269 -117
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/auto_forge.py +3 -3
- {autoforge-1.5.1 → autoforge-1.5.2}/README.md +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/setup.cfg +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/AutoForge.egg-info/SOURCES.txt +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/AutoForge.egg-info/dependency_links.txt +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/AutoForge.egg-info/entry_points.txt +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/AutoForge.egg-info/requires.txt +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/AutoForge.egg-info/top_level.txt +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Gui/Views/__init__.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Gui/Views/backend_database.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Gui/Views/filament_editor.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Gui/__init__.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Gui/app.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Gui/main.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Helper/CAdamW.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Helper/FilamentHelper.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Helper/ImageHelper.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Helper/OutputHelper.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Helper/PruningHelper.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Helper/__init__.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Helper/filamentcolors_library.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Loss/LossFunctions.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Loss/PerceptionLoss.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Loss/__init__.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Misc/__init__.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Misc/color_fit_all.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Modules/Optimizer.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Modules/__init__.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/__init__.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/__main__.py +0 -0
- {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/old_jax_files.py +0 -0
|
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "AutoForge"
|
|
7
|
-
version="1.5.
|
|
7
|
+
version="1.5.2"
|
|
8
8
|
description = "AutoForge is a Python tool for generating 3D printed layered models from an input image."
|
|
9
9
|
authors = [{ name = "Hendric Voss", email = "hvoss@techfak.uni-bielefeld.de" }]
|
|
10
10
|
readme = "README.md"
|
|
@@ -20,6 +20,16 @@ def adaptive_round(
|
|
|
20
20
|
) -> torch.Tensor:
|
|
21
21
|
"""
|
|
22
22
|
Smooth rounding based on temperature 'tau'.
|
|
23
|
+
|
|
24
|
+
Args:
|
|
25
|
+
x (torch.Tensor): The input tensor to be rounded.
|
|
26
|
+
tau (float): The current temperature parameter.
|
|
27
|
+
high_tau (float): The high threshold for the temperature.
|
|
28
|
+
low_tau (float): The low threshold for the temperature.
|
|
29
|
+
temp (float): The temperature parameter for the sigmoid function.
|
|
30
|
+
|
|
31
|
+
Returns:
|
|
32
|
+
torch.Tensor: The rounded tensor.
|
|
23
33
|
"""
|
|
24
34
|
if tau <= low_tau:
|
|
25
35
|
return torch.round(x)
|
|
@@ -40,6 +50,16 @@ def adaptive_round(
|
|
|
40
50
|
# A deterministic random generator that mimics torch.rand_like.
|
|
41
51
|
@torch.jit.script
|
|
42
52
|
def deterministic_rand_like(tensor: torch.Tensor, seed: int) -> torch.Tensor:
|
|
53
|
+
"""
|
|
54
|
+
Generate a deterministic random tensor that mimics torch.rand_like.
|
|
55
|
+
|
|
56
|
+
Args:
|
|
57
|
+
tensor (torch.Tensor): The input tensor whose shape and device will be used.
|
|
58
|
+
seed (int): The seed for the deterministic random generator.
|
|
59
|
+
|
|
60
|
+
Returns:
|
|
61
|
+
torch.Tensor: A tensor with the same shape as the input tensor, filled with deterministic random values.
|
|
62
|
+
"""
|
|
43
63
|
# Compute the total number of elements.
|
|
44
64
|
n: int = 1
|
|
45
65
|
for d in tensor.shape:
|
|
@@ -59,6 +79,18 @@ def deterministic_rand_like(tensor: torch.Tensor, seed: int) -> torch.Tensor:
|
|
|
59
79
|
def deterministic_gumbel_softmax(
|
|
60
80
|
logits: torch.Tensor, tau: float, hard: bool, rng_seed: int
|
|
61
81
|
) -> torch.Tensor:
|
|
82
|
+
"""
|
|
83
|
+
Apply the Gumbel-Softmax trick in a deterministic manner using a fixed random seed.
|
|
84
|
+
|
|
85
|
+
Args:
|
|
86
|
+
logits (torch.Tensor): The input logits tensor.
|
|
87
|
+
tau (float): The temperature parameter for the Gumbel-Softmax.
|
|
88
|
+
hard (bool): If True, the output will be one-hot encoded.
|
|
89
|
+
rng_seed (int): The seed for the deterministic random generator.
|
|
90
|
+
|
|
91
|
+
Returns:
|
|
92
|
+
torch.Tensor: The resulting tensor after applying the Gumbel-Softmax trick.
|
|
93
|
+
"""
|
|
62
94
|
eps: float = 1e-20
|
|
63
95
|
# Instead of torch.rand_like(..., generator=...), use our deterministic_rand_like.
|
|
64
96
|
U = deterministic_rand_like(logits, rng_seed)
|
|
@@ -83,57 +115,94 @@ def composite_image_cont(
|
|
|
83
115
|
tau_global: float,
|
|
84
116
|
h: float,
|
|
85
117
|
max_layers: int,
|
|
86
|
-
material_colors: torch.Tensor,
|
|
87
|
-
material_TDs: torch.Tensor,
|
|
88
|
-
background: torch.Tensor,
|
|
118
|
+
material_colors: torch.Tensor, # [n_materials, 3]
|
|
119
|
+
material_TDs: torch.Tensor, # [n_materials]
|
|
120
|
+
background: torch.Tensor, # [3]
|
|
89
121
|
) -> torch.Tensor:
|
|
90
122
|
"""
|
|
91
123
|
Continuous compositing over all pixels.
|
|
92
124
|
Uses Gumbel softmax with either hard or soft sampling depending on tau_global.
|
|
125
|
+
|
|
126
|
+
Args:
|
|
127
|
+
pixel_height_logits (torch.Tensor): Logits for pixel heights, shape [H, W].
|
|
128
|
+
global_logits (torch.Tensor): Logits for global material assignment, shape [max_layers, n_materials].
|
|
129
|
+
tau_height (float): Temperature parameter for height rounding.
|
|
130
|
+
tau_global (float): Temperature parameter for global material assignment.
|
|
131
|
+
h (float): Height of each layer.
|
|
132
|
+
max_layers (int): Maximum number of layers.
|
|
133
|
+
material_colors (torch.Tensor): Tensor of material colors, shape [n_materials, 3].
|
|
134
|
+
material_TDs (torch.Tensor): Tensor of material transmission/opacity parameters, shape [n_materials].
|
|
135
|
+
background (torch.Tensor): Background color tensor, shape [3].
|
|
136
|
+
|
|
137
|
+
Returns:
|
|
138
|
+
torch.Tensor: Composite image tensor, shape [H, W, 3].
|
|
93
139
|
"""
|
|
140
|
+
# Compute continuous and discrete layers for each pixel
|
|
94
141
|
pixel_height = (max_layers * h) * torch.sigmoid(pixel_height_logits)
|
|
95
142
|
continuous_layers = pixel_height / h
|
|
96
143
|
|
|
144
|
+
# adaptive rounding
|
|
97
145
|
adaptive_layers = adaptive_round(
|
|
98
146
|
continuous_layers, tau_height, high_tau=0.1, low_tau=0.01, temp=0.1
|
|
99
147
|
)
|
|
148
|
+
|
|
149
|
+
# "Straight-through" trick to keep gradient from adaptive rounding
|
|
100
150
|
discrete_layers_temp = torch.round(continuous_layers)
|
|
101
151
|
discrete_layers = (
|
|
102
152
|
discrete_layers_temp + (adaptive_layers - discrete_layers_temp).detach()
|
|
103
|
-
).to(torch.int32)
|
|
104
|
-
|
|
153
|
+
).to(torch.int32) # shape [H, W]
|
|
154
|
+
|
|
155
|
+
# Precompute global material assignments for each layer
|
|
156
|
+
# Use a single Gumbel Softmax call for all layers.
|
|
157
|
+
hard_flag = tau_global < 1e-3
|
|
158
|
+
# global_logits: [max_layers, n_materials]
|
|
159
|
+
p = F.gumbel_softmax(
|
|
160
|
+
global_logits, tau_global, hard=hard_flag, dim=1
|
|
161
|
+
) # [L, n_materials]
|
|
162
|
+
|
|
163
|
+
# Compute color_i and TD_i for each layer in one go.
|
|
164
|
+
# color_i -> [max_layers, 3], TD_i -> [max_layers]
|
|
165
|
+
layer_colors = p @ material_colors # [L, 3]
|
|
166
|
+
layer_TDs = p @ material_TDs # [L]
|
|
167
|
+
layer_TDs.clamp_(1e-8, 1e8)
|
|
168
|
+
|
|
169
|
+
# Prepare output and compositing variables
|
|
105
170
|
H, W = pixel_height.shape
|
|
106
171
|
comp = torch.zeros(H, W, 3, dtype=torch.float32, device=pixel_height.device)
|
|
107
172
|
remaining = torch.ones(H, W, dtype=torch.float32, device=pixel_height.device)
|
|
108
173
|
|
|
109
|
-
#
|
|
174
|
+
# Opacity function parameters
|
|
110
175
|
o = -1.2416557e-02
|
|
111
176
|
A = 9.6407950e-01
|
|
112
177
|
k = 3.4103447e01
|
|
113
178
|
b = -4.1554203e00
|
|
114
179
|
|
|
180
|
+
# Composite from top layer (max_layers-1) down to 0
|
|
115
181
|
for i in range(max_layers):
|
|
116
182
|
layer_idx = max_layers - 1 - i
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
# continuous mode: use hard gumbel if tau_global is very low
|
|
121
|
-
if tau_global < 1e-3:
|
|
122
|
-
p_i = F.gumbel_softmax(global_logits[layer_idx], tau_global, hard=True)
|
|
123
|
-
else:
|
|
124
|
-
p_i = F.gumbel_softmax(global_logits[layer_idx], tau_global, hard=False)
|
|
183
|
+
# Mask of which pixels actually print this layer:
|
|
184
|
+
p_print = (discrete_layers > layer_idx).to(dtype=comp.dtype)
|
|
125
185
|
|
|
126
|
-
|
|
127
|
-
|
|
128
|
-
TD_i = torch.clamp(TD_i, 1e-8, 1e8)
|
|
186
|
+
# Effective thickness for those pixels
|
|
187
|
+
eff_thick = p_print * h
|
|
129
188
|
|
|
189
|
+
# Opacity based on thickness and material TD
|
|
190
|
+
TD_i = layer_TDs[layer_idx]
|
|
130
191
|
opac = o + (A * torch.log1p(k * (eff_thick / TD_i)) + b * (eff_thick / TD_i))
|
|
131
192
|
opac = torch.clamp(opac, 0.0, 1.0)
|
|
132
193
|
|
|
133
|
-
|
|
134
|
-
|
|
194
|
+
# Add contribution to comp
|
|
195
|
+
color_i = layer_colors[layer_idx] # shape [3]
|
|
196
|
+
# We need to broadcast opac: [H, W] -> [H, W, 1]
|
|
197
|
+
comp = comp + (remaining * opac).unsqueeze(-1) * color_i
|
|
135
198
|
|
|
199
|
+
# Update the remaining factor
|
|
200
|
+
remaining = remaining * (1.0 - opac)
|
|
201
|
+
|
|
202
|
+
# Add background contribution
|
|
136
203
|
comp = comp + remaining.unsqueeze(-1) * background
|
|
204
|
+
|
|
205
|
+
# 6. Scale to [0, 255] range
|
|
137
206
|
return comp * 255.0
|
|
138
207
|
|
|
139
208
|
|
|
@@ -151,25 +220,34 @@ def composite_image_disc(
|
|
|
151
220
|
rng_seed: int = -1,
|
|
152
221
|
) -> torch.Tensor:
|
|
153
222
|
"""
|
|
154
|
-
|
|
155
|
-
|
|
156
|
-
|
|
157
|
-
|
|
158
|
-
|
|
159
|
-
|
|
160
|
-
|
|
161
|
-
|
|
162
|
-
|
|
163
|
-
|
|
164
|
-
|
|
165
|
-
|
|
223
|
+
Perform discrete compositing over all pixels.
|
|
224
|
+
|
|
225
|
+
Args:
|
|
226
|
+
pixel_height_logits (torch.Tensor): Logits for pixel heights, shape [H, W].
|
|
227
|
+
global_logits (torch.Tensor): Logits for global material assignment, shape [max_layers, n_materials].
|
|
228
|
+
tau_height (float): Temperature parameter for height rounding.
|
|
229
|
+
tau_global (float): Temperature parameter for global material assignment.
|
|
230
|
+
h (float): Height of each layer.
|
|
231
|
+
max_layers (int): Maximum number of layers.
|
|
232
|
+
material_colors (torch.Tensor): Tensor of material colors, shape [n_materials, 3].
|
|
233
|
+
material_TDs (torch.Tensor): Tensor of material transmission/opacity parameters, shape [n_materials].
|
|
234
|
+
background (torch.Tensor): Background color tensor, shape [3].
|
|
235
|
+
rng_seed (int, optional): Random seed for deterministic sampling. Defaults to -1.
|
|
236
|
+
|
|
237
|
+
Returns:
|
|
238
|
+
torch.Tensor: Composite image tensor, shape [H, W, 3].
|
|
166
239
|
"""
|
|
167
240
|
|
|
168
|
-
#
|
|
241
|
+
# -------------------------------------------------------------------------
|
|
242
|
+
# 1) Compute discrete per-pixel layer counts (discrete_layers).
|
|
243
|
+
# -------------------------------------------------------------------------
|
|
244
|
+
# pixel_height ~ [0, max_layers*h]
|
|
169
245
|
pixel_height = (max_layers * h) * torch.sigmoid(pixel_height_logits)
|
|
246
|
+
|
|
247
|
+
# continuous_layers ~ [0, max_layers]
|
|
170
248
|
continuous_layers = pixel_height / h
|
|
171
249
|
|
|
172
|
-
#
|
|
250
|
+
# Use your "adaptive rounding" trick if desired:
|
|
173
251
|
adaptive_layers = adaptive_round(
|
|
174
252
|
continuous_layers, tau_height, high_tau=0.1, low_tau=0.01, temp=0.1
|
|
175
253
|
)
|
|
@@ -178,8 +256,10 @@ def composite_image_disc(
|
|
|
178
256
|
discrete_layers_temp + (adaptive_layers - discrete_layers_temp).detach()
|
|
179
257
|
).to(torch.int32) # [H,W]
|
|
180
258
|
|
|
181
|
-
#
|
|
182
|
-
#
|
|
259
|
+
# -------------------------------------------------------------------------
|
|
260
|
+
# 2) Pick a single global material per layer, either deterministically
|
|
261
|
+
# or via gumbel softmax, as in original code.
|
|
262
|
+
# -------------------------------------------------------------------------
|
|
183
263
|
if rng_seed >= 0:
|
|
184
264
|
# Deterministic sampling for each layer
|
|
185
265
|
new_mats_list = []
|
|
@@ -198,16 +278,16 @@ def composite_image_disc(
|
|
|
198
278
|
p_all = F.gumbel_softmax(global_logits, tau_global, hard=True, dim=1) # [L, M]
|
|
199
279
|
new_mats = torch.argmax(p_all, dim=1).to(torch.int32) # [max_layers]
|
|
200
280
|
|
|
201
|
-
|
|
202
281
|
H, W = pixel_height.shape
|
|
203
282
|
device = pixel_height.device
|
|
204
283
|
|
|
205
284
|
comp = torch.zeros(H, W, 3, dtype=torch.float32, device=device)
|
|
206
285
|
remaining = torch.ones(H, W, dtype=torch.float32, device=device)
|
|
207
286
|
|
|
208
|
-
#
|
|
287
|
+
# Current material index for each pixel, or -1 for none
|
|
209
288
|
cur_mat = -torch.ones((H, W), dtype=torch.int32, device=device)
|
|
210
|
-
|
|
289
|
+
|
|
290
|
+
# Accumulated thickness for the current segment
|
|
211
291
|
acc_thick = torch.zeros((H, W), dtype=torch.float32, device=device)
|
|
212
292
|
|
|
213
293
|
# Opacity function parameters
|
|
@@ -218,102 +298,158 @@ def composite_image_disc(
|
|
|
218
298
|
|
|
219
299
|
# Main compositing loop: top to bottom
|
|
220
300
|
for layer_idx in range(max_layers - 1, -1, -1):
|
|
221
|
-
#
|
|
222
|
-
layer_mat = new_mats[layer_idx] #
|
|
223
|
-
|
|
224
|
-
|
|
225
|
-
|
|
226
|
-
#
|
|
227
|
-
|
|
228
|
-
#
|
|
229
|
-
#
|
|
230
|
-
#
|
|
231
|
-
#
|
|
232
|
-
#
|
|
233
|
-
|
|
234
|
-
|
|
235
|
-
|
|
236
|
-
|
|
237
|
-
|
|
238
|
-
|
|
239
|
-
|
|
240
|
-
|
|
241
|
-
|
|
242
|
-
|
|
243
|
-
|
|
244
|
-
|
|
245
|
-
|
|
246
|
-
|
|
247
|
-
|
|
248
|
-
|
|
249
|
-
|
|
250
|
-
|
|
251
|
-
|
|
252
|
-
|
|
253
|
-
#
|
|
254
|
-
#
|
|
255
|
-
#
|
|
256
|
-
#
|
|
257
|
-
|
|
301
|
+
# layer_mat is the global material chosen for this layer (int32 scalar).
|
|
302
|
+
layer_mat = new_mats[layer_idx] # shape []
|
|
303
|
+
|
|
304
|
+
# Which pixels actually print on this layer?
|
|
305
|
+
# p_print = (discrete_layers > layer_idx)
|
|
306
|
+
p_print = discrete_layers.gt(layer_idx) # bool
|
|
307
|
+
|
|
308
|
+
# ---------------------------------------------------------------------
|
|
309
|
+
# (A) "Finish" any existing segments that are now 'done'.
|
|
310
|
+
#
|
|
311
|
+
# A segment is done if:
|
|
312
|
+
# 1) cur_mat != -1, i.e. the pixel had an ongoing segment
|
|
313
|
+
# 2) EITHER
|
|
314
|
+
# - the pixel does not print now (~p_print),
|
|
315
|
+
# - OR the new layer material differs (cur_mat != layer_mat).
|
|
316
|
+
# ---------------------------------------------------------------------
|
|
317
|
+
mask_done = (cur_mat.ne(-1)) & ((~p_print) | (cur_mat.ne(layer_mat)))
|
|
318
|
+
|
|
319
|
+
# Convert to float for multiplications
|
|
320
|
+
mask_done_f = mask_done.to(torch.float32)
|
|
321
|
+
|
|
322
|
+
# Gather thickness densities & colors for the old segment
|
|
323
|
+
# We'll clamp cur_mat so -1 becomes 0 (doesn't matter since we multiply by 0).
|
|
324
|
+
old_inds_clamped = torch.clamp(cur_mat, min=0)
|
|
325
|
+
td_vals = material_TDs[old_inds_clamped] # [H, W]
|
|
326
|
+
col_vals = material_colors[old_inds_clamped] # [H, W, 3]
|
|
327
|
+
|
|
328
|
+
# Compute alpha from accumulated thickness
|
|
329
|
+
thick_ratio = acc_thick / td_vals
|
|
330
|
+
opac_vals = o + (A * torch.log1p(k * thick_ratio) + b * thick_ratio)
|
|
331
|
+
opac_vals = torch.clamp(opac_vals, 0.0, 1.0) # [H, W]
|
|
332
|
+
|
|
333
|
+
# Compositing the old segment:
|
|
334
|
+
# comp += mask_done_f * remaining * opac_vals * col_vals
|
|
335
|
+
# remaining *= (1 - mask_done_f * opac_vals) ...
|
|
336
|
+
# but we have to broadcast for color:
|
|
337
|
+
comp_add = (mask_done_f * remaining * opac_vals).unsqueeze(-1) * col_vals
|
|
338
|
+
comp = comp + comp_add
|
|
339
|
+
remaining = remaining - (mask_done_f * remaining * opac_vals)
|
|
340
|
+
|
|
341
|
+
# Reset old segment where mask_done is True
|
|
342
|
+
# cur_mat = -1, acc_thick = 0
|
|
343
|
+
# We'll do it by `torch.where(mask, val_if_true, val_if_false)`
|
|
344
|
+
cur_mat = torch.where(mask_done, torch.full_like(cur_mat, -1), cur_mat)
|
|
345
|
+
acc_thick = torch.where(mask_done, torch.zeros_like(acc_thick), acc_thick)
|
|
346
|
+
|
|
347
|
+
# ---------------------------------------------------------------------
|
|
348
|
+
# (B) For pixels that print this layer:
|
|
349
|
+
# - Start a new segment if cur_mat == -1
|
|
350
|
+
# - Accumulate thickness if cur_mat == layer_mat
|
|
351
|
+
# ---------------------------------------------------------------------
|
|
258
352
|
eff_thick = p_print.to(torch.float32) * h
|
|
259
353
|
|
|
260
|
-
# Start new segment
|
|
261
|
-
mask_new = p_print & (cur_mat
|
|
262
|
-
|
|
263
|
-
cur_mat[mask_new] = layer_mat
|
|
264
|
-
acc_thick[mask_new] = eff_thick[mask_new]
|
|
354
|
+
# (B1) Start new segment where cur_mat == -1
|
|
355
|
+
mask_new = p_print & (cur_mat.eq(-1))
|
|
356
|
+
mask_new_f = mask_new.to(torch.float32)
|
|
265
357
|
|
|
266
|
-
#
|
|
267
|
-
|
|
268
|
-
|
|
269
|
-
acc_thick[mask_same] += eff_thick[mask_same]
|
|
358
|
+
# Set cur_mat to layer_mat if mask_new is True
|
|
359
|
+
# (layer_mat is shape [], so it will broadcast)
|
|
360
|
+
cur_mat = torch.where(mask_new, layer_mat, cur_mat)
|
|
270
361
|
|
|
362
|
+
# We add thickness:
|
|
363
|
+
acc_thick = acc_thick + mask_new_f * eff_thick
|
|
271
364
|
|
|
365
|
+
# (B2) Accumulate thickness where cur_mat == layer_mat
|
|
366
|
+
# We do this in a second mask to avoid confusion, but you can combine.
|
|
367
|
+
mask_same = p_print & (cur_mat.eq(layer_mat))
|
|
368
|
+
acc_thick = acc_thick + (mask_same.to(torch.float32) * eff_thick)
|
|
272
369
|
|
|
273
|
-
#
|
|
274
|
-
|
|
275
|
-
|
|
276
|
-
|
|
277
|
-
|
|
278
|
-
col_vals = material_colors[old_inds]
|
|
370
|
+
# -------------------------------------------------------------------------
|
|
371
|
+
# 5) After the loop, composite any remaining segments (cur_mat != -1).
|
|
372
|
+
# -------------------------------------------------------------------------
|
|
373
|
+
mask_remain = cur_mat.ne(-1)
|
|
374
|
+
mask_remain_f = mask_remain.to(torch.float32)
|
|
279
375
|
|
|
280
|
-
|
|
281
|
-
|
|
282
|
-
|
|
376
|
+
old_inds_clamped = torch.clamp(cur_mat, min=0)
|
|
377
|
+
td_vals = material_TDs[old_inds_clamped]
|
|
378
|
+
col_vals = material_colors[old_inds_clamped]
|
|
379
|
+
|
|
380
|
+
thick_ratio = acc_thick / td_vals
|
|
381
|
+
opac_vals = o + (A * torch.log1p(k * thick_ratio) + b * thick_ratio)
|
|
382
|
+
opac_vals = torch.clamp(opac_vals, 0.0, 1.0)
|
|
283
383
|
|
|
284
|
-
|
|
285
|
-
|
|
286
|
-
|
|
287
|
-
remaining[mask_remain] *= 1.0 - opac_vals
|
|
384
|
+
comp_add = (mask_remain_f * remaining * opac_vals).unsqueeze(-1) * col_vals
|
|
385
|
+
comp = comp + comp_add
|
|
386
|
+
remaining = remaining - (mask_remain_f * remaining * opac_vals)
|
|
288
387
|
|
|
289
388
|
# -------------------------------------------------------------------------
|
|
290
|
-
# Composite background
|
|
389
|
+
# 6) Composite background
|
|
291
390
|
# -------------------------------------------------------------------------
|
|
292
|
-
comp
|
|
391
|
+
comp = comp + remaining.unsqueeze(-1) * background
|
|
293
392
|
return comp * 255.0
|
|
294
393
|
|
|
295
394
|
|
|
296
395
|
@torch.jit.script
|
|
297
|
-
def
|
|
298
|
-
|
|
396
|
+
def composite_image_combined(
|
|
397
|
+
pixel_height_logits: torch.Tensor, # [H,W]
|
|
398
|
+
global_logits: torch.Tensor, # [max_layers, n_materials]
|
|
399
|
+
tau_height: float,
|
|
400
|
+
tau_global: float,
|
|
401
|
+
h: float,
|
|
402
|
+
max_layers: int,
|
|
403
|
+
material_colors: torch.Tensor, # [n_materials, 3]
|
|
404
|
+
material_TDs: torch.Tensor, # [n_materials]
|
|
405
|
+
background: torch.Tensor, # [3]
|
|
406
|
+
rng_seed: int = -1,
|
|
299
407
|
) -> torch.Tensor:
|
|
300
408
|
"""
|
|
301
|
-
|
|
409
|
+
Combine continuous and discrete compositing over all pixels.
|
|
410
|
+
|
|
411
|
+
Args:
|
|
412
|
+
pixel_height_logits (torch.Tensor): Logits for pixel heights, shape [H, W].
|
|
413
|
+
global_logits (torch.Tensor): Logits for global material assignment, shape [max_layers, n_materials].
|
|
414
|
+
tau_height (float): Temperature parameter for height rounding.
|
|
415
|
+
tau_global (float): Temperature parameter for global material assignment.
|
|
416
|
+
h (float): Height of each layer.
|
|
417
|
+
max_layers (int): Maximum number of layers.
|
|
418
|
+
material_colors (torch.Tensor): Tensor of material colors, shape [n_materials, 3].
|
|
419
|
+
material_TDs (torch.Tensor): Tensor of material transmission/opacity parameters, shape [n_materials].
|
|
420
|
+
background (torch.Tensor): Background color tensor, shape [3].
|
|
421
|
+
rng_seed (int, optional): Random seed for deterministic sampling. Defaults to -1.
|
|
422
|
+
|
|
423
|
+
Returns:
|
|
424
|
+
torch.Tensor: Composite image tensor, shape [H, W, 3].
|
|
302
425
|
"""
|
|
303
|
-
|
|
304
|
-
|
|
305
|
-
|
|
306
|
-
|
|
307
|
-
|
|
308
|
-
|
|
309
|
-
|
|
426
|
+
cont = composite_image_cont(
|
|
427
|
+
pixel_height_logits,
|
|
428
|
+
global_logits,
|
|
429
|
+
tau_height,
|
|
430
|
+
tau_global,
|
|
431
|
+
h,
|
|
432
|
+
max_layers,
|
|
433
|
+
material_colors,
|
|
434
|
+
material_TDs,
|
|
435
|
+
background,
|
|
436
|
+
)
|
|
437
|
+
if tau_global < 1.0:
|
|
438
|
+
disc = composite_image_disc(
|
|
439
|
+
pixel_height_logits,
|
|
440
|
+
global_logits,
|
|
441
|
+
tau_height,
|
|
442
|
+
tau_global,
|
|
443
|
+
h,
|
|
444
|
+
max_layers,
|
|
445
|
+
material_colors,
|
|
446
|
+
material_TDs,
|
|
447
|
+
background,
|
|
448
|
+
rng_seed,
|
|
449
|
+
)
|
|
450
|
+
return cont * tau_global + disc * (1 - tau_global)
|
|
310
451
|
else:
|
|
311
|
-
|
|
312
|
-
hard_round = torch.round(x)
|
|
313
|
-
floor_val = torch.floor(x)
|
|
314
|
-
diff = x - floor_val
|
|
315
|
-
soft_round = floor_val + torch.sigmoid((diff - 0.5) / temp)
|
|
316
|
-
return ratio * soft_round + (1 - ratio) * hard_round
|
|
452
|
+
return cont
|
|
317
453
|
|
|
318
454
|
|
|
319
455
|
def discretize_solution(
|
|
@@ -321,6 +457,18 @@ def discretize_solution(
|
|
|
321
457
|
):
|
|
322
458
|
"""
|
|
323
459
|
Convert continuous logs to discrete layer counts and discrete color IDs.
|
|
460
|
+
|
|
461
|
+
Args:
|
|
462
|
+
params (dict): Dictionary containing the parameters 'pixel_height_logits' and 'global_logits'.
|
|
463
|
+
tau_global (float): Temperature parameter for global material assignment.
|
|
464
|
+
h (float): Height of each layer.
|
|
465
|
+
max_layers (int): Maximum number of layers.
|
|
466
|
+
rng_seed (int, optional): Random seed for deterministic sampling. Defaults to -1.
|
|
467
|
+
|
|
468
|
+
Returns:
|
|
469
|
+
tuple: A tuple containing:
|
|
470
|
+
- torch.Tensor: Discrete global material assignments, shape [max_layers].
|
|
471
|
+
- torch.Tensor: Discrete height image, shape [H, W].
|
|
324
472
|
"""
|
|
325
473
|
pixel_height_logits = params["pixel_height_logits"]
|
|
326
474
|
global_logits = params["global_logits"]
|
|
@@ -347,10 +495,10 @@ def initialize_pixel_height_logits(target):
|
|
|
347
495
|
Uses the formula: L = 0.299*R + 0.587*G + 0.114*B.
|
|
348
496
|
|
|
349
497
|
Args:
|
|
350
|
-
target (
|
|
498
|
+
target (np.ndarray): The target image array with shape (H, W, 3).
|
|
351
499
|
|
|
352
500
|
Returns:
|
|
353
|
-
|
|
501
|
+
np.ndarray: The initialized pixel height logits.
|
|
354
502
|
"""
|
|
355
503
|
# Compute normalized luminance in [0,1]
|
|
356
504
|
normalized_lum = (
|
|
@@ -730,6 +878,10 @@ def init_height_map(target, max_layers, h, eps=1e-6, random_seed=None):
|
|
|
730
878
|
|
|
731
879
|
Args:
|
|
732
880
|
target (jnp.ndarray): The target image array with shape (H, W, 3).
|
|
881
|
+
max_layers (int): Maximum number of layers.
|
|
882
|
+
h (float): Height of each layer.
|
|
883
|
+
eps (float): Small constant to avoid division by zero.
|
|
884
|
+
random_seed (int): Random seed for reproducibility.
|
|
733
885
|
|
|
734
886
|
Returns:
|
|
735
887
|
jnp.ndarray: The initialized pixel height logits.
|
|
@@ -20,7 +20,6 @@ from autoforge.Helper.OutputHelper import (
|
|
|
20
20
|
generate_project_file,
|
|
21
21
|
)
|
|
22
22
|
from autoforge.Helper.PruningHelper import prune_redundant_layers
|
|
23
|
-
from autoforge.Loss.PerceptionLoss import MultiLayerVGGPerceptualLoss
|
|
24
23
|
from autoforge.Modules.Optimizer import FilamentOptimizer
|
|
25
24
|
|
|
26
25
|
|
|
@@ -282,7 +281,8 @@ def main():
|
|
|
282
281
|
) # shape (H', W') cause of weird opencv quirk
|
|
283
282
|
|
|
284
283
|
# VGG Perceptual Loss
|
|
285
|
-
|
|
284
|
+
# We currently disable this as it is not used in the optimization.
|
|
285
|
+
perception_loss_module = None # MultiLayerVGGPerceptualLoss().to(device).eval()
|
|
286
286
|
|
|
287
287
|
# Create an optimizer instance
|
|
288
288
|
optimizer = FilamentOptimizer(
|
|
@@ -300,7 +300,7 @@ def main():
|
|
|
300
300
|
print("Starting optimization...")
|
|
301
301
|
tbar = tqdm(range(args.iterations))
|
|
302
302
|
for i in tbar:
|
|
303
|
-
loss_val = optimizer.step(record_best=i %
|
|
303
|
+
loss_val = optimizer.step(record_best=i % 10 == 0)
|
|
304
304
|
|
|
305
305
|
optimizer.visualize(interval=25)
|
|
306
306
|
optimizer.log_to_tensorboard(interval=100)
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|