autoforge 1.3.0__tar.gz → 1.3.1__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 (28) hide show
  1. {autoforge-1.3.0 → autoforge-1.3.1}/PKG-INFO +5 -3
  2. {autoforge-1.3.0 → autoforge-1.3.1}/README.md +4 -2
  3. {autoforge-1.3.0 → autoforge-1.3.1}/pyproject.toml +1 -1
  4. {autoforge-1.3.0 → autoforge-1.3.1}/src/AutoForge.egg-info/PKG-INFO +5 -3
  5. {autoforge-1.3.0 → autoforge-1.3.1}/src/autoforge/Helper/OptimizerHelper.py +103 -36
  6. {autoforge-1.3.0 → autoforge-1.3.1}/src/autoforge/Helper/OutputHelper.py +65 -52
  7. {autoforge-1.3.0 → autoforge-1.3.1}/src/autoforge/Loss/LossFunctions.py +2 -0
  8. {autoforge-1.3.0 → autoforge-1.3.1}/src/autoforge/Modules/Optimizer.py +3 -2
  9. {autoforge-1.3.0 → autoforge-1.3.1}/src/autoforge/auto_forge.py +18 -17
  10. {autoforge-1.3.0 → autoforge-1.3.1}/setup.cfg +0 -0
  11. {autoforge-1.3.0 → autoforge-1.3.1}/src/AutoForge.egg-info/SOURCES.txt +0 -0
  12. {autoforge-1.3.0 → autoforge-1.3.1}/src/AutoForge.egg-info/dependency_links.txt +0 -0
  13. {autoforge-1.3.0 → autoforge-1.3.1}/src/AutoForge.egg-info/entry_points.txt +0 -0
  14. {autoforge-1.3.0 → autoforge-1.3.1}/src/AutoForge.egg-info/requires.txt +0 -0
  15. {autoforge-1.3.0 → autoforge-1.3.1}/src/AutoForge.egg-info/top_level.txt +0 -0
  16. {autoforge-1.3.0 → autoforge-1.3.1}/src/autoforge/Helper/FilamentHelper.py +0 -0
  17. {autoforge-1.3.0 → autoforge-1.3.1}/src/autoforge/Helper/ImageHelper.py +0 -0
  18. {autoforge-1.3.0 → autoforge-1.3.1}/src/autoforge/Helper/PruningHelper.py +0 -0
  19. {autoforge-1.3.0 → autoforge-1.3.1}/src/autoforge/Helper/__init__.py +0 -0
  20. {autoforge-1.3.0 → autoforge-1.3.1}/src/autoforge/Loss/PerceptionLoss.py +0 -0
  21. {autoforge-1.3.0 → autoforge-1.3.1}/src/autoforge/Loss/__init__.py +0 -0
  22. {autoforge-1.3.0 → autoforge-1.3.1}/src/autoforge/Modules/__init__.py +0 -0
  23. {autoforge-1.3.0 → autoforge-1.3.1}/src/autoforge/__init__.py +0 -0
  24. {autoforge-1.3.0 → autoforge-1.3.1}/src/autoforge/__main__.py +0 -0
  25. {autoforge-1.3.0 → autoforge-1.3.1}/src/autoforge/auto_forge_gradio.py +0 -0
  26. {autoforge-1.3.0 → autoforge-1.3.1}/src/autoforge/color_fit_bayesian.py +0 -0
  27. {autoforge-1.3.0 → autoforge-1.3.1}/src/autoforge/filamentcolors_library.py +0 -0
  28. {autoforge-1.3.0 → autoforge-1.3.1}/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.1
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). \
@@ -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). \
@@ -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.1"
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.1
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). \
@@ -1,4 +1,4 @@
1
- from itertools import permutations
1
+ import random
2
2
 
3
3
  import numpy as np
4
4
  import torch
@@ -7,7 +7,6 @@ from skimage.color import rgb2lab
7
7
  from sklearn.cluster import KMeans
8
8
  from sklearn.exceptions import ConvergenceWarning
9
9
  from sklearn.utils._testing import ignore_warnings
10
- from tqdm import tqdm
11
10
  from transformers import pipeline
12
11
  from PIL import Image
13
12
  from sklearn.metrics import silhouette_score
