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.
Files changed (35) hide show
  1. {autoforge-1.5.1 → autoforge-1.5.2}/PKG-INFO +1 -1
  2. {autoforge-1.5.1 → autoforge-1.5.2}/pyproject.toml +1 -1
  3. {autoforge-1.5.1 → autoforge-1.5.2}/src/AutoForge.egg-info/PKG-INFO +1 -1
  4. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Helper/OptimizerHelper.py +269 -117
  5. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/auto_forge.py +3 -3
  6. {autoforge-1.5.1 → autoforge-1.5.2}/README.md +0 -0
  7. {autoforge-1.5.1 → autoforge-1.5.2}/setup.cfg +0 -0
  8. {autoforge-1.5.1 → autoforge-1.5.2}/src/AutoForge.egg-info/SOURCES.txt +0 -0
  9. {autoforge-1.5.1 → autoforge-1.5.2}/src/AutoForge.egg-info/dependency_links.txt +0 -0
  10. {autoforge-1.5.1 → autoforge-1.5.2}/src/AutoForge.egg-info/entry_points.txt +0 -0
  11. {autoforge-1.5.1 → autoforge-1.5.2}/src/AutoForge.egg-info/requires.txt +0 -0
  12. {autoforge-1.5.1 → autoforge-1.5.2}/src/AutoForge.egg-info/top_level.txt +0 -0
  13. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Gui/Views/__init__.py +0 -0
  14. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Gui/Views/backend_database.py +0 -0
  15. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Gui/Views/filament_editor.py +0 -0
  16. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Gui/__init__.py +0 -0
  17. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Gui/app.py +0 -0
  18. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Gui/main.py +0 -0
  19. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Helper/CAdamW.py +0 -0
  20. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Helper/FilamentHelper.py +0 -0
  21. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Helper/ImageHelper.py +0 -0
  22. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Helper/OutputHelper.py +0 -0
  23. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Helper/PruningHelper.py +0 -0
  24. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Helper/__init__.py +0 -0
  25. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Helper/filamentcolors_library.py +0 -0
  26. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Loss/LossFunctions.py +0 -0
  27. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Loss/PerceptionLoss.py +0 -0
  28. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Loss/__init__.py +0 -0
  29. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Misc/__init__.py +0 -0
  30. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Misc/color_fit_all.py +0 -0
  31. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Modules/Optimizer.py +0 -0
  32. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/Modules/__init__.py +0 -0
  33. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/__init__.py +0 -0
  34. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/__main__.py +0 -0
  35. {autoforge-1.5.1 → autoforge-1.5.2}/src/autoforge/old_jax_files.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.2
2
2
  Name: AutoForge
3
- Version: 1.5.1
3
+ Version: 1.5.2
4
4
  Summary: AutoForge is a Python tool for generating 3D printed layered models from an input image.
5
5
  Author-email: Hendric Voss <hvoss@techfak.uni-bielefeld.de>
6
6
  License: CC BY-NC-SA 4.0
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "AutoForge"
7
- version="1.5.1"
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"
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.2
2
2
  Name: AutoForge
3
- Version: 1.5.1
3
+ Version: 1.5.2
4
4
  Summary: AutoForge is a Python tool for generating 3D printed layered models from an input image.
5
5
  Author-email: Hendric Voss <hvoss@techfak.uni-bielefeld.de>
6
6
  License: CC BY-NC-SA 4.0
