hvoss-techfak commited on
Commit
91a801a
·
1 Parent(s): fe75e01

Autoforge wrapper app

Browse files

# Conflicts:
# README.md

Files changed (4) hide show
  1. README.md +36 -1
  2. app.py +783 -0
  3. default_materials.csv +4 -0
  4. requirements.txt +2 -0
README.md CHANGED
@@ -11,4 +11,39 @@ license: other
11
  short_description: Generating 3D printed layered models from an input image
12
  ---
13
 
14
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
11
  short_description: Generating 3D printed layered models from an input image
12
  ---
13
 
14
+ # AutoForge
15
+
16
+ AutoForge is a Python tool for generating 3D printed layered models from an input image. Using a learned optimization strategy with a Gumbel softmax formulation, AutoForge assigns materials per layer and produces both a discretized composite image and a 3D-printable STL file. It also generates swap instructions to guide the printer through material changes during a multi-material print. \
17
+
18
+ **TLDR:** It uses a picture to generate a 3D layer image that you can print with a 3d printer. Similar to [Hueforge](https://shop.thehueforge.com/), but without the manual work (and without the artistic control).
19
+
20
+ ## Example
21
+ All examples use only the 13 BambuLab Basic filaments, currently available in Hueforge, the background color is set to black.
22
+ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at most 8 different colors and swaps the filament at most 20 times.
23
+ <div style="display: flex; justify-content: center; gap: 20px;">
24
+ <div style="text-align: center;">
25
+ <h3>Input Image</h3>
26
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi.jpg" width="200" />
27
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature.jpg" width="200" />
28
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat.jpg" width="200" />
29
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/chameleon.jpg" width="200" />
30
+ </div>
31
+ <div style="text-align: center;">
32
+ <h3>Autoforge Output</h3>
33
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_discretized.png" width="200" />
34
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_discretized.png" width="200" />
35
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat_discretized.png" width="200" />
36
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/chameleon_discretized.png" width="200" />
37
+ </div>
38
+ </div>
39
+
40
+ ## Features
41
+
42
+ - **Image-to-Model Conversion**: Converts an input image into a layered model suitable for 3D printing.
43
+ - **Learned Optimization**: Optimizes per-pixel height and per-layer material assignments using PyTorch.
44
+ - **Learned Heightmap**: Optimizes the height of the layered model to create more detailed prints.
45
+ - **Gumbel Softmax Sampling**: Leverages the Gumbel softmax method to decide material assignments for each layer.
46
+ - **STL File Generation**: Exports an ASCII STL file based on the optimized height map.
47
+ - **Swap Instructions**: Generates clear swap instructions for changing materials during printing.
48
+ - **Live Visualization**: (Optional) Displays live composite images during the optimization process.
49
+ - **Hueforge export**: Outputs a project file that can be opened with hueforge.
app.py ADDED
@@ -0,0 +1,783 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import pandas as pd
3
+ import os
4
+ import subprocess
5
+ import time
6
+ import shutil
7
+ import sys
8
+ from datetime import datetime
9
+ import re
10
+ from PIL import Image
11
+
12
+ # --- Configuration ---
13
+ #AUTFORGE_SCRIPT_PATH = "auto_forge.py" # Make sure this points to your script
14
+ DEFAULT_MATERIALS_CSV = "default_materials.csv"
15
+ GRADIO_OUTPUT_BASE_DIR = "output"
16
+ os.makedirs(GRADIO_OUTPUT_BASE_DIR, exist_ok=True)
17
+
18
+ REQUIRED_SCRIPT_COLS = ["Brand", " Name", " TD", " Color"]
19
+ DISPLAY_COL_MAP = {
20
+ "Brand": "Brand",
21
+ " Name": "Name",
22
+ " TD": "TD",
23
+ " Color": "Color (Hex)",
24
+ }
25
+
26
+
27
+ def ensure_required_cols(df, *, in_display_space):
28
+ """
29
+ Return a copy of *df* with every required column present.
30
+ If *in_display_space* is True we use the display names
31
+ (Brand, Name, TD, Color (Hex)); otherwise we use the script names.
32
+ """
33
+ target_cols = (
34
+ DISPLAY_COL_MAP if in_display_space else {k: k for k in REQUIRED_SCRIPT_COLS}
35
+ )
36
+ df_fixed = df.copy()
37
+ for col_script, col_display in target_cols.items():
38
+ if col_display not in df_fixed.columns:
39
+ # sensible defaults
40
+ if "TD" in col_display:
41
+ default = 0.0
42
+ elif "Color" in col_display:
43
+ default = "#000000"
44
+ elif "Owned" in col_display: # NEW
45
+ default = "false"
46
+ else:
47
+ default = ""
48
+ df_fixed[col_display] = default
49
+ # order columns nicely
50
+ return df_fixed[list(target_cols.values())]
51
+
52
+
53
+ def rgba_to_hex(col: str) -> str:
54
+ """
55
+ Turn 'rgba(r, g, b, a)' or 'rgb(r, g, b)' into '#RRGGBB'.
56
+ If the input is already a hex code or anything unexpected,
57
+ return it unchanged.
58
+ """
59
+ if not isinstance(col, str):
60
+ return col
61
+ col = col.strip()
62
+ if col.startswith("#"): # already fine
63
+ return col.upper()
64
+
65
+ m = re.match(
66
+ r"rgba?\(\s*([\d.]+)\s*,\s*([\d.]+)\s*,\s*([\d.]+)(?:\s*,\s*[\d.]+)?\s*\)",
67
+ col,
68
+ )
69
+ if not m:
70
+ return col # not something we recognise
71
+
72
+ r, g, b = (int(float(x)) for x in m.groups()[:3])
73
+ return "#{:02X}{:02X}{:02X}".format(r, g, b)
74
+
75
+
76
+ # --- Helper Functions ---
77
+ def get_script_args_info(exclude_args=None):
78
+ if exclude_args is None:
79
+ exclude_args = []
80
+
81
+ all_args_info = [
82
+ # input_image is handled separately in the UI
83
+ {
84
+ "name": "--iterations",
85
+ "type": "number",
86
+ "default": 2000,
87
+ "help": "Number of optimization iterations",
88
+ },
89
+ {
90
+ "name": "--layer_height",
91
+ "type": "number",
92
+ "default": 0.04,
93
+ "step": 0.01,
94
+ "help": "Layer thickness in mm",
95
+ },
96
+ {
97
+ "name": "--max_layers",
98
+ "type": "number",
99
+ "default": 75,
100
+ "precision": 0,
101
+ "help": "Maximum number of layers",
102
+ },
103
+ {
104
+ "name": "--learning_rate",
105
+ "type": "number",
106
+ "default": 0.015,
107
+ "step": 0.001,
108
+ "help": "Learning rate for optimization",
109
+ },
110
+ {
111
+ "name": "--background_height",
112
+ "type": "number",
113
+ "default": 0.4,
114
+ "step": 0.01,
115
+ "help": "Height of the background in mm",
116
+ },
117
+ {
118
+ "name": "--background_color",
119
+ "type": "colorpicker",
120
+ "default": "#000000",
121
+ "help": "Background color",
122
+ },
123
+ {
124
+ "name": "--stl_output_size",
125
+ "type": "number",
126
+ "default": 100,
127
+ "precision": 0,
128
+ "help": "Size of the longest dimension of the output STL file in mm",
129
+ },
130
+ {
131
+ "name": "--nozzle_diameter",
132
+ "type": "number",
133
+ "default": 0.4,
134
+ "step": 0.1,
135
+ "help": "Diameter of the printer nozzle in mm",
136
+ },
137
+ {
138
+ "name": "--pruning_max_colors",
139
+ "type": "number",
140
+ "default": 10,
141
+ "precision": 0,
142
+ "help": "Max number of colors allowed after pruning",
143
+ },
144
+ {
145
+ "name": "--pruning_max_swaps",
146
+ "type": "number",
147
+ "default": 20,
148
+ "precision": 0,
149
+ "help": "Max number of swaps allowed after pruning",
150
+ },
151
+ {
152
+ "name": "--pruning_max_layer",
153
+ "type": "number",
154
+ "default": 75,
155
+ "precision": 0,
156
+ "help": "Max number of layers allowed after pruning",
157
+ },
158
+ {
159
+ "name": "--warmup_fraction",
160
+ "type": "slider",
161
+ "default": 1.0,
162
+ "min": 0.0,
163
+ "max": 1.0,
164
+ "step": 0.01,
165
+ "help": "Fraction of iterations for keeping the tau at the initial value",
166
+ },
167
+ {
168
+ "name": "--learning_rate_warmup_fraction",
169
+ "type": "slider",
170
+ "default": 0.25,
171
+ "min": 0.0,
172
+ "max": 1.0,
173
+ "step": 0.01,
174
+ "help": "Fraction of iterations that the learning rate is increasing (warmup)",
175
+ },
176
+ # {
177
+ # "name": "--init_tau",
178
+ # "type": "number",
179
+ # "default": 1.0,
180
+ # "help": "Initial tau value for Gumbel-Softmax",
181
+ # },
182
+ # {
183
+ # "name": "--final_tau",
184
+ # "type": "number",
185
+ # "default": 0.01,
186
+ # "help": "Final tau value for Gumbel-Softmax",
187
+ # },
188
+ # {
189
+ # "name": "--min_layers",
190
+ # "type": "number",
191
+ # "default": 0,
192
+ # "precision": 0,
193
+ # "help": "Minimum number of layers. Used for pruning.",
194
+ # },
195
+ {
196
+ "name": "--early_stopping",
197
+ "type": "number",
198
+ "default": 1500,
199
+ "precision": 0,
200
+ "help": "Number of steps without improvement before stopping",
201
+ },
202
+ {
203
+ "name": "--random_seed",
204
+ "type": "number",
205
+ "default": 0,
206
+ "precision": 0,
207
+ "help": "Specify the random seed, or use 0 for automatic generation",
208
+ },
209
+ {
210
+ "name": "--num_init_rounds",
211
+ "type": "number",
212
+ "default": 32,
213
+ "precision": 0,
214
+ "help": "Number of rounds to choose the starting height map from.",
215
+ },
216
+ ]
217
+ return [arg for arg in all_args_info if arg["name"] not in exclude_args]
218
+
219
+
220
+ # Initial filament data
221
+ initial_filament_data = {
222
+ "Brand": ["Generic", "Generic", "Generic"],
223
+ " Name": ["PLA Black", "PLA Grey", "PLA White"],
224
+ " TD": [1.0, 1.0, 1.0],
225
+ " Color": ["#000000", "#808080", "#FFFFFF"],
226
+ " Owned": ["true", "true", "true"], # ← add
227
+ }
228
+ initial_df = pd.DataFrame(initial_filament_data)
229
+
230
+ if os.path.exists(DEFAULT_MATERIALS_CSV):
231
+ try:
232
+ initial_df = pd.read_csv(DEFAULT_MATERIALS_CSV)
233
+ for col in ["Brand", " Name", " TD", " Color"]:
234
+ if col not in initial_df.columns:
235
+ initial_df[col] = None
236
+ initial_df = initial_df[["Brand", " Name", " TD", " Color"]].astype(
237
+ {" TD": float, " Color": str}
238
+ )
239
+ except Exception as e:
240
+ print(f"Warning: Could not load {DEFAULT_MATERIALS_CSV}: {e}. Using default.")
241
+ initial_df = pd.DataFrame(initial_filament_data)
242
+ else:
243
+ initial_df.to_csv(DEFAULT_MATERIALS_CSV, index=False)
244
+
245
+
246
+ # Helper for creating an empty 10-tuple for error returns
247
+ def create_empty_error_outputs(log_message=""):
248
+ return (
249
+ log_message, # progress_output
250
+ None, # final_image_preview
251
+ gr.update(visible=False, interactive=False), # ### ZIP: download_zip
252
+ )
253
+
254
+
255
+ # --- Gradio UI Definition ---
256
+ with gr.Blocks(theme=gr.themes.Soft()) as demo:
257
+ gr.Markdown("# Autoforge Web UI")
258
+
259
+ filament_df_state = gr.State(initial_df.copy())
260
+ current_run_output_dir = gr.State(None)
261
+
262
+ with gr.Tabs():
263
+ with gr.TabItem("Filament Management"):
264
+ gr.Markdown(
265
+ 'Manage your filament list. This list will be saved as a CSV and used by the Autoforge process. \n To remove a filament simply rightclick on any of the fields and select "Delete Row"'
266
+ )
267
+ with gr.Row():
268
+ load_csv_button = gr.UploadButton(
269
+ "Load Filaments CSV", file_types=[".csv"]
270
+ )
271
+ save_csv_button = gr.Button("Save Current Filaments to CSV")
272
+ filament_table = gr.DataFrame(
273
+ value=ensure_required_cols(
274
+ initial_df.copy().rename(
275
+ columns={" Name": "Name", " TD": "TD", " Color": "Color (Hex)"}
276
+ ),
277
+ in_display_space=True,
278
+ ),
279
+ headers=["Brand", "Name", "TD", "Color (Hex)"],
280
+ datatype=["str", "str", "number", "str"],
281
+ interactive=True,
282
+ label="Filaments",
283
+ )
284
+ gr.Markdown("### Add New Filament")
285
+ with gr.Row():
286
+ new_brand = gr.Textbox(label="Brand")
287
+ new_name = gr.Textbox(label="Name")
288
+ with gr.Row():
289
+ new_td = gr.Number(
290
+ label="TD (Transmission/Opacity)",
291
+ value=1.0,
292
+ minimum=0,
293
+ maximum=100,
294
+ step=0.1,
295
+ )
296
+ new_color_hex = gr.ColorPicker(label="Color", value="#FF0000")
297
+ add_filament_button = gr.Button("Add Filament to Table")
298
+ download_csv_trigger = gr.File(
299
+ label="Download Filament CSV", visible=False, interactive=False
300
+ )
301
+
302
+ def update_filament_df_state_from_table(display_df):
303
+ display_df = ensure_required_cols(display_df, in_display_space=True)
304
+
305
+ # make sure every colour is hex
306
+ if "Color (Hex)" in display_df.columns:
307
+ display_df["Color (Hex)"] = display_df["Color (Hex)"].apply(
308
+ rgba_to_hex
309
+ )
310
+
311
+ script_df = display_df.rename(
312
+ columns={"Name": " Name", "TD": " TD", "Color (Hex)": " Color"}
313
+ )
314
+ script_df = ensure_required_cols(script_df, in_display_space=False)
315
+ filament_df_state.value = script_df
316
+
317
+ def add_filament_to_table(current_display_df, brand, name, td, color_hex):
318
+ if not brand or not name:
319
+ gr.Warning("Brand and Name cannot be empty.")
320
+ return current_display_df
321
+
322
+ color_hex = rgba_to_hex(color_hex) # <-- new line
323
+
324
+ new_row = pd.DataFrame(
325
+ [{"Brand": brand, "Name": name, "TD": td, "Color (Hex)": color_hex}]
326
+ )
327
+ updated_display_df = pd.concat(
328
+ [current_display_df, new_row], ignore_index=True
329
+ )
330
+ update_filament_df_state_from_table(updated_display_df)
331
+ return updated_display_df
332
+
333
+ def load_filaments_from_csv_upload(file_obj):
334
+ if file_obj is None:
335
+ current_script_df = filament_df_state.value
336
+ if current_script_df is not None and not current_script_df.empty:
337
+ return current_script_df.rename(
338
+ columns={
339
+ " Name": "Name",
340
+ " TD": "TD",
341
+ " Color": "Color (Hex)",
342
+ }
343
+ )
344
+ return initial_df.copy().rename(
345
+ columns={" Name": "Name", " TD": "TD", " Color": "Color (Hex)"}
346
+ )
347
+ try:
348
+ loaded_script_df = pd.read_csv(file_obj.name)
349
+ loaded_script_df = ensure_required_cols(
350
+ loaded_script_df, in_display_space=False
351
+ )
352
+ expected_cols = ["Brand", " Name", " TD", " Color"]
353
+ if not all(
354
+ col in loaded_script_df.columns for col in expected_cols
355
+ ):
356
+ gr.Error(
357
+ f"CSV must contain columns: {', '.join(expected_cols)}. Found: {loaded_script_df.columns.tolist()}"
358
+ )
359
+ current_script_df = filament_df_state.value
360
+ if (
361
+ current_script_df is not None
362
+ and not current_script_df.empty
363
+ ):
364
+ return current_script_df.rename(
365
+ columns={
366
+ " Name": "Name",
367
+ " TD": "TD",
368
+ " Color": "Color (Hex)",
369
+ }
370
+ )
371
+ return initial_df.copy().rename(
372
+ columns={
373
+ " Name": "Name",
374
+ " TD": "TD",
375
+ " Color": "Color (Hex)",
376
+ }
377
+ )
378
+ filament_df_state.value = loaded_script_df.copy()
379
+ return loaded_script_df.rename(
380
+ columns={" Name": "Name", " TD": "TD", " Color": "Color (Hex)"}
381
+ )
382
+ except Exception as e:
383
+ gr.Error(f"Error loading CSV: {e}")
384
+ current_script_df = filament_df_state.value
385
+ if current_script_df is not None and not current_script_df.empty:
386
+ return current_script_df.rename(
387
+ columns={
388
+ " Name": "Name",
389
+ " TD": "TD",
390
+ " Color": "Color (Hex)",
391
+ }
392
+ )
393
+ return initial_df.copy().rename(
394
+ columns={" Name": "Name", " TD": "TD", " Color": "Color (Hex)"}
395
+ )
396
+
397
+ def save_filaments_to_file_for_download(current_script_df_from_state):
398
+ if (
399
+ current_script_df_from_state is None
400
+ or current_script_df_from_state.empty
401
+ ):
402
+ gr.Warning("Filament table is empty. Nothing to save.")
403
+ return None
404
+ df_to_save = current_script_df_from_state.copy()
405
+ required_cols = ["Brand", " Name", " TD", " Color"]
406
+ if not all(col in df_to_save.columns for col in required_cols):
407
+ gr.Error(
408
+ f"Cannot save. DataFrame missing required script columns. Expected: {required_cols}. Found: {df_to_save.columns.tolist()}"
409
+ )
410
+ return None
411
+ temp_dir = os.path.join(GRADIO_OUTPUT_BASE_DIR, "_temp_downloads")
412
+ os.makedirs(temp_dir, exist_ok=True)
413
+ temp_filament_csv_path = os.path.join(
414
+ temp_dir,
415
+ f"filaments_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
416
+ )
417
+ try:
418
+ df_to_save.to_csv(temp_filament_csv_path, index=False)
419
+ gr.Info("Filaments prepared for download.")
420
+ return gr.File(
421
+ value=temp_filament_csv_path,
422
+ label="Download Filament CSV",
423
+ interactive=True,
424
+ visible=True,
425
+ )
426
+ except Exception as e:
427
+ gr.Error(f"Error saving CSV for download: {e}")
428
+ return None
429
+
430
+ filament_table.change(
431
+ update_filament_df_state_from_table,
432
+ inputs=[filament_table],
433
+ outputs=None,
434
+ queue=False,
435
+ )
436
+ add_filament_button.click(
437
+ add_filament_to_table,
438
+ inputs=[filament_table, new_brand, new_name, new_td, new_color_hex],
439
+ outputs=[filament_table],
440
+ )
441
+ load_csv_button.upload(
442
+ load_filaments_from_csv_upload,
443
+ inputs=[load_csv_button],
444
+ outputs=[filament_table],
445
+ )
446
+ save_csv_button.click(
447
+ save_filaments_to_file_for_download,
448
+ inputs=[filament_df_state],
449
+ outputs=[download_csv_trigger],
450
+ )
451
+
452
+ with gr.TabItem("Run Autoforge"):
453
+ accordion_params_dict = {}
454
+ accordion_params_ordered_names = []
455
+
456
+ with gr.Row():
457
+ with gr.Column(scale=1):
458
+ gr.Markdown("### Input Image (Required)")
459
+ input_image_component = gr.Image(
460
+ type="filepath",
461
+ image_mode="RGBA",
462
+ label="Upload Image",
463
+ sources=["upload"],
464
+ interactive=True,
465
+ )
466
+ with gr.Column(scale=2):
467
+ gr.Markdown("### Autoforge Parameters")
468
+ with gr.Accordion("Progress & Output", open=True):
469
+ final_image_preview = gr.Image(
470
+ label="Final Model Preview",
471
+ type="filepath",
472
+ interactive=False,
473
+ )
474
+ with gr.Row():
475
+ download_zip = gr.File( # was visible=True
476
+ label="Download all results (.zip)",
477
+ interactive=True,
478
+ visible=False,
479
+ )
480
+ with gr.Row():
481
+ with gr.Accordion("Adjust Parameters", open=False):
482
+ args_for_accordion = get_script_args_info(
483
+ exclude_args=["--input_image"]
484
+ )
485
+
486
+ for arg in args_for_accordion:
487
+ label, info, default_val = (
488
+ f"{arg['name']}",
489
+ arg["help"],
490
+ arg.get("default"),
491
+ )
492
+ if arg["type"] == "number":
493
+ accordion_params_dict[arg["name"]] = gr.Number(
494
+ label=label,
495
+ value=default_val,
496
+ info=info,
497
+ minimum=arg.get("min"),
498
+ maximum=arg.get("max"),
499
+ step=arg.get(
500
+ "step",
501
+ 0.001 if isinstance(default_val, float) else 1,
502
+ ),
503
+ precision=arg.get("precision", None),
504
+ )
505
+ elif arg["type"] == "slider":
506
+ accordion_params_dict[arg["name"]] = gr.Slider(
507
+ label=label,
508
+ value=default_val,
509
+ info=info,
510
+ minimum=arg.get("min", 0),
511
+ maximum=arg.get("max", 1),
512
+ step=arg.get("step", 0.01),
513
+ )
514
+ elif arg["type"] == "checkbox":
515
+ accordion_params_dict[arg["name"]] = gr.Checkbox(
516
+ label=label, value=default_val, info=info
517
+ )
518
+ elif arg["type"] == "colorpicker":
519
+ accordion_params_dict[arg["name"]] = gr.ColorPicker(
520
+ label=label, value=default_val, info=info
521
+ )
522
+ else:
523
+ accordion_params_dict[arg["name"]] = gr.Textbox(
524
+ label=label, value=str(default_val), info=info
525
+ )
526
+ accordion_params_ordered_names.append(arg["name"])
527
+
528
+ run_button = gr.Button(
529
+ "Run Autoforge Process",
530
+ variant="primary",
531
+ elem_id="run_button_full_width",
532
+ )
533
+
534
+
535
+ progress_output = gr.Textbox(
536
+ label="Console Output",
537
+ lines=15,
538
+ autoscroll=True,
539
+ show_copy_button=False,
540
+ )
541
+
542
+ # --- Backend Function for Running the Script ---
543
+ def execute_autoforge_script(
544
+ current_filaments_df_state_val, input_image_path, *accordion_param_values
545
+ ):
546
+ # 0. Validate Inputs
547
+ if (
548
+ not input_image_path
549
+ ): # Covers None and empty string from gr.Image(type="filepath")
550
+ gr.Error("Input Image is required! Please upload an image.")
551
+ return create_empty_error_outputs("Error: Input Image is required!")
552
+
553
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
554
+ run_output_dir_val = os.path.join(GRADIO_OUTPUT_BASE_DIR, f"run_{timestamp}")
555
+ os.makedirs(run_output_dir_val, exist_ok=True)
556
+ current_run_output_dir.value = run_output_dir_val
557
+
558
+ # 1. Save current filaments
559
+ if (
560
+ current_filaments_df_state_val is None
561
+ or current_filaments_df_state_val.empty
562
+ ):
563
+ gr.Error("Filament table is empty. Please add filaments.")
564
+ return create_empty_error_outputs("Error: Filament table is empty.")
565
+
566
+ temp_filament_csv = os.path.join(run_output_dir_val, "materials.csv")
567
+ df_to_save = current_filaments_df_state_val.copy()
568
+ required_cols = ["Brand", " Name", " TD", " Color"]
569
+ missing_cols = [col for col in required_cols if col not in df_to_save.columns]
570
+ if missing_cols:
571
+ err_msg = (
572
+ f"Error: Filament data is missing columns: {', '.join(missing_cols)}."
573
+ )
574
+ gr.Error(err_msg)
575
+ return create_empty_error_outputs(err_msg)
576
+ try:
577
+ df_to_save.to_csv(temp_filament_csv, index=False)
578
+ except Exception as e:
579
+ err_msg = f"Error saving temporary filament CSV: {e}"
580
+ gr.Error(err_msg)
581
+ return create_empty_error_outputs(err_msg)
582
+
583
+ # 2. Construct command
584
+ python_executable = sys.executable or "python"
585
+ command = ["autoforge",]
586
+ command.extend(["--csv_file", temp_filament_csv])
587
+ command.extend(["--output_folder", run_output_dir_val])
588
+ command.extend(["--disable_visualization_for_gradio","1"])
589
+
590
+ base_filename = os.path.basename(input_image_path)
591
+ script_input_image_path = os.path.join(run_output_dir_val, base_filename)
592
+ try:
593
+ img = Image.open(input_image_path)
594
+ # decide where to store the image we pass to Autoforge
595
+ base_no_ext, _ = os.path.splitext(os.path.basename(input_image_path))
596
+ script_input_image_path = os.path.join(
597
+ run_output_dir_val, f"{base_no_ext}.png"
598
+ )
599
+
600
+ if img.mode in ("RGBA", "LA") or (
601
+ img.mode == "P" and "transparency" in img.info
602
+ ):
603
+ # the uploaded file has an alpha channel – save it as PNG
604
+ img.save(script_input_image_path, format="PNG")
605
+ else:
606
+ # no alpha present – just copy the file in whatever format it was
607
+ script_input_image_path = os.path.join(
608
+ run_output_dir_val, os.path.basename(input_image_path)
609
+ )
610
+ shutil.copy(input_image_path, script_input_image_path)
611
+
612
+ command.extend(["--input_image", script_input_image_path])
613
+ except Exception as e:
614
+ err_msg = f"Error handling input image: {e}"
615
+ gr.Error(err_msg)
616
+ return create_empty_error_outputs(err_msg)
617
+
618
+ param_dict = dict(zip(accordion_params_ordered_names, accordion_param_values))
619
+ for arg_name, arg_widget_val in param_dict.items():
620
+ if arg_widget_val is None or arg_widget_val == "":
621
+ arg_info_list = [
622
+ item for item in get_script_args_info() if item["name"] == arg_name
623
+ ] # get full list to check type
624
+ if (
625
+ arg_info_list
626
+ and arg_info_list[0]["type"] == "checkbox"
627
+ and arg_widget_val is False
628
+ ):
629
+ continue
630
+ else:
631
+ continue
632
+ if isinstance(arg_widget_val, bool):
633
+ if arg_widget_val:
634
+ command.append(arg_name)
635
+ else:
636
+ command.extend([arg_name, str(arg_widget_val)])
637
+
638
+ # 3. Run script
639
+ log_output = (
640
+ f"Starting Autoforge process at "
641
+ f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
642
+ f"Output directory: {run_output_dir_val}\n"
643
+ f"Command: {' '.join(command)}\n\n"
644
+ )
645
+
646
+ yield create_empty_error_outputs(log_output) # clear UI and show header
647
+
648
+ process = subprocess.Popen(
649
+ command,
650
+ stdout=subprocess.PIPE,
651
+ stderr=subprocess.PIPE,
652
+ text=True,
653
+ bufsize=1,
654
+ universal_newlines=True,
655
+ )
656
+
657
+ # ---- helper: read stdout in a background thread -------------------
658
+ from threading import Thread
659
+ from queue import Queue, Empty
660
+
661
+ def _enqueue(pipe, q):
662
+ """Forward stdout/stderr to a queue, emitting on both '\n' and '\r'."""
663
+ buf = ""
664
+ while True:
665
+ ch = pipe.read(1) # read a single character
666
+ if ch == "": # EOF
667
+ if buf:
668
+ q.put(buf) # flush whatever is left
669
+ break
670
+ buf += ch
671
+ if ch in ("\n", "\r"): # tqdm uses '\r'
672
+ q.put(buf)
673
+ buf = ""
674
+ pipe.close()
675
+
676
+ q_out = Queue()
677
+ Thread(target=_enqueue, args=(process.stdout, q_out), daemon=True).start()
678
+ Thread(target=_enqueue, args=(process.stderr, q_out), daemon=True).start()
679
+
680
+ preview_mtime = 0
681
+ last_push = 0
682
+
683
+ def _maybe_new_preview():
684
+ """
685
+ If vis_temp.png has a newer mtime than last time, copy it to a
686
+ stamped name (to defeat browser cache) and return that path.
687
+ Otherwise return gr.update() so the image stays as-is.
688
+ """
689
+ from gradio import update # local import for clarity
690
+
691
+ nonlocal preview_mtime
692
+
693
+ src = os.path.join(run_output_dir_val, "vis_temp.png")
694
+ if not os.path.exists(src):
695
+ return update() # nothing new, keep old
696
+
697
+ mtime = os.path.getmtime(src)
698
+ if mtime <= preview_mtime: # unchanged
699
+ return update() # → no UI update
700
+
701
+ return src # → refresh image
702
+
703
+ # ---- main loop: poll every 0.5 s ----------------------------------
704
+ while process.poll() is None or not q_out.empty():
705
+ # drain whatever is waiting in stdout
706
+ try:
707
+ while True:
708
+ log_output += q_out.get_nowait()
709
+ except Empty:
710
+ pass
711
+
712
+ now = time.time()
713
+ if now - last_push >= 1.0: # 500 ms tick
714
+ current_preview = _maybe_new_preview()
715
+ yield (
716
+ log_output,
717
+ current_preview,
718
+ gr.update(), # ### ZIP PATCH: placeholder for zip widget
719
+ )
720
+ last_push = now
721
+
722
+ time.sleep(0.05) # keep CPU load low
723
+
724
+ return_code = process.wait()
725
+ log_output += (
726
+ "\nAutoforge process completed successfully!"
727
+ if return_code == 0
728
+ else f"\nAutoforge process failed with exit code {return_code}."
729
+ )
730
+
731
+ # make sure we show the final preview (if any)
732
+ final_preview = _maybe_new_preview() or os.path.join(
733
+ run_output_dir_val, "final_model.png"
734
+ )
735
+
736
+ zip_base = os.path.join(
737
+ run_output_dir_val, "autoforge_results"
738
+ ) # ### ZIP PATCH
739
+ zip_path = shutil.make_archive(zip_base, "zip", run_output_dir_val)
740
+
741
+ # 4. Prepare output file paths
742
+ png_path = os.path.join(run_output_dir_val, "final_model.png")
743
+ stl_path = os.path.join(run_output_dir_val, "final_model.stl")
744
+ txt_path = os.path.join(run_output_dir_val, "swap_instructions.txt")
745
+ hfp_path = os.path.join(run_output_dir_val, "project_file.hfp")
746
+
747
+ out_png = png_path if os.path.exists(png_path) else None
748
+ out_stl = stl_path if os.path.exists(stl_path) else None
749
+ out_txt = txt_path if os.path.exists(txt_path) else None
750
+ out_hfp = hfp_path if os.path.exists(hfp_path) else None
751
+
752
+ if out_png is None:
753
+ log_output += "\nWarning: final_model.png not found in output."
754
+
755
+ yield (
756
+ log_output, # progress_output
757
+ out_png, # final_image_preview
758
+ gr.update(
759
+ value=zip_path, visible=True, interactive=True
760
+ ), # ### ZIP PATCH: download_zip
761
+ )
762
+
763
+ run_inputs = [filament_df_state, input_image_component] + [
764
+ accordion_params_dict[name] for name in accordion_params_ordered_names
765
+ ]
766
+ run_outputs = [
767
+ progress_output,
768
+ final_image_preview,
769
+ download_zip, # ### ZIP PATCH: only three outputs now
770
+ ]
771
+
772
+ run_button.click(execute_autoforge_script, inputs=run_inputs, outputs=run_outputs)
773
+
774
+ css = """ #run_button_full_width { width: 100%; } """
775
+ if __name__ == "__main__":
776
+ if not os.path.exists(DEFAULT_MATERIALS_CSV):
777
+ print(f"Creating default filament file: {DEFAULT_MATERIALS_CSV}")
778
+ try:
779
+ initial_df.to_csv(DEFAULT_MATERIALS_CSV, index=False)
780
+ except Exception as e:
781
+ print(f"Could not write default {DEFAULT_MATERIALS_CSV}: {e}")
782
+ print("To run the UI, execute: python app.py") # Corrected to python app.py
783
+ demo.queue().launch(share=False)
default_materials.csv ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ Brand, Name, TD, Color, Owned
2
+ Generic,PLA Black,1.0,#000000,true
3
+ Generic,PLA Grey,1.0,#808080,true
4
+ Generic,PLA White,1.0,#FFFFFF,true
requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ autoforge
2
+ gradio