@@ -487,6 +486,104 @@ def init_height_map_depth_color_adjusted(
487
486
  return pixel_height_logits
488
487
 
489
488
 
489
+ def tsp_simulated_annealing(
490
+ band_reps,
491
+ start_band,
492
+ end_band,
493
+ initial_order=None,
494
+ initial_temp=1000,
495
+ cooling_rate=0.995,
496
+ num_iter=10000,
497
+ ):
498
+ """
499
+ Solve the band ordering problem using simulated annealing.
500
+
501
+ Args:
502
+ band_reps (list or np.ndarray): List of Lab color representations.
503
+ start_band (int): Index for the darkest band.
504
+ end_band (int): Index for the brightest band.
505
+ initial_order (list, optional): Initial ordering of band indices.
506
+ initial_temp (float): Starting temperature.
507
+ cooling_rate (float): Factor to cool the temperature.
508
+ num_iter (int): Maximum number of iterations.
509
+
510
+ Returns:
511
+ list: An ordering of band indices from start_band to end_band.
512
+ """
513
+ if initial_order is None:
514
+ # Use a simple ordering: start, middle bands as given, then end.
515
+ middle_indices = [
516
+ i for i in range(len(band_reps)) if i not in (start_band, end_band)
517
+ ]
518
+ order = [start_band] + middle_indices + [end_band]
519
+ else:
520
+ order = initial_order.copy()
521
+
522
+ def total_distance(order):
523
+ return sum(
524
+ np.linalg.norm(band_reps[order[i]] - band_reps[order[i + 1]])
525
+ for i in range(len(order) - 1)
526
+ )
527
+
528
+ current_distance = total_distance(order)
529
+ best_order = order.copy()
530
+ best_distance = current_distance
531
+ temp = initial_temp
532
+
533
+ for _ in range(num_iter):
534
+ # Randomly swap two indices in the middle of the order
535
+ new_order = order.copy()
536
+ idx1, idx2 = random.sample(range(1, len(order) - 1), 2)
537
+ new_order[idx1], new_order[idx2] = new_order[idx2], new_order[idx1]
538
+
539
+ new_distance = total_distance(new_order)
540
+ delta = new_distance - current_distance
541
+
542
+ # Accept the new order if it improves or with a probability to escape local minima
543
+ if delta < 0 or np.exp(-delta / temp) > random.random():
544
+ order = new_order.copy()
545
+ current_distance = new_distance
546
+ if current_distance < best_distance:
547
+ best_order = order.copy()
548
+ best_distance = current_distance
549
+
550
+ temp *= cooling_rate
551
+ if temp < 1e-6:
552
+ break
553
+ return best_order
554
+
555
+
556
+ def choose_optimal_num_bands(centroids, min_bands=2, max_bands=15, random_seed=None):
557
+ """
558
+ Determine the optimal number of clusters (bands) for the centroids
559
+ by maximizing the silhouette score.
560
+
561
+ Args:
562
+ centroids (np.ndarray): Array of centroid colors (e.g., shape (n_clusters, 3)).
563
+ min_bands (int): Minimum number of clusters to try.
564
+ max_bands (int): Maximum number of clusters to try.
565
+ random_seed (int, optional): Random seed for reproducibility.
566
+
567
+ Returns:
568
+ int: Optimal number of bands.
569
+ """
570
+ best_num = min_bands
571
+ best_score = -1
572
+
573
+ for num in range(min_bands, max_bands + 1):
574
+ kmeans = KMeans(n_clusters=num, random_state=random_seed).fit(centroids)
575
+ labels = kmeans.labels_
576
+ # If there's only one unique label, skip to avoid errors.
577
+ if len(np.unique(labels)) < 2:
578
+ continue
579
+ score = silhouette_score(centroids, labels)
580
+ if score > best_score:
581
+ best_score = score
582
+ best_num = num
583
+
584
+ return best_num
585
+
586
+
490
587
  def init_height_map(target, max_layers, h, eps=1e-6, random_seed=None):
491
588
  """
492
589
  Initialize pixel height logits based on the luminance of the target image.
@@ -512,7 +609,9 @@ def init_height_map(target, max_layers, h, eps=1e-6, random_seed=None):
512
609
  return 0.299 * col[0] + 0.587 * col[1] + 0.114 * col[2]
513
610
 
514
611
  # --- Step 2: Second clustering of centroids into bands ---
515
- num_bands = 9
612
+ num_bands = choose_optimal_num_bands(
613
+ centroids, min_bands=5, max_bands=max_layers // 2, random_seed=random_seed
614
+ )
516
615
  band_kmeans = KMeans(n_clusters=num_bands, random_state=random_seed).fit(centroids)
517
616
  band_labels = band_kmeans.labels_
518
617
 
@@ -544,39 +643,7 @@ def init_height_map(target, max_layers, h, eps=1e-6, random_seed=None):
544
643
  L_values = [lab[0] for lab in band_reps]
545
644
  start_band = np.argmin(L_values) # darkest band index
546
645
  end_band = np.argmax(L_values) # brightest band index
547
-
548
- # --- Step 5: Find the best ordering for the middle bands ---
549
- # We want to order the bands so that the total perceptual difference (Euclidean distance in Lab)
550
- # between consecutive bands is minimized, while forcing the darkest band first and brightest band last.
551
- all_indices = list(range(len(bands)))
552
- middle_indices = [i for i in all_indices if i not in (start_band, end_band)]
553
-
554
- min_total_distance = np.inf
555
- best_order = None
556
- total = len(middle_indices) * len(middle_indices)
557
- # Try all permutations of the middle bands
558
- ie = 0
559
- tbar = tqdm(
560
- permutations(middle_indices),
561
- total=total,
562
- desc="Finding best ordering for color bands:",
563
- )
564
- for perm in tbar:
565
- candidate = [start_band] + list(perm) + [end_band]
566
- total_distance = 0
567
- for i in range(len(candidate) - 1):
568
- total_distance += np.linalg.norm(
569
- band_reps[candidate[i]] - band_reps[candidate[i + 1]]
570
- )
571
- if total_distance < min_total_distance:
572
- min_total_distance = total_distance
573
- best_order = candidate
574
- tbar.set_description(
575
- f"Finding best ordering for color bands: Total distance = {min_total_distance:.2f}"
576
- )
577
- ie += 1
578
- if ie > 500000:
579
- break
646
+ best_order = tsp_simulated_annealing(band_reps, start_band, end_band)
580
647
 
581
648
  new_order = []
582
649
  for band_idx in best_order:
@@ -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
 
@@ -105,7 +105,7 @@ class FilamentOptimizer:
105
105
  "Please dont panic if the best discretize composite is blank for a while."
106
106
  )
107
107
  print(
108
- "We only record the best discrete function at the end to allow for penalty loss to function correctly."
108
+ "We only record the best discrete function after some time to allow for the penalty loss to function correctly."
109
109
  )
110
110
  plt.ion()
111
111
  self.fig, self.ax = plt.subplots(1, 4, figsize=(14, 6))
@@ -173,6 +173,7 @@ class FilamentOptimizer:
173
173
  material_TDs=self.material_TDs,
174
174
  background=self.background,
175
175
  perception_loss_module=self.perception_loss_module,
176
+ add_penalty_loss=False,
176
177
  )
177
178
  loss.backward()
178
179
  self.optimizer.step()
@@ -370,7 +371,7 @@ class FilamentOptimizer:
370
371
  and update the best solution if it improves.
371
372
  """
372
373
 
373
- for i in range(3):
374
+ for i in range(1):
374
375
  # draw random integer seed
375
376
  seed = np.random.randint(0, 1000000)
376
377
 
@@ -44,12 +44,6 @@ def main():
44
44
  parser.add_argument(
45
45
  "--iterations", type=int, default=5000, help="Number of optimization iterations"
46
46
  )
47
- parser.add_argument(
48
- "--best_loss_iterations",
49
- type=float,
50
- default=0.9,
51
- help="Percentage of optimization iterations after which we start to record the best loss",
52
- )
53
47
 
54
48
  parser.add_argument(
55
49
  "--learning_rate",
@@ -250,9 +244,7 @@ def main():
250
244
  print("Starting optimization...")
251
245
  tbar = tqdm(range(args.iterations))
252
246
  for i in tbar:
253
- loss_val = optimizer.step(
254
- record_best=i > args.best_loss_iterations * args.iterations and i % 5 == 0
255
- )
247
+ loss_val = optimizer.step(record_best=i % 3 == 0)
256
248
  optimizer.visualize(interval=25)
257
249
  if (i + 1) % 100 == 0:
258
250
  tbar.set_description(f"Iteration {i + 1}, Loss = {loss_val:.4f}")
@@ -260,6 +252,15 @@ def main():
260
252
  # After we finish, we can get the final discrete solution at solver resolution
261
253
  disc_global, disc_height_image = optimizer.get_discretized_solution(best=True)
262
254
 
255
+ # Do a quick search for a better rng seed
256
+ new_rng_seed, new_loss = optimizer.rng_seed_search(
257
+ optimizer.best_discrete_loss, 1000
258
+ )
259
+ if new_loss < optimizer.best_discrete_loss:
260
+ optimizer.best_seed = new_rng_seed
261
+ optimizer.best_discrete_loss = new_loss
262
+ print(f"New best seed found: {new_rng_seed} with loss: {new_loss}")
263
+
263
264
  # Optionally prune
264
265
  if args.perform_pruning:
265
266
  disc_global = optimizer.prune(
@@ -270,14 +271,14 @@ def main():
270
271
  tau_g=optimizer.best_tau,
271
272
  )
272
273
 
273
- # Do a quick search for a better rng seed
274
- new_rng_seed, new_loss = optimizer.rng_seed_search(
275
- optimizer.best_discrete_loss, 1000
276
- )
277
- if new_loss < optimizer.best_discrete_loss:
278
- optimizer.best_seed = new_rng_seed
279
- optimizer.best_discrete_loss = new_loss
280
- print(f"New best seed found: {new_rng_seed} with loss: {new_loss}")
274
+ # Do a quick search for a better rng seed
275
+ new_rng_seed, new_loss = optimizer.rng_seed_search(
276
+ optimizer.best_discrete_loss, 1000
277
+ )
278
+ if new_loss < optimizer.best_discrete_loss:
279
+ optimizer.best_seed = new_rng_seed
280
+ optimizer.best_discrete_loss = new_loss
281
+ print(f"New best seed found: {new_rng_seed} with loss: {new_loss}")
281
282
  else:
282
283
  print(
283
284
  f"No better seed found. Best seed loss: {optimizer.best_discrete_loss}. Best searched loss: {new_loss}"
File without changes