autoforge 1.3.0__tar.gz → 1.3.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 (31) hide show
  1. {autoforge-1.3.0 → autoforge-1.3.2}/PKG-INFO +5 -4
  2. {autoforge-1.3.0 → autoforge-1.3.2}/README.md +4 -3
  3. {autoforge-1.3.0 → autoforge-1.3.2}/pyproject.toml +1 -1
  4. {autoforge-1.3.0 → autoforge-1.3.2}/src/AutoForge.egg-info/PKG-INFO +5 -4
  5. {autoforge-1.3.0 → autoforge-1.3.2}/src/AutoForge.egg-info/SOURCES.txt +2 -1
  6. autoforge-1.3.2/src/autoforge/Helper/ImageHelper.py +22 -0
  7. {autoforge-1.3.0 → autoforge-1.3.2}/src/autoforge/Helper/OptimizerHelper.py +111 -6
  8. {autoforge-1.3.0 → autoforge-1.3.2}/src/autoforge/Helper/OutputHelper.py +65 -52
  9. {autoforge-1.3.0 → autoforge-1.3.2}/src/autoforge/Loss/LossFunctions.py +2 -0
  10. autoforge-1.3.2/src/autoforge/Misc/color_fit_all.py +1877 -0
  11. {autoforge-1.3.0 → autoforge-1.3.2}/src/autoforge/Modules/Optimizer.py +2 -7
  12. autoforge-1.3.2/src/autoforge/__init__.py +0 -0
  13. {autoforge-1.3.0 → autoforge-1.3.2}/src/autoforge/auto_forge.py +32 -19
  14. {autoforge-1.3.0 → autoforge-1.3.2}/src/autoforge/auto_forge_gradio.py +1 -1
  15. autoforge-1.3.0/src/autoforge/Helper/ImageHelper.py +0 -13
  16. autoforge-1.3.0/src/autoforge/color_fit_bayesian.py +0 -225
  17. {autoforge-1.3.0 → autoforge-1.3.2}/setup.cfg +0 -0
  18. {autoforge-1.3.0 → autoforge-1.3.2}/src/AutoForge.egg-info/dependency_links.txt +0 -0
  19. {autoforge-1.3.0 → autoforge-1.3.2}/src/AutoForge.egg-info/entry_points.txt +0 -0
  20. {autoforge-1.3.0 → autoforge-1.3.2}/src/AutoForge.egg-info/requires.txt +0 -0
  21. {autoforge-1.3.0 → autoforge-1.3.2}/src/AutoForge.egg-info/top_level.txt +0 -0
  22. {autoforge-1.3.0 → autoforge-1.3.2}/src/autoforge/Helper/FilamentHelper.py +0 -0
  23. {autoforge-1.3.0 → autoforge-1.3.2}/src/autoforge/Helper/PruningHelper.py +0 -0
  24. {autoforge-1.3.0 → autoforge-1.3.2}/src/autoforge/Helper/__init__.py +0 -0
  25. {autoforge-1.3.0 → autoforge-1.3.2}/src/autoforge/Loss/PerceptionLoss.py +0 -0
  26. {autoforge-1.3.0 → autoforge-1.3.2}/src/autoforge/Loss/__init__.py +0 -0
  27. {autoforge-1.3.0/src/autoforge/Modules → autoforge-1.3.2/src/autoforge/Misc}/__init__.py +0 -0
  28. {autoforge-1.3.0/src/autoforge → autoforge-1.3.2/src/autoforge/Modules}/__init__.py +0 -0
  29. {autoforge-1.3.0 → autoforge-1.3.2}/src/autoforge/__main__.py +0 -0
  30. {autoforge-1.3.0 → autoforge-1.3.2}/src/autoforge/filamentcolors_library.py +0 -0
  31. {autoforge-1.3.0 → autoforge-1.3.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.3.0
3
+ Version: 1.3.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
@@ -30,6 +30,10 @@ AutoForge is a Python tool for generating 3D printed layered models from an inpu
30
30
  We recently switched from JAX to PyTorch, to allow for a more streamlined development process. \
31
31
  This can have some unforeseen consequences, so please report any bugs you find.
32
32
 
33
+ This release should give better gpu support for more users, especially on windows. \
34
+ If you have problems running the code on your gpu, please refer to the [Pytorch Homepage](https://pytorch.org/) for help. \
35
+ Both CUDA and ROCm are supported, but you need to install the correct version of pytorch for your system.
36
+
33
37
  ## Example
34
38
  All examples use only the 13 BambuLab Basic filaments, currently available in Hueforge.
35
39
  <div style="display: flex; justify-content: center; gap: 20px;">
@@ -95,8 +99,6 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
95
99
  - `--csv_file`: **(Required)** Path to the CSV file containing material data. The CSV should include columns for the brand, name, color (hex code), and TD values.
96
100
  - `--output_folder`: Folder where output files will be saved. (default: `./outputs`).
97
101
  - `--iterations`: Number of optimization iterations (default: 5000).
98
- - `--best_loss_iterations`: Percentage of optimization iterations after which we start to record the best loss (default: 4000). \
99
- **Note:** We currently do this as we increase a penalty loss for color and color swaps as the time goes on. Without this the penalties would have no impact.
100
102
  - `--learning_rate`: Learning rate for the optimizer (default: 1e-2).
101
103
  - `--layer_height`: Layer thickness in millimeters (default: 0.04).
102
104
  - `--max_layers`: Maximum number of layers (default: 75). \
@@ -160,7 +162,6 @@ Although I would love to be it fully color compatible with hueforge, I don't thi
160
162
  AutoForge © 2025 by Hendric Voss is licensed under [CC BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/).
161
163
  The software is provided as-is and comes with no warranty or guarantee of support.
162
164
 
163
- The
164
165
 
165
166
  ## Acknowledgements
166
167
 
@@ -10,6 +10,10 @@ AutoForge is a Python tool for generating 3D printed layered models from an inpu
10
10
  We recently switched from JAX to PyTorch, to allow for a more streamlined development process. \
11
11
  This can have some unforeseen consequences, so please report any bugs you find.
12
12
 
13
+ This release should give better gpu support for more users, especially on windows. \
14
+ If you have problems running the code on your gpu, please refer to the [Pytorch Homepage](https://pytorch.org/) for help. \
15
+ Both CUDA and ROCm are supported, but you need to install the correct version of pytorch for your system.
16
+
13
17
  ## Example
14
18
  All examples use only the 13 BambuLab Basic filaments, currently available in Hueforge.
15
19
  <div style="display: flex; justify-content: center; gap: 20px;">
@@ -75,8 +79,6 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
75
79
  - `--csv_file`: **(Required)** Path to the CSV file containing material data. The CSV should include columns for the brand, name, color (hex code), and TD values.
76
80
  - `--output_folder`: Folder where output files will be saved. (default: `./outputs`).
77
81
  - `--iterations`: Number of optimization iterations (default: 5000).
78
- - `--best_loss_iterations`: Percentage of optimization iterations after which we start to record the best loss (default: 4000). \
79
- **Note:** We currently do this as we increase a penalty loss for color and color swaps as the time goes on. Without this the penalties would have no impact.
80
82
  - `--learning_rate`: Learning rate for the optimizer (default: 1e-2).
81
83
  - `--layer_height`: Layer thickness in millimeters (default: 0.04).
82
84
  - `--max_layers`: Maximum number of layers (default: 75). \
@@ -140,7 +142,6 @@ Although I would love to be it fully color compatible with hueforge, I don't thi
140
142
  AutoForge © 2025 by Hendric Voss is licensed under [CC BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/).
141
143
  The software is provided as-is and comes with no warranty or guarantee of support.
142
144
 
143
- The
144
145
 
145
146
  ## Acknowledgements
146
147
 
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "AutoForge"
7
- version="1.3.0"
7
+ version="1.3.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.3.0
3
+ Version: 1.3.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
@@ -30,6 +30,10 @@ AutoForge is a Python tool for generating 3D printed layered models from an inpu
30
30
  We recently switched from JAX to PyTorch, to allow for a more streamlined development process. \
31
31
  This can have some unforeseen consequences, so please report any bugs you find.
32
32
 
33
+ This release should give better gpu support for more users, especially on windows. \
34
+ If you have problems running the code on your gpu, please refer to the [Pytorch Homepage](https://pytorch.org/) for help. \
35
+ Both CUDA and ROCm are supported, but you need to install the correct version of pytorch for your system.
36
+
33
37
  ## Example
34
38
  All examples use only the 13 BambuLab Basic filaments, currently available in Hueforge.
35
39
  <div style="display: flex; justify-content: center; gap: 20px;">
@@ -95,8 +99,6 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
95
99
  - `--csv_file`: **(Required)** Path to the CSV file containing material data. The CSV should include columns for the brand, name, color (hex code), and TD values.
96
100
  - `--output_folder`: Folder where output files will be saved. (default: `./outputs`).
97
101
  - `--iterations`: Number of optimization iterations (default: 5000).
98
- - `--best_loss_iterations`: Percentage of optimization iterations after which we start to record the best loss (default: 4000). \
99
- **Note:** We currently do this as we increase a penalty loss for color and color swaps as the time goes on. Without this the penalties would have no impact.
100
102
  - `--learning_rate`: Learning rate for the optimizer (default: 1e-2).
101
103
  - `--layer_height`: Layer thickness in millimeters (default: 0.04).
102
104
  - `--max_layers`: Maximum number of layers (default: 75). \
@@ -160,7 +162,6 @@ Although I would love to be it fully color compatible with hueforge, I don't thi
160
162
  AutoForge © 2025 by Hendric Voss is licensed under [CC BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/).
161
163
  The software is provided as-is and comes with no warranty or guarantee of support.
162
164
 
163
- The
164
165
 
165
166
  ## Acknowledgements
166
167
 
@@ -10,7 +10,6 @@ src/autoforge/__init__.py
10
10
  src/autoforge/__main__.py
11
11
  src/autoforge/auto_forge.py
12
12
  src/autoforge/auto_forge_gradio.py
13
- src/autoforge/color_fit_bayesian.py
14
13
  src/autoforge/filamentcolors_library.py
15
14
  src/autoforge/old_jax_files.py
16
15
  src/autoforge/Helper/FilamentHelper.py
@@ -22,5 +21,7 @@ src/autoforge/Helper/__init__.py
22
21
  src/autoforge/Loss/LossFunctions.py
23
22
  src/autoforge/Loss/PerceptionLoss.py
24
23
  src/autoforge/Loss/__init__.py
24
+ src/autoforge/Misc/__init__.py
25
+ src/autoforge/Misc/color_fit_all.py
25
26
  src/autoforge/Modules/Optimizer.py
26
27
  src/autoforge/Modules/__init__.py
@@ -0,0 +1,22 @@
1
+ import cv2
2
+
3
+
4
+ def resize_image(img, max_size):
5
+ h_img, w_img, _ = img.shape
6
+
7
+ # Compute the scaling factor based on the larger dimension
8
+ if w_img >= h_img:
9
+ scale = max_size / w_img
10
+ else:
11
+ scale = max_size / h_img
12
+
13
+ # Compute new dimensions and round them
14
+ new_w = int(round(w_img * scale))
15
+ new_h = int(round(h_img * scale))
16
+
17
+ # Resize the image with an appropriate interpolation method
18
+ img_out = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
19
+ # INTER_CUBIC is actually not a good choice for shrinking images. Except in our case.
20
+ # As we shrink the image, we lose information and the image becomes blurry with INTER_AREA. This is detrimental for the solver. If you want to try it out be my guest.
21
+ # INTER_AREA actually destroys a lot of the colors of our solver.
22
+ return img_out
@@ -1,3 +1,4 @@
1
+ import random
1
2
  from itertools import permutations
2
3
 
3
4
  import numpy as np
@@ -111,9 +112,10 @@ def composite_image(
111
112
  comp = torch.zeros(H, W, 3, dtype=torch.float32, device=pixel_height.device)
112
113
  remaining = torch.ones(H, W, dtype=torch.float32, device=pixel_height.device)
113
114
 
114
- A = 0.1215
115
- k = 61.6970
116
- b = 0.4773
115
+ o = -1.1021105e-02
116
+ A = 1.0786635e00
117
+ k = 3.0878708e01
118
+ b = -5.0812101e00
117
119
 
118
120
  for i in range(max_layers):
119
121
  layer_idx = max_layers - 1 - i
@@ -147,8 +149,11 @@ def composite_image(
147
149
  color_i = torch.matmul(p_i, material_colors)
148
150
  TD_i = torch.matmul(p_i, material_TDs) * 0.1
149
151
  TD_i = torch.clamp(TD_i, 1e-8, 1e8)
150
- opac = A * torch.log1p(k * (eff_thick / TD_i)) + b * (eff_thick / TD_i)
152
+
153
+ opac = o + (A * torch.log1p(k * (eff_thick / TD_i)) + b * (eff_thick / TD_i))
151
154
  opac = torch.clamp(opac, 0.0, 1.0)
155
+ # Setting opac to 1.0 for the layers where the thickness is less than 2*h to avoid single layer artifacts
156
+ opac = torch.where(TD_i <= h * 2, torch.tensor(1.0), opac)
152
157
 
153
158
  comp = comp + ((remaining * opac).unsqueeze(-1) * color_i)
154
159
  remaining = remaining * (1 - opac)
@@ -487,6 +492,104 @@ def init_height_map_depth_color_adjusted(
487
492
  return pixel_height_logits
488
493
 
489
494
 
495
+ def tsp_simulated_annealing(
496
+ band_reps,
497
+ start_band,
498
+ end_band,
499
+ initial_order=None,
500
+ initial_temp=1000,
501
+ cooling_rate=0.995,
502
+ num_iter=10000,
503
+ ):
504
+ """
505
+ Solve the band ordering problem using simulated annealing.
506
+
507
+ Args:
508
+ band_reps (list or np.ndarray): List of Lab color representations.
509
+ start_band (int): Index for the darkest band.
510
+ end_band (int): Index for the brightest band.
511
+ initial_order (list, optional): Initial ordering of band indices.
512
+ initial_temp (float): Starting temperature.
513
+ cooling_rate (float): Factor to cool the temperature.
514
+ num_iter (int): Maximum number of iterations.
515
+
516
+ Returns:
517
+ list: An ordering of band indices from start_band to end_band.
518
+ """
519
+ if initial_order is None:
520
+ # Use a simple ordering: start, middle bands as given, then end.
521
+ middle_indices = [
522
+ i for i in range(len(band_reps)) if i not in (start_band, end_band)
523
+ ]
524
+ order = [start_band] + middle_indices + [end_band]
525
+ else:
526
+ order = initial_order.copy()
527
+
528
+ def total_distance(order):
529
+ return sum(
530
+ np.linalg.norm(band_reps[order[i]] - band_reps[order[i + 1]])
531
+ for i in range(len(order) - 1)
532
+ )
533
+
534
+ current_distance = total_distance(order)
535
+ best_order = order.copy()
536
+ best_distance = current_distance
537
+ temp = initial_temp
538
+
539
+ for _ in range(num_iter):
540
+ # Randomly swap two indices in the middle of the order
541
+ new_order = order.copy()
542
+ idx1, idx2 = random.sample(range(1, len(order) - 1), 2)
543
+ new_order[idx1], new_order[idx2] = new_order[idx2], new_order[idx1]
544
+
545
+ new_distance = total_distance(new_order)
546
+ delta = new_distance - current_distance
547
+
548
+ # Accept the new order if it improves or with a probability to escape local minima
549
+ if delta < 0 or np.exp(-delta / temp) > random.random():
550
+ order = new_order.copy()
551
+ current_distance = new_distance
552
+ if current_distance < best_distance:
553
+ best_order = order.copy()
554
+ best_distance = current_distance
555
+
556
+ temp *= cooling_rate
557
+ if temp < 1e-6:
558
+ break
559
+ return best_order
560
+
561
+
562
+ def choose_optimal_num_bands(centroids, min_bands=2, max_bands=15, random_seed=None):
563
+ """
564
+ Determine the optimal number of clusters (bands) for the centroids
565
+ by maximizing the silhouette score.
566
+
567
+ Args:
568
+ centroids (np.ndarray): Array of centroid colors (e.g., shape (n_clusters, 3)).
569
+ min_bands (int): Minimum number of clusters to try.
570
+ max_bands (int): Maximum number of clusters to try.
571
+ random_seed (int, optional): Random seed for reproducibility.
572
+
573
+ Returns:
574
+ int: Optimal number of bands.
575
+ """
576
+ best_num = min_bands
577
+ best_score = -1
578
+
579
+ for num in range(min_bands, max_bands + 1):
580
+ kmeans = KMeans(n_clusters=num, random_state=random_seed).fit(centroids)
581
+ labels = kmeans.labels_
582
+ # If there's only one unique label, skip to avoid errors.
583
+ if len(np.unique(labels)) < 2:
584
+ continue
585
+ score = silhouette_score(centroids, labels)
586
+ if score > best_score:
587
+ best_score = score
588
+ best_num = num
589
+
590
+ return best_num
591
+
592
+
490
593
  def init_height_map(target, max_layers, h, eps=1e-6, random_seed=None):
491
594
  """
492
595
  Initialize pixel height logits based on the luminance of the target image.
@@ -512,7 +615,9 @@ def init_height_map(target, max_layers, h, eps=1e-6, random_seed=None):
512
615
  return 0.299 * col[0] + 0.587 * col[1] + 0.114 * col[2]
513
616
 
514
617
  # --- Step 2: Second clustering of centroids into bands ---
515
- num_bands = 9
618
+ num_bands = choose_optimal_num_bands(
619
+ centroids, min_bands=5, max_bands=10, random_seed=random_seed
620
+ )
516
621
  band_kmeans = KMeans(n_clusters=num_bands, random_state=random_seed).fit(centroids)
517
622
  band_labels = band_kmeans.labels_
518
623
 
@@ -589,7 +694,7 @@ def init_height_map(target, max_layers, h, eps=1e-6, random_seed=None):
589
694
 
590
695
  new_labels = new_labels.astype(np.float32) / new_labels.max()
591
696
 
592
- normalized_lum = np.array(new_labels, dtype=np.float64)
697
+ normalized_lum = np.array(new_labels, dtype=np.float32)
593
698
  # convert out to inverse sigmoid logit function
594
699
  pixel_height_logits = np.log((normalized_lum + eps) / (1 - normalized_lum + eps))
595
700
 
@@ -188,85 +188,98 @@ def generate_project_file(
188
188
 
189
189
  def generate_stl(height_map, filename, background_height, scale=1.0):
190
190
  """
191
- Generate a binary STL file from a height map.
191
+ Generate a binary STL file from a height map with shared boundary vertices
192
+ to fix non-manifold edge warnings.
192
193
 
193
194
  Args:
194
195
  height_map (np.ndarray): 2D array representing the height map.
195
196
  filename (str): The name of the output STL file.
196
197
  background_height (float): The height of the background in the STL model.
197
- scale (float, optional): Scale factor for the x and y dimensions. Defaults to 1.0.
198
+ scale (float, optional): Scale factor for the x dimension (and optionally y). Defaults to 1.0.
198
199
  """
199
200
  H, W = height_map.shape
200
- vertices = np.zeros((H, W, 3), dtype=np.float32)
201
+
202
+ # Create the top vertices grid.
203
+ # Note: y is computed as H-1-i. (Apply scale to y if desired.)
204
+ top_vertices = np.zeros((H, W, 3), dtype=np.float32)
201
205
  for i in range(H):
202
206
  for j in range(W):
203
- # Original coordinates: x = j*scale, y = (H - 1 - i), z = height + background
204
- vertices[i, j, 0] = j * scale
205
- vertices[i, j, 1] = H - 1 - i # (Consider applying scale if needed)
206
- vertices[i, j, 2] = height_map[i, j] + background_height
207
+ top_vertices[i, j, 0] = j * scale
208
+ top_vertices[i, j, 1] = H - 1 - i # or (H - 1 - i) * scale if you want uniform scaling
209
+ top_vertices[i, j, 2] = height_map[i, j] + background_height
210
+
211
+ # Create the bottom vertices for the boundary by copying the x and y coordinates,
212
+ # but setting z to 0. Using these ensures that adjacent triangles share the same vertices.
213
+ bottom_vertices = top_vertices.copy()
214
+ bottom_vertices[:, :, 2] = 0
207
215
 
208
216
  triangles = []
209
217
 
210
218
  def add_triangle(v1, v2, v3):
211
- """
212
- Add a triangle to the list of triangles.
213
-
214
- Args:
215
- v1 (np.ndarray): First vertex of the triangle.
216
- v2 (np.ndarray): Second vertex of the triangle.
217
- v3 (np.ndarray): Third vertex of the triangle.
218
- """
219
+ """Add a triangle defined by vertices v1, v2, and v3."""
219
220
  triangles.append((v1, v2, v3))
220
221
 
222
+ # --- Top Surface ---
221
223
  for i in range(H - 1):
222
224
  for j in range(W - 1):
223
- v0 = vertices[i, j]
224
- v1 = vertices[i, j + 1]
225
- v2 = vertices[i + 1, j + 1]
226
- v3 = vertices[i + 1, j]
227
- # Reversed order so normals face upward
225
+ v0 = top_vertices[i, j]
226
+ v1 = top_vertices[i, j + 1]
227
+ v2 = top_vertices[i + 1, j + 1]
228
+ v3 = top_vertices[i + 1, j]
229
+ # Using reversed order so that normals face upward.
228
230
  add_triangle(v2, v1, v0)
229
231
  add_triangle(v3, v2, v0)
230
232
 
233
+ # --- Side Walls ---
234
+ # Top edge (i = 0)
231
235
  for j in range(W - 1):
232
- v0 = vertices[0, j]
233
- v1 = vertices[0, j + 1]
234
- v0b = np.array([v0[0], v0[1], 0], dtype=np.float32)
235
- v1b = np.array([v1[0], v1[1], 0], dtype=np.float32)
236
- add_triangle(v0, v1, v1b)
237
- add_triangle(v0, v1b, v0b)
236
+ vt0 = top_vertices[0, j]
237
+ vt1 = top_vertices[0, j + 1]
238
+ vb0 = bottom_vertices[0, j]
239
+ vb1 = bottom_vertices[0, j + 1]
240
+ add_triangle(vt0, vt1, vb1)
241
+ add_triangle(vt0, vb1, vb0)
242
+
243
+ # Bottom edge (i = H - 1)
238
244
  for j in range(W - 1):
239
- v0 = vertices[H - 1, j]
240
- v1 = vertices[H - 1, j + 1]
241
- v0b = np.array([v0[0], v0[1], 0], dtype=np.float32)
242
- v1b = np.array([v1[0], v1[1], 0], dtype=np.float32)
243
- add_triangle(v1, v0, v1b)
244
- add_triangle(v0, v0b, v1b)
245
+ vt0 = top_vertices[H - 1, j]
246
+ vt1 = top_vertices[H - 1, j + 1]
247
+ vb0 = bottom_vertices[H - 1, j]
248
+ vb1 = bottom_vertices[H - 1, j + 1]
249
+ add_triangle(vt1, vt0, vb0)
250
+ add_triangle(vt1, vb0, vb1)
251
+
252
+ # Left edge (j = 0)
245
253
  for i in range(H - 1):
246
- v0 = vertices[i, 0]
247
- v1 = vertices[i + 1, 0]
248
- v0b = np.array([v0[0], v0[1], 0], dtype=np.float32)
249
- v1b = np.array([v1[0], v1[1], 0], dtype=np.float32)
250
- add_triangle(v1, v0, v1b)
251
- add_triangle(v0, v0b, v1b)
254
+ vt0 = top_vertices[i, 0]
255
+ vt1 = top_vertices[i + 1, 0]
256
+ vb0 = bottom_vertices[i, 0]
257
+ vb1 = bottom_vertices[i + 1, 0]
258
+ add_triangle(vt1, vt0, vb0)
259
+ add_triangle(vt1, vb0, vb1)
260
+
261
+ # Right edge (j = W - 1)
252
262
  for i in range(H - 1):
253
- v0 = vertices[i, W - 1]
254
- v1 = vertices[i + 1, W - 1]
255
- v0b = np.array([v0[0], v0[1], 0], dtype=np.float32)
256
- v1b = np.array([v1[0], v1[1], 0], dtype=np.float32)
257
- add_triangle(v0, v1, v1b)
258
- add_triangle(v0, v1b, v0b)
259
-
260
- v0 = np.array([0, 0, 0], dtype=np.float32)
261
- v1 = np.array([(W - 1) * scale, 0, 0], dtype=np.float32)
262
- v2 = np.array([(W - 1) * scale, (H - 1) * scale, 0], dtype=np.float32)
263
- v3 = np.array([0, (H - 1) * scale, 0], dtype=np.float32)
264
- add_triangle(v2, v1, v0)
265
- add_triangle(v3, v2, v0)
263
+ vt0 = top_vertices[i, W - 1]
264
+ vt1 = top_vertices[i + 1, W - 1]
265
+ vb0 = bottom_vertices[i, W - 1]
266
+ vb1 = bottom_vertices[i + 1, W - 1]
267
+ add_triangle(vt0, vt1, vb1)
268
+ add_triangle(vt0, vb1, vb0)
269
+
270
+ # --- Bottom Face ---
271
+ # Use the four corner vertices from the bottom_vertices grid.
272
+ b_top_left = bottom_vertices[0, 0]
273
+ b_top_right = bottom_vertices[0, W - 1]
274
+ b_bottom_right = bottom_vertices[H - 1, W - 1]
275
+ b_bottom_left = bottom_vertices[H - 1, 0]
276
+ # Triangulate the bottom face (ensure the normal faces downward).
277
+ add_triangle(b_bottom_right, b_top_right, b_top_left)
278
+ add_triangle(b_bottom_left, b_bottom_right, b_top_left)
266
279
 
267
280
  num_triangles = len(triangles)
268
281
 
269
- # Write the binary STL file.
282
+ # --- Write Binary STL ---
270
283
  with open(filename, "wb") as f:
271
284
  header_str = "Binary STL generated from heightmap"
272
285
  header = header_str.encode("utf-8")
@@ -15,6 +15,7 @@ def loss_fn(
15
15
  material_TDs: torch.Tensor,
16
16
  background: torch.Tensor,
17
17
  perception_loss_module: torch.nn.Module,
18
+ add_penalty_loss: bool = True,
18
19
  ) -> torch.Tensor:
19
20
  """
20
21
  Full forward pass for continuous assignment:
@@ -39,6 +40,7 @@ def loss_fn(
39
40
  perception_loss_module=perception_loss_module,
40
41
  tau_global=tau_global,
41
42
  num_materials=material_colors.shape[0],
43
+ add_penalty_loss=add_penalty_loss,
42
44
  )
43
45
 
44
46