autoforge 1.1.2__tar.gz → 1.1.3__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.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.2
2
2
  Name: AutoForge
3
- Version: 1.1.2
3
+ Version: 1.1.3
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
@@ -34,18 +34,21 @@ All examples use only the 13 BambuLab Basic filaments, currently available in Hu
34
34
  <div style="display: flex; justify-content: center; gap: 20px;">
35
35
  <div style="text-align: center;">
36
36
  <h3>Input Image</h3>
37
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi.jpg" width="300" />
38
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature.jpg" width="300" />
37
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi.jpg" width="200" />
38
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature.jpg" width="200" />
39
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat.jpg" width="200" />
39
40
  </div>
40
41
  <div style="text-align: center;">
41
42
  <h3>Autoforge Output (75 color layers (3.0mm) + 10 gray background layers (0.4mm)) </h3>
42
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_discretized.jpg" width="300" />
43
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_discretized.jpg" width="300" />
43
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_discretized.jpg" width="200" />
44
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_discretized.jpg" width="200" />
45
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat_discretized.jpg" width="200" />
44
46
  </div>
45
47
  <div style="text-align: center;">
46
48
  <h3>Imported Hueforge View</h3>
47
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_hueforge.png" width="300" />
48
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_hueforge.png" width="300" />
49
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_hueforge.png" width="200" />
50
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_hueforge.png" width="200" />
51
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat_hueforge.png" width="200" />
49
52
  </div>
50
53
  </div>
51
54
 
@@ -66,33 +69,17 @@ I would love to see what you made!
66
69
 
67
70
  ## Installation
68
71
 
69
- 1. **Clone the Repository**
70
-
71
- ```bash
72
- git clone https://github.com/hvoss-techfak/AutoForge.git
73
- cd AutoForge
74
- ```
75
-
76
- 2. **Set Up a Virtual Environment (Optional but Recommended)**
77
-
78
- ```bash
79
- python -m venv venv
80
- source venv/bin/activate # On Windows: venv\Scripts\activate
81
- ```
82
-
83
- 3. **Install Dependencies**
84
-
85
- ```bash
86
- pip install -r requirements.txt
87
- ```
72
+ To install AutoForge, simply install the current version from PyPI:
73
+ ```bash
74
+ pip install autoforge
75
+ ```
88
76
 
89
- See the [requirements.txt](requirements.txt) file for a complete list of dependencies.
90
- The optimizer is built using JAX, which benefits from a CUDA-compatible GPU for optimal performance.
91
- Without a dedicated GPU the process can take significantly longer (up to 10-20x slower).
92
- If you have a GPU, you can install the GPU version of JAX by running:
77
+ The optimizer is built using JAX, which benefits from a CUDA-compatible GPU for optimal performance. \
78
+ Without a dedicated GPU the process can take significantly longer (up to 10-20x slower). \
79
+ If you have a GPU, you can install the GPU version of autoforge by running:\
93
80
 
94
81
  ```bash
95
- pip install -U "jax[cuda12]"
82
+ pip install autoforge[gpu]
96
83
  ```
97
84
 
98
85
  Currently, jax gpu support is only available under linux. \
@@ -108,7 +95,7 @@ The script is run from the command line and accepts several arguments. Below is
108
95
  > If you want to limit the amount of colors the program can use, you can simply delete the colors you don't want from the CSV file.
109
96
 
110
97
  ```bash
111
- python auto_forge.py --input_image path/to/input_image.jpg --csv_file path/to/materials.csv --output_folder outputs
98
+ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv --output_folder outputs
112
99
  ```
113
100
 
114
101
  ### Command Line Arguments
@@ -174,5 +161,6 @@ AutoForge makes use of several open source libraries:
174
161
  Example Images:
175
162
  <a href="https://www.vecteezy.com/free-photos/anime-girl">Anime Girl Stock photos by Vecteezy</a>
176
163
  <a href="https://www.vecteezy.com/free-photos/nature">Nature Stock photos by Vecteezy</a>
164
+ <a href="https://www.vecteezy.com/free-photos/ai-generated">Ai Generated Stock photos by Vecteezy</a>
177
165
 
178
166
  Happy printing!
@@ -12,18 +12,21 @@ All examples use only the 13 BambuLab Basic filaments, currently available in Hu
12
12
  <div style="display: flex; justify-content: center; gap: 20px;">
13
13
  <div style="text-align: center;">
14
14
  <h3>Input Image</h3>
15
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi.jpg" width="300" />
16
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature.jpg" width="300" />
15
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi.jpg" width="200" />
16
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature.jpg" width="200" />
17
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat.jpg" width="200" />
17
18
  </div>
18
19
  <div style="text-align: center;">
19
20
  <h3>Autoforge Output (75 color layers (3.0mm) + 10 gray background layers (0.4mm)) </h3>
20
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_discretized.jpg" width="300" />
21
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_discretized.jpg" width="300" />
21
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_discretized.jpg" width="200" />
22
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_discretized.jpg" width="200" />
23
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat_discretized.jpg" width="200" />
22
24
  </div>
23
25
  <div style="text-align: center;">
24
26
  <h3>Imported Hueforge View</h3>
25
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_hueforge.png" width="300" />
26
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_hueforge.png" width="300" />
27
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_hueforge.png" width="200" />
28
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_hueforge.png" width="200" />
29
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat_hueforge.png" width="200" />
27
30
  </div>
28
31
  </div>
29
32
 
@@ -44,33 +47,17 @@ I would love to see what you made!
44
47
 
45
48
  ## Installation
46
49
 
47
- 1. **Clone the Repository**
48
-
49
- ```bash
50
- git clone https://github.com/hvoss-techfak/AutoForge.git
51
- cd AutoForge
52
- ```
53
-
54
- 2. **Set Up a Virtual Environment (Optional but Recommended)**
55
-
56
- ```bash
57
- python -m venv venv
58
- source venv/bin/activate # On Windows: venv\Scripts\activate
59
- ```
60
-
61
- 3. **Install Dependencies**
62
-
63
- ```bash
64
- pip install -r requirements.txt
65
- ```
50
+ To install AutoForge, simply install the current version from PyPI:
51
+ ```bash
52
+ pip install autoforge
53
+ ```
66
54
 
67
- See the [requirements.txt](requirements.txt) file for a complete list of dependencies.
68
- The optimizer is built using JAX, which benefits from a CUDA-compatible GPU for optimal performance.
69
- Without a dedicated GPU the process can take significantly longer (up to 10-20x slower).
70
- If you have a GPU, you can install the GPU version of JAX by running:
55
+ The optimizer is built using JAX, which benefits from a CUDA-compatible GPU for optimal performance. \
56
+ Without a dedicated GPU the process can take significantly longer (up to 10-20x slower). \
57
+ If you have a GPU, you can install the GPU version of autoforge by running:\
71
58
 
72
59
  ```bash
73
- pip install -U "jax[cuda12]"
60
+ pip install autoforge[gpu]
74
61
  ```
75
62
 
76
63
  Currently, jax gpu support is only available under linux. \
@@ -86,7 +73,7 @@ The script is run from the command line and accepts several arguments. Below is
86
73
  > If you want to limit the amount of colors the program can use, you can simply delete the colors you don't want from the CSV file.