@@ -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
- # opacity function parameters
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
- p_print = (discrete_layers > layer_idx).float() # [H,W]
118
- eff_thick = p_print * h # effective thickness per pixel
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
- color_i = torch.matmul(p_i, material_colors)
127
- TD_i = torch.matmul(p_i, material_TDs)
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
- comp = comp + ((remaining * opac).unsqueeze(-1) * color_i)
134
- remaining = remaining * (1 - opac)
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
- An optimized rewrite of the original compositing code.
155
-
156
- The main changes:
157
- - We combine the old (mask_diff) and (mask_not_printed) logic
158
- into a single step that composites any pixel which is 'done'
159
- with its current material (either because it's not printing
160
- this layer or because the layer material is different).
161
- - Then we update (or start) new segments for any pixel that
162
- actually prints this layer and whose material matches or
163
- doesn't match.
164
- - Fewer conditionals, fewer index_select calls, and no repeated
165
- if mask.any(): checks.
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
- # Determine how many discrete layers each pixel prints
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
- # Example "adaptive rounding" approach
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
- # Precompute the chosen material for each layer
182
- # (One material per layer for the entire image.)
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
- # Tracks the material currently being accumulated (or -1 for none)
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
- # Tracks the accumulated thickness for the current segment
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
- # Material chosen for this layer
222
- layer_mat = new_mats[layer_idx] # a single int
223
- # Which pixels actually print this layer?
224
- p_print = discrete_layers > layer_idx # [H,W] bool
225
-
226
- # --------------------------------------------------
227
- # (A) Composite any "finished" segment, i.e. the pixel
228
- # is *done* with its old segment if:
229
- # (1) we already had a segment (cur_mat != -1), AND
230
- # (2) either the pixel is not printing this layer,
231
- # or the pixel's new layer material != cur_mat
232
- # --------------------------------------------------
233
- mask_done = (cur_mat != -1) & (~p_print | (cur_mat != layer_mat))
234
- if mask_done.any():
235
- old_inds = cur_mat[mask_done].to(torch.int64)
236
- td_vals = material_TDs[old_inds] # gather thickness densities
237
- col_vals = material_colors[old_inds] # gather old material colors
238
-
239
- thick_ratio = acc_thick[mask_done] / td_vals
240
- opac_vals = o + (A * torch.log1p(k * thick_ratio) + b * thick_ratio)
241
- opac_vals = torch.clamp(opac_vals, 0.0, 1.0)
242
-
243
- # composite old segment
244
- comp[mask_done] += (remaining[mask_done] * opac_vals).unsqueeze(
245
- -1
246
- ) * col_vals
247
- remaining[mask_done] *= 1.0 - opac_vals
248
-
249
- # Reset old segment
250
- cur_mat[mask_done] = -1
251
- acc_thick[mask_done] = 0.0
252
-
253
- # --------------------------------------------------
254
- # (B) For pixels that do print this layer:
255
- # - start a new segment if cur_mat == -1
256
- # - accumulate thickness if cur_mat == layer_mat
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 == -1)
262
- if mask_new.any():
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
- # Accumulate thickness
267
- mask_same = p_print & (cur_mat == layer_mat)
268
- if mask_same.any():
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
- # After the loop, composite any still-pending segments
274
- mask_remain = cur_mat != -1
275
- if mask_remain.any():
276
- old_inds = cur_mat[mask_remain].to(torch.int64)
277
- td_vals = material_TDs[old_inds]
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
- thick_ratio = acc_thick[mask_remain] / td_vals
281
- opac_vals = o + (A * torch.log1p(k * thick_ratio) + b * thick_ratio)
282
- opac_vals = torch.clamp(opac_vals, 0.0, 1.0)
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
- comp[mask_remain] += (remaining[mask_remain] * opac_vals).unsqueeze(
285
- -1
286
- ) * col_vals
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 += remaining.unsqueeze(-1) * background
391
+ comp = comp + remaining.unsqueeze(-1) * background
293
392
  return comp * 255.0
294
393
 
295
394
 
296
395
  @torch.jit.script
297
- def adaptive_round(
298
- x: torch.Tensor, tau: float, high_tau: float, low_tau: float, temp: float
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
- Smooth rounding based on temperature 'tau'.
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
- if tau <= low_tau:
304
- return torch.round(x)
305
- elif tau >= high_tau:
306
- floor_val = torch.floor(x)
307
- diff = x - floor_val
308
- soft_round = floor_val + torch.sigmoid((diff - 0.5) / temp)
309
- return soft_round
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
- ratio = (tau - low_tau) / (high_tau - low_tau)
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 (jnp.ndarray): The target image array with shape (H, W, 3).
498
+ target (np.ndarray): The target image array with shape (H, W, 3).
351
499
 
352
500
  Returns:
353
- jnp.ndarray: The initialized pixel height logits.
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
- perception_loss_module = MultiLayerVGGPerceptualLoss().to(device).eval()
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 % 3 == 0)
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