87
74
 
88
75
  ```bash
89
- python auto_forge.py --input_image path/to/input_image.jpg --csv_file path/to/materials.csv --output_folder outputs
76
+ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv --output_folder outputs
90
77
  ```
91
78
 
92
79
  ### Command Line Arguments
@@ -152,5 +139,6 @@ AutoForge makes use of several open source libraries:
152
139
  Example Images:
153
140
  <a href="https://www.vecteezy.com/free-photos/anime-girl">Anime Girl Stock photos by Vecteezy</a>
154
141
  <a href="https://www.vecteezy.com/free-photos/nature">Nature Stock photos by Vecteezy</a>
142
+ <a href="https://www.vecteezy.com/free-photos/ai-generated">Ai Generated Stock photos by Vecteezy</a>
155
143
 
156
144
  Happy printing!
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "AutoForge"
7
- version="1.1.2"
7
+ version="1.1.3"
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.1.2
3
+ Version: 1.1.3
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
@@ -34,18 +34,21 @@ All examples use only the 13 BambuLab Basic filaments, currently available in Hu
34
34
  <div style="display: flex; justify-content: center; gap: 20px;">
35
35
  <div style="text-align: center;">
36
36
  <h3>Input Image</h3>
37
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi.jpg" width="300" />
38
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature.jpg" width="300" />
37
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi.jpg" width="200" />
38
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature.jpg" width="200" />
39
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat.jpg" width="200" />
39
40
  </div>
40
41
  <div style="text-align: center;">
41
42
  <h3>Autoforge Output (75 color layers (3.0mm) + 10 gray background layers (0.4mm)) </h3>
42
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_discretized.jpg" width="300" />
43
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_discretized.jpg" width="300" />
43
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_discretized.jpg" width="200" />
44
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_discretized.jpg" width="200" />
45
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat_discretized.jpg" width="200" />
44
46
  </div>
45
47
  <div style="text-align: center;">
46
48
  <h3>Imported Hueforge View</h3>
47
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_hueforge.png" width="300" />
48
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_hueforge.png" width="300" />
49
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_hueforge.png" width="200" />
50
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_hueforge.png" width="200" />
51
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat_hueforge.png" width="200" />
49
52
  </div>
50
53
  </div>
51
54
 
@@ -66,33 +69,17 @@ I would love to see what you made!
66
69
 
67
70
  ## Installation
68
71
 
69
- 1. **Clone the Repository**
70
-
71
- ```bash
72
- git clone https://github.com/hvoss-techfak/AutoForge.git
73
- cd AutoForge
74
- ```
75
-
76
- 2. **Set Up a Virtual Environment (Optional but Recommended)**
77
-
78
- ```bash
79
- python -m venv venv
80
- source venv/bin/activate # On Windows: venv\Scripts\activate
81
- ```
82
-
83
- 3. **Install Dependencies**
84
-
85
- ```bash
86
- pip install -r requirements.txt
87
- ```
72
+ To install AutoForge, simply install the current version from PyPI:
73
+ ```bash
74
+ pip install autoforge
75
+ ```
88
76
 
89
- See the [requirements.txt](requirements.txt) file for a complete list of dependencies.
90
- The optimizer is built using JAX, which benefits from a CUDA-compatible GPU for optimal performance.
91
- Without a dedicated GPU the process can take significantly longer (up to 10-20x slower).
92
- If you have a GPU, you can install the GPU version of JAX by running:
77
+ The optimizer is built using JAX, which benefits from a CUDA-compatible GPU for optimal performance. \
78
+ Without a dedicated GPU the process can take significantly longer (up to 10-20x slower). \
79
+ If you have a GPU, you can install the GPU version of autoforge by running:\
93
80
 
94
81
  ```bash
95
- pip install -U "jax[cuda12]"
82
+ pip install autoforge[gpu]
96
83
  ```
97
84
 
98
85
  Currently, jax gpu support is only available under linux. \
@@ -108,7 +95,7 @@ The script is run from the command line and accepts several arguments. Below is
108
95
  > If you want to limit the amount of colors the program can use, you can simply delete the colors you don't want from the CSV file.
109
96
 
110
97
  ```bash
111
- python auto_forge.py --input_image path/to/input_image.jpg --csv_file path/to/materials.csv --output_folder outputs
98
+ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv --output_folder outputs
112
99
  ```
113
100
 
114
101
  ### Command Line Arguments
@@ -174,5 +161,6 @@ AutoForge makes use of several open source libraries:
174
161
  Example Images:
175
162
  <a href="https://www.vecteezy.com/free-photos/anime-girl">Anime Girl Stock photos by Vecteezy</a>
176
163
  <a href="https://www.vecteezy.com/free-photos/nature">Nature Stock photos by Vecteezy</a>
164
+ <a href="https://www.vecteezy.com/free-photos/ai-generated">Ai Generated Stock photos by Vecteezy</a>
177
165
 
178
166
  Happy printing!
@@ -9,5 +9,7 @@ src/AutoForge.egg-info/top_level.txt
9
9
  src/autoforge/__init__.py
10
10
  src/autoforge/__main__.py
11
11
  src/autoforge/auto_forge.py
12
+ src/autoforge/auto_forge_gradio.py
13
+ src/autoforge/color_fit_bayesian.py
12
14
  src/autoforge/filamentcolors_library.py
13
15
  src/autoforge/helper_functions.py
@@ -67,9 +67,9 @@ def composite_pixel_combined(pixel_height_logit, global_logits, tau_height, tau_
67
67
  if mode == "pruning":
68
68
  global_logits = jax.nn.one_hot(global_logits, len(material_colors))
69
69
 
70
- A = 0.178763
71
- k = 39.302848
72
- b = 0.351177
70
+ A = 0.1215
71
+ k = 61.6970
72
+ b = 0.4773
73
73
 
74
74
  def step_fn(carry, i):
75
75
  comp, remaining = carry
@@ -274,7 +274,7 @@ def discretize_solution_jax(params, tau_global, gumbel_keys, h, max_layers):
274
274
 
275
275
 
276
276
 
277
- def run_optimizer(rng_key, target, H, W, max_layers, h, material_colors, material_TDs, background,
277
+ def run_optimizer(rng_key, target,pixel_height_logits, H, W, max_layers, h, material_colors, material_TDs, background,
278
278
  num_iters, learning_rate, decay_v, loss_function, visualize=False,
279
279
  output_folder=None, save_interval_pct=None,
280
280
  img_width=None, img_height=None, background_height=None,
@@ -313,7 +313,7 @@ def run_optimizer(rng_key, target, H, W, max_layers, h, material_colors, materia
313
313
  rng_key, subkey = random.split(rng_key)
314
314
 
315
315
  #pixel_height_logits = initialize_pixel_height_logits(target)
316
- pixel_height_logits = init_height_map(target,max_layers,h)
316
+
317
317
  #pixel_height_logits = random.uniform(subkey, (H, W), minval=-2.0, maxval=2.0)
318
318
 
319
319
  params = {'global_logits': global_logits, 'pixel_height_logits': pixel_height_logits}
@@ -494,7 +494,7 @@ def save_intermediate_outputs(iteration, params, tau_global, gumbel_keys, h, max
494
494
  np.array(disc_height_image),
495
495
  img_width, img_height, stl_filename, csv_file)
496
496
 
497
- def pruning(target,best_params,disc_global,tau_global_disc,gumbel_keys_disc,h,max_layers,material_colors,material_TDs,background,max_loss_increase=0.05):
497
+ def pruning(target,best_params,disc_global,tau_global_disc,gumbel_keys_disc,h,max_layers,material_colors,material_TDs,background,max_loss_increase=0.1):
498
498
  initial_image = composite_image_combined_jit(best_params['pixel_height_logits'], disc_global,
499
499
  tau_global_disc, tau_global_disc, gumbel_keys_disc,
500
500
  h, max_layers, material_colors, material_TDs, background, mode="pruning")
@@ -524,40 +524,69 @@ def pruning(target,best_params,disc_global,tau_global_disc,gumbel_keys_disc,h,ma
524
524
  # Pruning test 1 - Try removing entire color bands
525
525
  current_bands = len(color_bands)
526
526
 
527
- for _ in tqdm(range(100),desc="Pruning - Removing entire color bands"):
527
+ def get_image_loss(disc_global):
528
+ disc_global = jnp.array(disc_global)
529
+ return jnp.mean((composite_image_combined_jit(best_params['pixel_height_logits'], disc_global,
530
+ tau_global_disc, tau_global_disc, gumbel_keys_disc,
531
+ h, max_layers, material_colors, material_TDs, background, mode="pruning") - target) ** 2)
532
+
533
+ for _ in tqdm(range(10),desc="Pruning - Removing entire color bands"):
534
+ change = False
535
+ #testing removing color bands forward
528
536
  for i in range(1,len(color_bands)):
529
537
  band = color_bands[i]
530
538
  disc_global_copy = np.asarray(disc_global).copy()
531
539
  disc_global_copy[band[0]:band[1]+1] = color_bands[i-1][2]
532
- disc_global_copy = jnp.array(disc_global_copy)
533
- new_image = composite_image_combined_jit(best_params['pixel_height_logits'], disc_global_copy,
534
- tau_global_disc, tau_global_disc, gumbel_keys_disc,
535
- h, max_layers, material_colors, material_TDs, background, mode="pruning")
536
- new_loss = jnp.mean((new_image - target) ** 2)
540
+ new_loss = get_image_loss(disc_global)
537
541
  if new_loss < max_loss:
542
+ change = True
538
543
  disc_global = disc_global_copy
539
- max_loss = new_loss
540
544
  color_bands[i] = color_bands[i-1]
541
545
 
542
- # Same as above but in reverse
546
+ # testing removing color bands backwards
543
547
  for i in reversed(range(0,len(color_bands)-1)):
544
548
  band = color_bands[i]
545
549
  disc_global_copy = np.asarray(disc_global).copy()
546
550
  disc_global_copy[band[0]:band[1]+1] = color_bands[i+1][2]
547
- disc_global_copy = jnp.array(disc_global_copy)
548
- new_image = composite_image_combined_jit(best_params['pixel_height_logits'], disc_global_copy,
549
- tau_global_disc, tau_global_disc, gumbel_keys_disc,
550
- h, max_layers, material_colors, material_TDs, background, mode="pruning")
551
- new_loss = jnp.mean((new_image - target) ** 2)
551
+ new_loss = get_image_loss(disc_global)
552
552
  if new_loss < max_loss:
553
+ change = True
553
554
  disc_global = disc_global_copy
554
- max_loss = new_loss
555
555
  color_bands[i] = color_bands[i+1]
556
- if current_bands == get_color_bands(disc_global):
556
+
557
+ #we only check for these one if the loss decreases as otherwise we risk color band shifting
558
+ p2_initial_loss = get_image_loss(disc_global)
559
+ #testing border shifting forward
560
+ for i in range(0,len(color_bands)):
561
+ band = color_bands[i]
562
+ disc_global_copy = np.asarray(disc_global).copy()
563
+ disc_global_copy[band[0]:min(band[1]+2,max_layers)] = band[2]
564
+ new_loss = get_image_loss(disc_global)
565
+ if new_loss < p2_initial_loss:
566
+ #print(f"color band {i} shifted forwards - {band} - {new_loss}")
567
+ p2_initial_loss = new_loss
568
+ change = True
569
+ disc_global = disc_global_copy
570
+ color_bands[i] = (band[0],band[1]+1,band[2])
571
+
572
+ #testing border shifting backwards
573
+ for i in reversed(range(0,len(color_bands))):
574
+ band = color_bands[i]
575
+ disc_global_copy = np.asarray(disc_global).copy()
576
+ disc_global_copy[max(0,band[0]-1):band[1]+1] = band[2]
577
+ new_loss = get_image_loss(disc_global)
578
+ if new_loss < p2_initial_loss:
579
+ #print(f"color band {i} shifted backwards - {band} - {new_loss}")
580
+ p2_initial_loss = new_loss
581
+ change = True
582
+ disc_global = disc_global_copy
583
+ color_bands[i] = (band[0]-1,band[1],band[2])
584
+
585
+
586
+
587
+ if not change:
557
588
  print("No changes in color bands. Exiting pruning.")
558
589
  break
559
- else:
560
- current_bands = get_color_bands(disc_global)
561
590
  #calculate color band difference
562
591
  new_color_bands = get_color_bands(disc_global)
563
592
  print(f"Color bands before pruning: {len(color_bands)}")
@@ -607,6 +636,7 @@ def main():
607
636
  parser.add_argument("--visualize", action="store_true", help="Enable visualization during optimization")
608
637
  parser.add_argument("--perform_gumbal_search", type=bool, default=True, help="Perform gumbal search after optimization")
609
638
  parser.add_argument("--perform_pruning", type=bool, default=True, help="Perform pruning after optimization")
639
+ parser.add_argument("--pruning_max_loss_increase", type=float, default=0.1, help="Maximum loss increase for pruning")
610
640
  parser.add_argument("--save_interval_pct", type=float, default=20,help="Percentage interval to save intermediate results")
611
641
 
612
642
  args = parser.parse_args()
@@ -638,12 +668,22 @@ def main():
638
668
  new_h, new_w, _ = target.shape
639
669
 
640
670
  output_target = resize_image(img,args.output_size)
671
+ print(target.shape)
672
+
673
+ pixel_height_logits = init_height_map(output_target, args.max_layers, h_value)
674
+ pixel_height_logits = np.asarray(pixel_height_logits)
675
+ print(pixel_height_logits.shape)
676
+
677
+ #resize to target size
678
+ o_shape = pixel_height_logits.shape
679
+ pixel_height_logits_solver = resize_image(pixel_height_logits.reshape(o_shape[0],o_shape[1],1),args.solver_size)
641
680
 
642
681
  target = jnp.array(target, dtype=jnp.float64)
643
682
 
683
+
644
684
  rng_key = random.PRNGKey(int(time.time()))
645
685
  best_params, _, val_gumbel_keys = run_optimizer(
646
- rng_key, target, new_h, new_w, max_layers_value, h_value,
686
+ rng_key, target,pixel_height_logits_solver, new_h, new_w, max_layers_value, h_value,
647
687
  material_colors, material_TDs, background,
648
688
  args.iterations, args.learning_rate, decay_v_value,
649
689
  loss_function=loss_fn,
@@ -664,7 +704,7 @@ def main():
664
704
  max_layers_value, rng_key, target, val_gumbel_keys,iterations=1000,desc="Searching small size image for best gumbal keys")
665
705
 
666
706
  if args.solver_size != args.output_size and args.perform_gumbal_search:
667
- best_params["pixel_height_logits"] = init_height_map(output_target, args.max_layers, h_value)
707
+ best_params["pixel_height_logits"] = pixel_height_logits
668
708
 
669
709
  val_gumbel_keys = gumbal_bruteforce(background, best_params, decay_v_value, h_value, material_TDs, material_colors,
670
710
  max_layers_value, rng_key, output_target, val_gumbel_keys, iterations=1000,
@@ -673,7 +713,7 @@ def main():
673
713
  disc_global, disc_height_image = discretize_solution_jax(best_params, tau_global_disc, val_gumbel_keys, h_value, max_layers_value)
674
714
 
675
715
  if args.perform_pruning:
676
- disc_global = pruning(output_target,best_params,disc_global,tau_global_disc,val_gumbel_keys,h_value,args.max_layers,material_colors,material_TDs,background)
716
+ disc_global = pruning(output_target,best_params,disc_global,tau_global_disc,val_gumbel_keys,h_value,args.max_layers,material_colors,material_TDs,background,max_loss_increase=args.pruning_max_loss_increase)
677
717
 
678
718
  disc_comp = composite_image_combined_jit(best_params['pixel_height_logits'], disc_global,
679
719
  tau_global_disc, tau_global_disc, val_gumbel_keys,
@@ -0,0 +1,495 @@
1
+ #!/usr/bin/env python
2
+ """
3
+ gradio_interface.py
4
+
5
+ This file demonstrates a Gradio interface that wraps your optimization code
6
+ and also adds a "Height Map Annotation" tab. In that tab the initial height map
7
+ is previewed (using your init_height_map function) and you can annotate it manually
8
+ with color labels loaded from your CSV file.
9
+ """
10
+ import json
11
+ import threading
12
+ import time
13
+ import queue
14
+ import os
15
+ import cv2
16
+ import numpy as np
17
+ import gradio as gr
18
+ import jax
19
+ import jax.numpy as jnp
20
+ import optax
21
+ import math
22
+ import time
23
+
24
+ import pandas as pd
25
+
26
+ # Import your functions from your original module.
27
+ # (Make sure these functions are available in your PYTHONPATH or adjust the import as needed.)
28
+ from autoforge.auto_forge import (
29
+ run_optimizer,
30
+ loss_fn,
31
+ composite_image_combined_jit,
32
+ init_height_map,
33
+ hex_to_rgb,
34
+ load_materials,
35
+ generate_stl,
36
+ generate_swap_instructions,
37
+ generate_project_file,
38
+ gumbal_bruteforce,
39
+ pruning
40
+ )
41
+ from autoforge.filamentcolors_library import download_filament_info
42
+ from autoforge.helper_functions import extract_colors_from_swatches, swatch_data_to_table
43
+
44
+ # -------------------------------
45
+ # Global state for background optimization
46
+ # -------------------------------
47
+ pause_flag = threading.Event()
48
+ stop_flag = threading.Event()
49
+ update_queue = queue.Queue()
50
+ optimizer_thread = None
51
+ current_optimizer_state = {} # Holds the latest state (iteration, loss, composite image, etc.)
52
+
53
+
54
+ # =============================================================================
55
+ # Helper functions for height map preview and annotation
56
+ # =============================================================================
57
+ def compute_initial_height_map(input_image, max_layers, layer_height):
58
+ """
59
+ Compute and return an initial height map preview from the input image.
60
+ The height map is normalized to 0-255 for display.
61
+ """
62
+ # Convert the input image to a jax array (assumed to be RGB)
63
+ target = jnp.array(input_image, dtype=jnp.float64)
64
+ height_map = init_height_map(target, max_layers, layer_height)
65
+ height_map_np = np.array(height_map)
66
+ # Normalize for display purposes.
67
+ norm = (height_map_np - height_map_np.min())
68
+ norm = norm / (norm.max() + 1e-6)
69
+ normalized = (255 * norm).astype(np.uint8)
70
+ return normalized
71
+
72
+
73
+ def get_material_labels(csv_file):
74
+ """
75
+ Load the material names from the CSV file.
76
+ """
77
+ _, _, material_names, _ = load_materials(csv_file)
78
+ return material_names
79
+
80
+
81
+ def init_height_map_preview(input_image, max_layers, layer_height, csv_file):
82
+ """
83
+ Compute the initial height map from the uploaded image and update the color labels.
84
+ Returns:
85
+ - The height map image for preview.
86
+ - An update for the dropdown choices containing the material labels.
87
+ """
88
+ if input_image is None:
89
+ return None, gr.Dropdown.update(choices=[])
90
+ hm = compute_initial_height_map(input_image, max_layers, layer_height)
91
+ labels = get_material_labels(csv_file)
92
+ default_val = labels[0] if labels else None
93
+ return hm, gr.Dropdown.update(choices=labels, value=default_val)
94
+
95
+
96
+ def save_annotated_height_map(annotated_hm):
97
+ """
98
+ Save the annotated height map image to a file and return the file path.
99
+ """
100
+ output_path = "annotated_height_map.png"
101
+ # Assume the annotated image is in RGB format.
102
+ cv2.imwrite(output_path, cv2.cvtColor(annotated_hm, cv2.COLOR_RGB2BGR))
103
+ return output_path
104
+
105
+
106
+ # =============================================================================
107
+ # Background Optimization Function (Simulated)
108
+ # =============================================================================
109
+ def run_optimization_background(
110
+ input_image, csv_file,
111
+ iterations, learning_rate,
112
+ layer_height, max_layers,
113
+ background_height, background_color,
114
+ output_size, solver_size,
115
+ decay, perform_gumbal_search, perform_pruning,
116
+ save_interval_pct
117
+ ):
118
+ """
119
+ This function loads the image and material data, then runs the optimizer.
120
+ It updates a shared queue with intermediate outputs for the UI.
121
+ (Here we simulate the iterative updates; integrate your actual optimizer loop as needed.)
122
+ """
123
+ global current_optimizer_state
124
+
125
+ pause_flag.clear()
126
+ stop_flag.clear()
127
+
128
+ # Load and prepare the image.
129
+ img = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
130
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
131
+ H, W, _ = img.shape
132
+ solver_size = int(solver_size)
133
+ target_solver = cv2.resize(img, (solver_size, solver_size))
134
+ target_solver = jnp.array(target_solver, dtype=jnp.float64)
135
+
136
+ # Load materials.
137
+ material_colors, material_TDs, material_names, material_hex = load_materials(csv_file)
138
+ background = jnp.array(hex_to_rgb(background_color), dtype=jnp.float64)
139
+
140
+ # Initialize optimizer state (simulate initialization).
141
+ rng_key = jax.random.PRNGKey(int(time.time()))
142
+
143
+ # Simulated iterative optimization loop.
144
+ for i in range(iterations):
145
+ if stop_flag.is_set():
146
+ current_optimizer_state["message"] = "Optimization stopped."
147
+ update_queue.put(current_optimizer_state)
148
+ return
149
+ # Pause if needed.
150
+ while pause_flag.is_set():
151
+ current_optimizer_state["message"] = f"Paused at iteration {i}."
152
+ update_queue.put(current_optimizer_state)
153
+ time.sleep(0.2)
154
+
155
+ # Simulate an update step.
156
+ simulated_loss = np.exp(-i / iterations) * 100 + np.random.rand() * 5
157
+ simulated_comp = np.clip(np.random.randn(solver_size, solver_size, 3) * 20 + 128, 0, 255).astype(np.uint8)
158
+ current_optimizer_state = {
159
+ "iteration": i,
160
+ "loss": simulated_loss,
161
+ "composite": simulated_comp,
162
+ "message": ""
163
+ }
164
+ update_queue.put(current_optimizer_state)
165
+ time.sleep(0.1)
166
+
167
+ current_optimizer_state["message"] = "Optimization complete."
168
+ update_queue.put(current_optimizer_state)
169
+
170
+
171
+ # =============================================================================
172
+ # Live update generator for Gradio streaming
173
+ # =============================================================================
174
+ def live_updates():
175
+ """
176
+ Generator that yields the latest composite image and status text.
177
+ """
178
+ while True:
179
+ try:
180
+ state = update_queue.get(timeout=1.0)
181
+ composite = state.get("composite", np.zeros((128, 128, 3), dtype=np.uint8))
182
+ iteration = state.get("iteration", 0)
183
+ loss = state.get("loss", 0)
184
+ message = state.get("message", "")
185
+ status_text = f"Iteration: {iteration} | Loss: {loss:.4f} {message}"
186
+ yield composite, status_text
187
+ if message in ["Optimization complete.", "Optimization stopped."]:
188
+ break
189
+ except queue.Empty:
190
+ if stop_flag.is_set():
191
+ break
192
+
193
+
194
+ # =============================================================================
195
+ # Button callback functions for optimization control.
196
+ # =============================================================================
197
+ def start_optimization(
198
+ input_image, csv_file,
199
+ iterations, learning_rate,
200
+ layer_height, max_layers,
201
+ background_height, background_color,
202
+ output_size, solver_size,
203
+ decay, perform_gumbal_search, perform_pruning,
204
+ save_interval_pct
205
+ ):
206
+ global optimizer_thread
207
+ if optimizer_thread is None or not optimizer_thread.is_alive():
208
+ optimizer_thread = threading.Thread(
209
+ target=run_optimization_background,
210
+ args=(
211
+ input_image, csv_file, iterations, learning_rate,
212
+ layer_height, max_layers, background_height, background_color,
213
+ output_size, solver_size, decay, perform_gumbal_search, perform_pruning,
214
+ save_interval_pct
215
+ ),
216
+ daemon=True
217
+ )
218
+ optimizer_thread.start()
219
+ return "Optimization started."
220
+
221
+
222
+ def pause_optimization():
223
+ pause_flag.set()
224
+ return "Optimization paused."
225
+
226
+
227
+ def resume_optimization():
228
+ pause_flag.clear()
229
+ return "Optimization resumed."
230
+
231
+
232
+ def stop_optimization():
233
+ stop_flag.set()
234
+ return "Optimization stopped."
235
+
236
+ # =============================================================================
237
+ # Filament Tab Helper
238
+ # =============================================================================
239
+ def filter_table(search_text, table):
240
+ if not search_text:
241
+ return table
242
+ filtered = [
243
+ row for row in table
244
+ if search_text.lower() in row["Brand"].lower() or search_text.lower() in row["Name"].lower()
245
+ ]
246
+ return filtered
247
+
248
+ # Function to update the active list based on selection.
249
+ def update_active(selected, table):
250
+ # 'selected' is a list of strings in the format "Brand - Name"
251
+ active = [row for row in table if f'{row["Brand"]} - {row["Name"]}' in selected]
252
+ return active
253
+
254
+ # =============================================================================
255
+ # Postprocessing callbacks (simulated)
256
+ # =============================================================================
257
+ def run_gumbal_search_callback():
258
+ updated_comp = np.clip(np.random.randn(128, 128, 3) * 20 + 150, 0, 255).astype(np.uint8)
259
+ status = "Gumbal search completed."
260
+ return updated_comp, status
261
+
262
+
263
+ def run_pruning_callback():
264
+ updated_comp = np.clip(np.random.randn(128, 128, 3) * 20 + 170, 0, 255).astype(np.uint8)
265
+ status = "Pruning completed."
266
+ return updated_comp, status
267
+
268
+
269
+ # =============================================================================
270
+ # Export callback.
271
+ # =============================================================================
272
+ def export_outputs_callback():
273
+ output_folder = "temp_output"
274
+ os.makedirs(output_folder, exist_ok=True)
275
+ stl_path = os.path.join(output_folder, "final_model.stl")
276
+ swap_path = os.path.join(output_folder, "swap_instructions.txt")
277
+ project_path = os.path.join(output_folder, "project_file.hfp")
278
+
279
+ with open(stl_path, "w") as f:
280
+ f.write("STL file content")
281
+ with open(swap_path, "w") as f:
282
+ f.write("Swap instructions content")
283
+ with open(project_path, "w") as f:
284
+ f.write("Project file content")
285
+
286
+ # For this example, we return the STL file for download.
287
+ return stl_path
288
+
289
+ #download filament info
290
+ print(f"Downloading filament info...")
291
+ print("If this is your first time running this script, it may take a few minutes.")
292
+ download_filament_info()
293
+ #load swatch json file
294
+ with open("swatches.json", "r") as f:
295
+ swatch_data = json.load(f)
296
+
297
+ table = swatch_data_to_table(swatch_data)
298
+ table_df = pd.DataFrame(table)
299
+
300
+
301
+ def filter_swatches_callback(search_query, table_data):
302
+ """
303
+ Given a search query and the full swatch table data,
304
+ returns a filtered table (including the header) where the query is
305
+ found in any field (case-insensitive).
306
+ """
307
+ if table_data is None or len(table_data) == 0:
308
+ return []
309
+ header = table_data[0]
310
+ rows = table_data[1:]
311
+ if search_query is None or search_query.strip() == "":
312
+ filtered = table_data
313
+ else:
314
+ filtered_rows = [
315
+ row for row in rows
316
+ if any(search_query.lower() in str(cell).lower() for cell in row)
317
+ ]
318
+ filtered = [header] + filtered_rows
319
+ return filtered
320
+
321
+
322
+ def add_to_active_callback(selected, table_data, active_data):
323
+ """
324
+ Given a list of selected material names (e.g., "Brand - Name"),
325
+ adds those rows from the full table_data (if not already present)
326
+ to the active list.
327
+ """
328
+ if table_data is None or len(table_data) == 0:
329
+ return active_data
330
+ header = table_data[0]
331
+ if active_data is None or len(active_data) == 0:
332
+ active_data = [header]
333
+ else:
334
+ # Ensure header is set
335
+ if active_data[0] != header:
336
+ active_data[0] = header
337
+ # Build a set of already-added materials (by "Brand - Name")
338
+ existing = {row[0] + " - " + row[1] for row in active_data[1:]}
339
+ # Loop through table_data rows and add those whose "Brand - Name" is in selected
340
+ for row in table_data[1:]:
341
+ material = row[0] + " - " + row[1]
342
+ if material in selected and material not in existing:
343
+ active_data.append(row)
344
+ return active_data
345
+
346
+ def highlight_cols(x):
347
+ df = x.copy()
348
+ df.loc[:, :] = 'color: purple'
349
+ df[['B', 'C', 'E']] = 'color: green'
350
+ return df
351
+
352
+ # ------------------------------------------------------------------------------
353
+ # Gradio Interface (including the new Color Filament Tab)
354
+ # ------------------------------------------------------------------------------
355
+
356
+ with gr.Blocks() as demo:
357
+ gr.Markdown("## Autoforge")
358
+
359
+ # -------------------------
360
+ # Color Filament Tab
361
+ # -------------------------
362
+ with gr.Tab("Color Filament"):
363
+ gr.Markdown("### Manage and Select Filament Colors")
364
+
365
+ with gr.Row():
366
+ #fill with data from table_df
367
+ material_select = gr.Dropdown(label="Select Filaments", multiselect=True,interactive=Trueasdasdasdsseee, choices=[brand + " - " + name + " - " + hex_color for brand,name,hex_color in table_df[["Brand", "Name", "Hex Color"]].values.tolist()])
368
+ add_active_btn = gr.Button("Add to Active List")
369
+ with gr.Row():
370
+ active_filaments = gr.DataFrame(label="Active Filaments", interactive=True)
371
+
372
+ with gr.Row():
373
+ #s = table_df.grsdfstyle.apply(highlight_cols, axis=None)
374
+ swatch_table = gr.DataFrame(value=table_df,label="All Filaments", interactive=True)
375
+
376
+ # 3. Add the selected materials to the active list.grgr
377
+ add_active_btn.click(
378
+ add_to_active_callback,
379
+ inputs=[material_select, swatch_table, active_filaments],
380
+ outputs=active_filaments
381
+ )
382
+
383
+ with gr.Tab("Height Map Annotation"):
384
+ gr.Markdown("### Preview and Annotate Height Map")
385
+ with gr.Row():
386
+ hm_input_image = gr.Image(label="Upload Image",
387
+ sources=['upload'],
388
+ type='numpy'
389
+ )
390
+ hm_csv_file = gr.Textbox(label="CSV File Path", value="materials.csv")
391
+ with gr.Row():
392
+ hm_max_layers = gr.Slider(10, 150, label="Max Layers", value=75, step=1)
393
+ hm_layer_height = gr.Slider(0.01, 0.1, label="Layer Height (mm)", value=0.04, step=0.01)
394
+ load_hm_btn = gr.Button("Load Height Map Preview")
395
+ with gr.Row():
396
+ height_map_preview = gr.ImageEditor(
397
+ label="Height Map",
398
+ sources=['upload'],
399
+ type='pil'
400
+ )
401
+
402
+ color_label_dropdown = gr.Dropdown(label="Select Material Label", choices=[], value=None)
403
+ save_hm_btn = gr.Button("Save Annotated Height Map")
404
+ hm_download = gr.File(label="Download Annotated Height Map")
405
+
406
+ # When "Load Height Map Preview" is clicked, compute the initial height map and update the dropdown.
407
+ load_hm_btn.click(
408
+ init_height_map_preview,
409
+ inputs=[hm_input_image, hm_max_layers, hm_layer_height, hm_csv_file],
410
+ outputs=[height_map_preview, color_label_dropdown]
411
+ )
412
+ # When "Save Annotated Height Map" is clicked, save the current annotated image.
413
+ save_hm_btn.click(
414
+ save_annotated_height_map,
415
+ inputs=[height_map_preview],
416
+ outputs=hm_download
417
+ )
418
+
419
+ with gr.Tab("Optimization"):
420
+ with gr.Row():
421
+ input_image = gr.ImageEditor(label="Upload Image",
422
+ sources=['upload'],
423
+ type= 'pil'
424
+ )
425
+ csv_file = gr.Textbox(label="CSV File Path", value="materials.csv")
426
+ with gr.Row():
427
+ iterations_slider = gr.Slider(100, 10000, step=100, label="Iterations", value=5000)
428
+ learning_rate_slider = gr.Slider(1e-4, 1e-1, label="Learning Rate", value=1e-2, step=1e-4)
429
+ with gr.Row():
430
+ layer_height_slider = gr.Slider(0.01, 0.1, label="Layer Height (mm)", value=0.04, step=0.01)
431
+ max_layers_slider = gr.Slider(10, 150, label="Max Layers", value=75, step=1)
432
+ with gr.Row():
433
+ background_height_slider = gr.Slider(0.1, 2.0, label="Background Height (mm)", value=0.4, step=0.1)
434
+ output_size_slider = gr.Slider(256, 2048, label="Output Size", value=1024, step=64)
435
+ solver_size_slider = gr.Slider(64, 512, label="Solver Size", value=128, step=16)
436
+ with gr.Row():
437
+ decay_slider = gr.Slider(0.001, 0.1, label="Decay", value=0.01, step=0.001)
438
+ with gr.Row():
439
+ perform_gumbal_search_checkbox = gr.Checkbox(label="Perform Gumbal Search", value=True)
440
+ perform_pruning_checkbox = gr.Checkbox(label="Perform Pruning", value=True)
441
+ with gr.Row():
442
+ save_interval_slider = gr.Slider(1, 100, label="Save Interval (%)", value=20, step=1)
443
+ # Create a separate textbox component for background color
444
+ background_color_textbox = gr.Textbox(label="Background Color", value="#000000")
445
+ with gr.Row():
446
+ start_btn = gr.Button("Start Optimization")
447
+ pause_btn = gr.Button("Pause")
448
+ resume_btn = gr.Button("Resume")
449
+ stop_btn = gr.Button("Stop")
450
+ status_box = gr.Textbox(label="Status", interactive=False)
451
+ live_image = gr.Image(label="Current Composite")
452
+ live_status = gr.Textbox(label="Live Status", interactive=False)
453
+
454
+ start_btn.click(
455
+ start_optimization,
456
+ inputs=[
457
+ input_image, csv_file,
458
+ iterations_slider, learning_rate_slider,
459
+ layer_height_slider, max_layers_slider,
460
+ background_height_slider, background_color_textbox,
461
+ output_size_slider, solver_size_slider,
462
+ decay_slider, perform_gumbal_search_checkbox, perform_pruning_checkbox,
463
+ save_interval_slider
464
+ ],
465
+ outputs=status_box
466
+ )
467
+ pause_btn.click(pause_optimization, outputs=status_box)
468
+ resume_btn.click(resume_optimization, outputs=status_box)
469
+ stop_btn.click(stop_optimization, outputs=status_box)
470
+
471
+ live_interface = gr.Interface(
472
+ fn=live_updates,
473
+ inputs=[],
474
+ outputs=[live_image, live_status],
475
+ live=True,
476
+ title="Live Optimization Updates"
477
+ )
478
+
479
+ with gr.Tab("Post-Processing"):
480
+ gr.Markdown("### Run Post-Processing Steps")
481
+ with gr.Row():
482
+ gumbal_btn = gr.Button("Run Gumbal Search")
483
+ pruning_btn = gr.Button("Run Pruning")
484
+ post_image = gr.Image(label="Post-Processed Composite")
485
+ post_status = gr.Textbox(label="Post-Processing Status", interactive=False)
486
+ gumbal_btn.click(run_gumbal_search_callback, inputs=[], outputs=[post_image, post_status])
487
+ pruning_btn.click(run_pruning_callback, inputs=[], outputs=[post_image, post_status])
488
+
489
+ with gr.Tab("Export Outputs"):
490
+ gr.Markdown("### Export Files")
491
+ export_btn = gr.Button("Export Outputs")
492
+ export_file = gr.File(label="Download Output (STL)")
493
+ export_btn.click(export_outputs_callback, inputs=[], outputs=export_file)
494
+
495
+ demo.launch()
@@ -0,0 +1,225 @@
1
+ import math
2
+ import pandas as pd
3
+ import numpy as np
4
+ import matplotlib.pyplot as plt
5
+ from matplotlib.patches import Rectangle
6
+ import optuna
7
+ import multiprocessing
8
+
9
+ # ===== Helper Functions: Color Conversion =====
10
+ def hex_to_rgb(hex_str):
11
+ """Convert a hex color (e.g. "#d9e0e9") to an RGB numpy array (floats 0-255)."""
12
+ hex_str = hex_str.lstrip('#')
13
+ return np.array([int(hex_str[i:i+2], 16) for i in (0, 2, 4)], dtype=float)
14
+
15
+ def rgb_to_hex(rgb):
16
+ """Convert an RGB numpy array (floats 0-255) to a hex color string."""
17
+ return '#' + ''.join(f'{int(round(x)):02X}' for x in rgb)
18
+
19
+ # ---- sRGB <-> Linear conversion functions ----
20
+ def srgb_to_linear(rgb):
21
+ """Convert an sRGB color (0-255) to linear space (0-1)."""
22
+ rgb = rgb / 255.0
23
+ linear = np.where(rgb <= 0.04045, rgb / 12.92, ((rgb + 0.055) / 1.055) ** 2.4)
24
+ return linear
25
+
26
+ def linear_to_srgb(linear):
27
+ """Convert a linear color (0-1) to sRGB (0-255)."""
28
+ srgb = np.where(linear <= 0.0031308, linear * 12.92, 1.055 * (linear ** (1/2.4)) - 0.055)
29
+ srgb = np.clip(srgb, 0, 1)
30
+ return srgb * 255
31
+
32
+ def model7(params, t, T):
33
+ A, k = params
34
+ if k < 1e-6:
35
+ return 0.0
36
+ return A * math.log(1 + k * (t / T)) / math.log(1 + k)
37
+
38
+ def new_model4(params, t, T):
39
+ A, k, B = params
40
+ return A * math.log(1 + k*(t/T)) + B * (t/T)
41
+
42
+ # ===== Dictionary of Extra Candidate Models =====
43
+ extra_models = {
44
+ "New Model 4: Log-Linear": {
45
+ "func": new_model4,
46
+ "ranges": [(0.0, 1.0), (0.0, 100.0), (0.0, 1.0)],
47
+ "param_names": ["A", "k", "B"]
48
+ },
49
+ }
50
+
51
+ # ===== Loss Calculation =====
52
+ def total_loss_for_model(model_func, params, df, layer_thickness, blending_mode="sRGB"):
53
+ loss = 0.0
54
+ for idx, row in df.iterrows():
55
+ T = float(row['Transmission Distance']) * 0.1
56
+ bg = hex_to_rgb(row['Background Material'])
57
+ fg = hex_to_rgb(row['Layer Material'])
58
+ for layer in range(1, 17):
59
+ t = layer * layer_thickness
60
+ alpha = model_func(params, t, T)
61
+ # Ensure alpha is within [0, 1]
62
+ if np.isscalar(alpha):
63
+ alpha = max(0.0, min(1.0, alpha))
64
+ else:
65
+ alpha = np.clip(alpha, 0.0, 1.0)
66
+ pred = bg + alpha * (fg - bg)
67
+ meas = hex_to_rgb(row[f"Layer {layer}"])
68
+ loss += np.sum((pred - meas) ** 2)
69
+ loss /= len(df) * 16
70
+ return loss, params
71
+
72
+ # ===== Optuna Bayesian Optimization using Process-Based Parallelism =====
73
+ def optuna_search_model_process(model_func, model_name, param_ranges, param_names, df,
74
+ layer_thickness, total_trials, blending_mode="sRGB"):
75
+ print(f"Starting Optuna Bayesian Optimization (process-based) for {model_name} ({blending_mode}) ...")
76
+
77
+ # Define the objective function for Optuna
78
+ def objective(trial):
79
+ params = []
80
+ for i, name in enumerate(param_names):
81
+ low, high = param_ranges[i]
82
+ params.append(trial.suggest_float(name, low, high))
83
+ loss, _ = total_loss_for_model(model_func, params, df, layer_thickness, blending_mode)
84
+ return loss
85
+
86
+ # Create a study with a shared SQLite storage so multiple processes can add trials.
87
+ study = optuna.create_study(
88
+ direction="minimize",
89
+ study_name=model_name,
90
+ storage="sqlite:///{}.db".format(model_name.replace(" ", "_")),
91
+ load_if_exists=True,
92
+ )
93
+
94
+ # Determine the number of worker processes and how many trials each will run.
95
+ cpu_count = multiprocessing.cpu_count()
96
+ trials_per_worker = total_trials // cpu_count
97
+ print(f"Launching {cpu_count} processes, each with {trials_per_worker} trials.")
98
+
99
+ # Function for each process to run its share of trials.
100
+ def run_worker():
101
+ study.optimize(objective, n_trials=trials_per_worker, n_jobs=1)
102
+
103
+ # Launch the worker processes.
104
+ processes = []
105
+ for _ in range(cpu_count):
106
+ p = multiprocessing.Process(target=run_worker)
107
+ p.start()
108
+ processes.append(p)
109
+ for p in processes:
110
+ p.join()
111
+
112
+ best_params_dict = study.best_params
113
+ best_params = [best_params_dict[name] for name in param_names]
114
+ best_loss = study.best_value
115
+ print(f"Best Loss={best_loss:.2f} for {model_name} with params {best_params}")
116
+ return best_params, best_loss, study
117
+
118
+ # ===== Worker Function for Model Search =====
119
+ def run_search_for_model(task):
120
+ model_name, model_data, df, layer_thickness, total_trials, blending_mode = task
121
+ func = model_data["func"]
122
+ param_ranges = model_data["ranges"]
123
+ param_names = model_data["param_names"]
124
+ best_params, best_loss, study = optuna_search_model_process(
125
+ func, model_name, param_ranges, param_names,
126
+ df, layer_thickness, total_trials, blending_mode
127
+ )
128
+ return {
129
+ "model": model_name,
130
+ "blending_mode": blending_mode,
131
+ "best_params": best_params,
132
+ "best_loss": best_loss,
133
+ "param_names": param_names,
134
+ "func": func,
135
+ "study": study
136
+ }
137
+
138
+ # ===== Main Function =====
139
+ def main():
140
+ # Load CSV (adjust path as needed)
141
+ df = pd.read_csv("printed_colors.csv")
142
+ layer_thickness = 0.04 # mm per layer
143
+ blending_mode = "sRGB" # using sRGB blending
144
+
145
+ # --- Phase 1: Process-Based Optuna Bayesian Optimization search for all models ---
146
+ total_trials = 5000 # Total trials per model (will be divided among processes)
147
+
148
+ results = []
149
+ for model_name, model_data in extra_models.items():
150
+ result = run_search_for_model((model_name, model_data, df, layer_thickness, total_trials, blending_mode))
151
+ results.append(result)
152
+
153
+ results.sort(key=lambda x: x["best_loss"])
154
+ print("\n=== Phase 1: Results (Optuna Process-Based Optimization) ===")
155
+ for res in results:
156
+ param_str = ", ".join(f"{name}={val:.4f}" for name, val in zip(res["param_names"], res["best_params"]))
157
+ print(f"{res['model']} ({res['blending_mode']}): Loss = {res['best_loss']:.2f} with {param_str}")
158
+
159
+ # --- Plot Sample Comparison for the Best Model (all CSV rows) ---
160
+ best_result = results[0]
161
+ func = best_result["func"]
162
+ best_params = best_result["best_params"]
163
+
164
+ # Optional: plot the model curve for new_model4 for a range of t values.
165
+ y = [new_model4(best_params, t/10, 1) for t in range(1, 11)]
166
+ x = list(range(1, 11))
167
+ plt.figure()
168
+ plt.plot(x, y, marker='o')
169
+ plt.title(f"{best_result['model']} Curve")
170
+ plt.xlabel("t/10")
171
+ plt.ylabel("Model Output")
172
+ plt.show()
173
+
174
+ n_csv = len(df)
175
+ # Create a figure with 2 rows per CSV row (measured and predicted) for the best model.
176
+ n_rows = 2 * n_csv
177
+ n_cols = 1
178
+
179
+ fig, axes = plt.subplots(nrows=n_rows, ncols=n_cols, figsize=(8, 3 * n_csv))
180
+ fig.suptitle("Best Model: All CSV Rows\nMeasured (top row) vs. Predicted (bottom row)", fontsize=16)
181
+
182
+ for idx in range(n_csv):
183
+ row_data = df.iloc[idx]
184
+ T = float(row_data['Transmission Distance'])
185
+ bg = hex_to_rgb(row_data['Background Material'])
186
+ fg = hex_to_rgb(row_data['Layer Material'])
187
+ measured = [hex_to_rgb(row_data[f"Layer {layer}"]) for layer in range(1, 17)]
188
+ predicted = []
189
+ for layer in range(1, 17):
190
+ t = layer * layer_thickness
191
+ alpha = func(best_params, t, T)
192
+ if np.isscalar(alpha):
193
+ alpha = max(0.0, min(1.0, alpha))
194
+ else:
195
+ alpha = np.clip(alpha, 0.0, 1.0)
196
+ pred = bg + alpha * (fg - bg)
197
+ predicted.append(pred)
198
+
199
+ # Plot measured colors
200
+ ax_meas = axes[2*idx]
201
+ for j in range(16):
202
+ rect = Rectangle((j, 0), 1, 1, color=np.clip(measured[j] / 255, 0, 1))
203
+ ax_meas.add_patch(rect)
204
+ ax_meas.set_xlim(0, 16)
205
+ ax_meas.set_ylim(0, 1)
206
+ ax_meas.set_xticks([])
207
+ ax_meas.set_yticks([])
208
+ ax_meas.set_ylabel(f"Row {idx+1}\nMeasured", fontsize=10)
209
+
210
+ # Plot predicted colors
211
+ ax_pred = axes[2*idx + 1]
212
+ for j in range(16):
213
+ rect = Rectangle((j, 0), 1, 1, color=np.clip(predicted[j] / 255, 0, 1))
214
+ ax_pred.add_patch(rect)
215
+ ax_pred.set_xlim(0, 16)
216
+ ax_pred.set_ylim(0, 1)
217
+ ax_pred.set_xticks([])
218
+ ax_pred.set_yticks([])
219
+ ax_pred.set_ylabel(f"Row {idx+1}\nPredicted", fontsize=10)
220
+
221
+ plt.tight_layout(rect=[0, 0, 1, 0.95])
222
+ plt.show()
223
+
224
+ if __name__ == '__main__':
225
+ main()
@@ -80,7 +80,7 @@ def download_all_pages(url):
80
80
  tbar.update(1)
81
81
  return all_results
82
82
 
83
- def main():
83
+ def download_filament_info():
84
84
  """
85
85
  Main function to manage the swatch data update process.
86
86
 
@@ -120,4 +120,4 @@ def main():
120
120
  save_local_version(api_version)
121
121
  print("Download complete, data saved to", DATA_FILE)
122
122
  if __name__ == "__main__":
123
- main()
123
+ download_filament_info()
@@ -27,6 +27,50 @@ import numpy as np
27
27
  import struct
28
28
 
29
29
 
30
+ def extract_colors_from_swatches(swatch_data):
31
+
32
+ # we keep only data with transmission distance
33
+ swatch_data = [swatch for swatch in swatch_data if swatch["td"]]
34
+
35
+ #For now we load it and convert it in the same way as the hueforge csv files
36
+ out = {}
37
+ for swatch in swatch_data:
38
+ brand = swatch["manufacturer"]["name"]
39
+ name = swatch["color_name"]
40
+ color = swatch["hex_color"]
41
+ td = swatch["td"]
42
+ out[(brand, name)] = (color, td)
43
+
44
+ #convert to the same format as the hueforge csv files
45
+ material_names = [brand + " - " + name for (brand, name) in out.keys()]
46
+ material_colors = jnp.array([hex_to_rgb("#"+color) for color, _ in out.values()], dtype=jnp.float64)
47
+ material_TDs = jnp.array([td for _, td in out.values()], dtype=jnp.float64)
48
+ colors_list = [color for color, _ in out.values()]
49
+
50
+ return material_colors, material_TDs, material_names, colors_list
51
+
52
+ def swatch_data_to_table(swatch_data):
53
+ """
54
+ Converts swatch JSON data into a table (list of dicts) with columns:
55
+ "Brand", "Name", "Transmission Distance", "Hex Color".
56
+ """
57
+ table = []
58
+ for swatch in swatch_data:
59
+ if not swatch["td"]:
60
+ continue
61
+ brand = swatch["manufacturer"]["name"]
62
+ name = swatch["color_name"]
63
+ hex_color = swatch["hex_color"]
64
+ td = swatch["td"]
65
+ table.append({
66
+ "Brand": brand,
67
+ "Name": name,
68
+ "Transmission Distance": td,
69
+ "Hex Color": f"#{hex_color}"
70
+ })
71
+ return table
72
+
73
+
30
74
  def resize_image(img, max_size):
31
75
  h_img, w_img, _ = img.shape
32
76
  if w_img >= h_img:
File without changes