autoforge 1.6.0__tar.gz → 1.6.2__tar.gz

{autoforge-1.6.0 → autoforge-1.6.2}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: AutoForge
-Version: 1.6.0
+Version: 1.6.2
 Summary: AutoForge is a Python tool for generating 3D printed layered models from an input image.
 Author-email: Hendric Voss <hvoss@techfak.uni-bielefeld.de>
 License: CC BY-NC-SA 4.0
@@ -48,10 +48,6 @@ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at m
   </div>
 </div>
 
-If someone uses this program and 3d prints something with it, please let me know :) \
-I would love to see what you made!
-
-
 ## Features
 
 - **Image-to-Model Conversion**: Converts an input image into a layered model suitable for 3D printing.
@@ -63,8 +59,23 @@ I would love to see what you made!
 - **Live Visualization**: (Optional) Displays live composite images during the optimization process.
 - **Hueforge export**: Outputs a project file that can be opened with hueforge.
 
+## [Google Colab Notebook](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih)
+
+If you're new to Google Colab, follow these simple steps to run the notebook:
+
+1. **Open the Notebook:**  
+   Click [here](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih) to open the notebook in Google Colab.
+
+2. **Select GPU Runtime:**  
+   In Colab, go to the menu and click **Runtime → Change runtime type**. Under **Hardware accelerator**, choose **T4 GPU** (or any other available GPU type).
+
+3. **Run All Cells:**  
+   Once the notebook loads, click **Runtime → Run all** to execute the cells in sequence. Follow the on-screen prompts to configure the materials CSV, upload your image, and run Autoforge.
 
-## Installation
+4. **Review and Download:**  
+   After the process completes, you'll see a preview of the output image. Finally, the notebook will zip the output folder for you to download.
+
+## Manual Installation
 
 To install AutoForge, simply install the current version from PyPI:
 ```bash
@@ -81,13 +92,18 @@ The script is run from the command line and accepts several arguments. Below is
 > **Note:** You will need [Hueforge](https://shop.thehueforge.com/) installed to export your filament CSV.  
 > To get your CSV file, simply go to the "Filaments" menu in Hueforge, click the export button, select your filaments, and export them as a CSV file.
 
-> 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.
 
 ```bash
-autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv --output_folder outputs 
+autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv 
+```
+
+> If you want to limit the amount of colors the program can use, you can set these as command line arguments. \
+> For Example: 8 colors and a maximum of 20 swaps:
+
+```bash
+autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv --pruning_max_colors 8 --pruning_max_swaps 20
 ```
 
-### Command Line Arguments
 
 ### Command Line Arguments
 
@@ -99,10 +115,10 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
 
 - `--output_folder` Folder where output files will be saved (default: `output`).
 - `--iterations` Number of optimization iterations (default: 5000).
-- `--learning_rate` Learning rate for optimization (default: 1e-2).
+- `--learning_rate` Learning rate for optimization (default: 0.015).
 - `--layer_height` Layer thickness in millimeters (default: 0.04).
-- `--max_layers` Maximum number of layers (default: 50).  
-  **Note:** This is about 2mm + the background height
+- `--max_layers` Maximum number of layers (default: 75).  
+  **Note:** This is about 3mm + the background height
 - `--background_height` Height of the background in millimeters (default: 0.4).  
   **Note:** The background height must be divisible by the layer height.
 - `--background_color` Background color in hexadecimal format (default: `#000000` aka Black).  

{autoforge-1.6.0 → autoforge-1.6.2}/README.md

@@ -24,10 +24,6 @@ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at m
   </div>
 </div>
 
-If someone uses this program and 3d prints something with it, please let me know :) \
-I would love to see what you made!
-
-
 ## Features
 
 - **Image-to-Model Conversion**: Converts an input image into a layered model suitable for 3D printing.
@@ -39,8 +35,23 @@ I would love to see what you made!
 - **Live Visualization**: (Optional) Displays live composite images during the optimization process.
 - **Hueforge export**: Outputs a project file that can be opened with hueforge.
 
+## [Google Colab Notebook](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih)
+
+If you're new to Google Colab, follow these simple steps to run the notebook:
+
+1. **Open the Notebook:**  
+   Click [here](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih) to open the notebook in Google Colab.
+
+2. **Select GPU Runtime:**  
+   In Colab, go to the menu and click **Runtime → Change runtime type**. Under **Hardware accelerator**, choose **T4 GPU** (or any other available GPU type).
+
+3. **Run All Cells:**  
+   Once the notebook loads, click **Runtime → Run all** to execute the cells in sequence. Follow the on-screen prompts to configure the materials CSV, upload your image, and run Autoforge.
 
-## Installation
+4. **Review and Download:**  
+   After the process completes, you'll see a preview of the output image. Finally, the notebook will zip the output folder for you to download.
+
+## Manual Installation
 
 To install AutoForge, simply install the current version from PyPI:
 ```bash
@@ -57,13 +68,18 @@ The script is run from the command line and accepts several arguments. Below is
 > **Note:** You will need [Hueforge](https://shop.thehueforge.com/) installed to export your filament CSV.  
 > To get your CSV file, simply go to the "Filaments" menu in Hueforge, click the export button, select your filaments, and export them as a CSV file.
 
-> 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.
 
 ```bash
-autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv --output_folder outputs 
+autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv 
+```
+
+> If you want to limit the amount of colors the program can use, you can set these as command line arguments. \
+> For Example: 8 colors and a maximum of 20 swaps:
+
+```bash
+autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv --pruning_max_colors 8 --pruning_max_swaps 20
 ```
 
-### Command Line Arguments
 
 ### Command Line Arguments
 
@@ -75,10 +91,10 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
 
 - `--output_folder` Folder where output files will be saved (default: `output`).
 - `--iterations` Number of optimization iterations (default: 5000).
-- `--learning_rate` Learning rate for optimization (default: 1e-2).
+- `--learning_rate` Learning rate for optimization (default: 0.015).
 - `--layer_height` Layer thickness in millimeters (default: 0.04).
-- `--max_layers` Maximum number of layers (default: 50).  
-  **Note:** This is about 2mm + the background height
+- `--max_layers` Maximum number of layers (default: 75).  
+  **Note:** This is about 3mm + the background height
 - `--background_height` Height of the background in millimeters (default: 0.4).  
   **Note:** The background height must be divisible by the layer height.
 - `--background_color` Background color in hexadecimal format (default: `#000000` aka Black).  
@@ -178,4 +194,4 @@ Example Images: \
 <a href="https://www.vecteezy.com/free-photos/animal">Animal Stock photos by Vecteezy</a> 
 
 
-Happy printing!
+Happy printing!

{autoforge-1.6.0 → autoforge-1.6.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "AutoForge"
-version="1.6.0"
+version="1.6.2"
 description = "AutoForge is a Python tool for generating 3D printed layered models from an input image."
 authors = [{ name = "Hendric Voss", email = "hvoss@techfak.uni-bielefeld.de" }]
 readme = "README.md"

{autoforge-1.6.0 → autoforge-1.6.2}/src/AutoForge.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: AutoForge
-Version: 1.6.0
+Version: 1.6.2
 Summary: AutoForge is a Python tool for generating 3D printed layered models from an input image.
 Author-email: Hendric Voss <hvoss@techfak.uni-bielefeld.de>
 License: CC BY-NC-SA 4.0
@@ -48,10 +48,6 @@ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at m
   </div>
 </div>
 
-If someone uses this program and 3d prints something with it, please let me know :) \
-I would love to see what you made!
-
-
 ## Features
 
 - **Image-to-Model Conversion**: Converts an input image into a layered model suitable for 3D printing.
@@ -63,8 +59,23 @@ I would love to see what you made!
 - **Live Visualization**: (Optional) Displays live composite images during the optimization process.
 - **Hueforge export**: Outputs a project file that can be opened with hueforge.
 
+## [Google Colab Notebook](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih)
+
+If you're new to Google Colab, follow these simple steps to run the notebook:
+
+1. **Open the Notebook:**  
+   Click [here](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih) to open the notebook in Google Colab.
+
+2. **Select GPU Runtime:**  
+   In Colab, go to the menu and click **Runtime → Change runtime type**. Under **Hardware accelerator**, choose **T4 GPU** (or any other available GPU type).
+
+3. **Run All Cells:**  
+   Once the notebook loads, click **Runtime → Run all** to execute the cells in sequence. Follow the on-screen prompts to configure the materials CSV, upload your image, and run Autoforge.
 
-## Installation
+4. **Review and Download:**  
+   After the process completes, you'll see a preview of the output image. Finally, the notebook will zip the output folder for you to download.
+
+## Manual Installation
 
 To install AutoForge, simply install the current version from PyPI:
 ```bash
@@ -81,13 +92,18 @@ The script is run from the command line and accepts several arguments. Below is
 > **Note:** You will need [Hueforge](https://shop.thehueforge.com/) installed to export your filament CSV.  
 > To get your CSV file, simply go to the "Filaments" menu in Hueforge, click the export button, select your filaments, and export them as a CSV file.
 
-> 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.
 
 ```bash
-autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv --output_folder outputs 
+autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv 
+```
+
+> If you want to limit the amount of colors the program can use, you can set these as command line arguments. \
+> For Example: 8 colors and a maximum of 20 swaps:
+
+```bash
+autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv --pruning_max_colors 8 --pruning_max_swaps 20
 ```
 
-### Command Line Arguments
 
 ### Command Line Arguments
 
@@ -99,10 +115,10 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
 
 - `--output_folder` Folder where output files will be saved (default: `output`).
 - `--iterations` Number of optimization iterations (default: 5000).
-- `--learning_rate` Learning rate for optimization (default: 1e-2).
+- `--learning_rate` Learning rate for optimization (default: 0.015).
 - `--layer_height` Layer thickness in millimeters (default: 0.04).
-- `--max_layers` Maximum number of layers (default: 50).  
-  **Note:** This is about 2mm + the background height
+- `--max_layers` Maximum number of layers (default: 75).  
+  **Note:** This is about 3mm + the background height
 - `--background_height` Height of the background in millimeters (default: 0.4).  
   **Note:** The background height must be divisible by the layer height.
 - `--background_color` Background color in hexadecimal format (default: `#000000` aka Black).  

{autoforge-1.6.0 → autoforge-1.6.2}/src/autoforge/Init_comparer/auto_forge_init_check.py

@@ -24,15 +24,20 @@ class Config:
     csv_file = "default_materials.csv"  # Path to CSV file with material data
     output_folder = "output"
     iterations = 2000
-    learning_rate = 1e-2
+    learning_rate = 0.015
+
+    warmup_fraction = 0.1 #[0.0, 1.0]
+    height_logits_learning_start_fraction = 0.1 #[0.0, 1.0] but should be less than height_logits_learning_full_fraction
+    height_logits_learning_full_fraction = 0.5 # [0.0, 1.0] but should be more than height_logits_learning_start_fraction
+    init_tau = 1.0 # [0.0, 1.0] needs to be more than final_tau
+    final_tau = 0.01 # [0.0, 1.0] needs to be less than init_tau
+
     layer_height = 0.04
     max_layers = 75
     min_layers = 0
     background_height = 0.4
     background_color = "#000000"
     output_size = 128
-    init_tau = 1.0
-    final_tau = 0.01
     visualize = False
     stl_output_size = 200
     perform_pruning = True
@@ -52,12 +57,19 @@ class Config:
     tensorboard = False
 
 
-def main(input_image, csv_file, init_method, cluster_layers, lab_space):
-    # Create config object using default values
+def main(input_image, csv_file, warmup, h_start, h_full, tau_init, tau_final):
+    # Create config object using default values and override with given hyperparameters.
     args = Config()
     args.input_image = input_image
     args.csv_file = csv_file
 
+    # Set the hyperparameters for the grid search.
+    args.warmup_fraction = warmup
+    args.height_logits_learning_start_fraction = h_start
+    args.height_logits_learning_full_fraction = h_full
+    args.init_tau = tau_init
+    args.final_tau = tau_final
+
     if torch.cuda.is_available():
         device = torch.device("cuda")
     elif args.mps and torch.backends.mps.is_available():
@@ -128,9 +140,9 @@ def main(input_image, csv_file, init_method, cluster_layers, lab_space):
         args.layer_height,
         bgr_tuple,
         random_seed=random_seed,
-        init_method=init_method,
-        cluster_layers=cluster_layers,
-        lab_space=lab_space,
+        init_method="kmeans",
+        cluster_layers=20,
+        lab_space=False,
         num_threads=8,
     )
 
@@ -186,47 +198,41 @@ def main(input_image, csv_file, init_method, cluster_layers, lab_space):
     return mse_loss.item()
 
 
-def main_suppressed(input_image, csv_file, init_method, cluster_layers, lab_space):
+def main_suppressed(input_image, csv_file, warmup, h_start, h_full, tau_init, tau_final):
     with open(os.devnull, "w") as fnull:
         with redirect_stdout(fnull), redirect_stderr(fnull):
-            result = main(input_image, csv_file, init_method, cluster_layers, lab_space)
+            result = main(input_image, csv_file, warmup, h_start, h_full, tau_init, tau_final)
     return result
 
 
 if __name__ == "__main__":
     folder = "../../../images/test_images/"
     csv_file = "../../../bambulab.csv"
-    images = [folder + "/" + img for img in os.listdir(folder) if img.endswith(".jpg")]
-    parallel_limit = os.cpu_count()
-    methods = [
-        "kmeans",
-        "quantize_median",
-        "quantize_maxcoverage",
-        "quantize_fastoctree",
-    ]
-    max_layers = 75
-    cluster_layers = [
-        max_layers // 4,
-        max_layers // 2,
-        max_layers,
-        max_layers * 2,
-        max_layers * 4,
-    ]
-    use_lab_space = [True, False]
-
-    # test every permutation using itertools
+    images = [os.path.join(folder, img) for img in os.listdir(folder) if img.endswith(".jpg")]
+    fixed_lr = 1e-1
+    parallel_limit = 10
+
+    # Define grid search parameters
+    warmup_values = np.linspace(0, 1, 5)  # 5 values for warmup_fraction
+    height_values = np.linspace(0, 1, 5)  # 5 values for height fractions
+    tau_values = np.linspace(0, 1, 5)     # 5 values for tau
+
+    # Build valid pairs for height fractions (start < full)
+    height_pairs = [(hs, hf) for hs in height_values for hf in height_values if hs < hf]
+    # Build valid pairs for tau (init_tau > final_tau)
+    tau_pairs = [(ti, tf) for ti in tau_values for tf in tau_values if ti > tf]
+
+    # Create grid: warmup x height_pairs x tau_pairs = 5 * 10 * 10 = 500 combinations
     from itertools import product
+    hyperparam_grid = list(product(warmup_values, height_pairs, tau_pairs))
+    print(f"Total hyperparameter combinations: {len(hyperparam_grid)}")
 
     out_dict = {}
-    do_list = list(product(methods, cluster_layers, use_lab_space))
-    shuffle(do_list)
-
-    for i, (method, cluster, lab) in enumerate(do_list):
+    # Loop over each hyperparameter combination
+    for idx, (warmup, (h_start, h_full), (tau_init, tau_final)) in enumerate(hyperparam_grid):
         try:
-            out_dict_str = f"{method}_{cluster}_{lab}"
-            print(
-                f"Running {method} with {cluster} clusters and lab={lab}, {i + 1}/{len(do_list)}"
-            )
+            out_dict_str = f"w={warmup}_height_start={h_start}_height_full={h_full}_tau_init={tau_init}_tau_final={tau_final}"
+            print(f"Running {idx+1}/{len(hyperparam_grid)}: {out_dict_str}")
             exec = ProcessPoolExecutor(max_workers=parallel_limit)
             tlist = []
             for img in images:
@@ -236,15 +242,19 @@ if __name__ == "__main__":
                             main_suppressed,
                             img,
                             csv_file,
-                            method,
-                            cluster,
-                            lab,
+                            warmup,
+                            h_start,
+                            h_full,
+                            tau_init,
+                            tau_final,
                         )
                     )
             for t in tqdm(concurrent.futures.as_completed(tlist), total=len(tlist)):
                 result_list = out_dict.get(out_dict_str, [])
                 result_list.append(t.result())
                 out_dict[out_dict_str] = result_list
+
+            exec.shutdown()
             # save out_dict as json
             import json
 

{autoforge-1.6.0 → autoforge-1.6.2}/src/autoforge/Init_comparer/plot_results.py

@@ -14,10 +14,12 @@ for key, values in data.items():
     arr = np.array(values)
     mean_val = np.mean(arr)
     std_val = np.std(arr)
+    #key is formatted as lr=float, we need to extract the float value and round it to 4 decimal values
+    key = round(float(key.split("=")[1]), 5)
     stats.append((key, mean_val, std_val, arr))
 
 # Sort the list of tuples based on the mean (ascending order)
-stats.sort(key=lambda x: x[1])
+#stats.sort(key=lambda x: x[1])
 
 # Extract sorted lists of keys, means, stds, and store the sorted arrays in a dictionary
 sorted_keys = [item[0] for item in stats]
@@ -53,7 +55,6 @@ fig.update_layout(
     barmode="group",
 )
 
-
 # Function to generate significance stars based on the p-value.
 def significance_label(p):
     if p < 0.001:
@@ -65,50 +66,57 @@ def significance_label(p):
     else:
         return ""
 
+# First, compute all significant pairwise comparisons
+significant_pairs = []  # will hold tuples of (p_value, i, j)
+for i in range(n_keys):
+    for j in range(i + 1, n_keys):
+        group1 = data_sorted[sorted_keys[i]]
+        group2 = data_sorted[sorted_keys[j]]
+        stat, p_value = ttest_ind(group1, group2)
+        if p_value < 0.05:
+            significant_pairs.append((p_value, i, j))
+
+# Sort pairs by p-value (most significant first)
+significant_pairs.sort(key=lambda x: x[0])
+
+# To ensure each bar is annotated only once, track used indices.
+used_bars = set()
+base_margin = max(means) * 0.05  # margin above the error bars
+
+significant_pairs = []
+for p_value, i, j in significant_pairs:
+    # If either bar is already annotated, skip this pair.
+    if i in used_bars or j in used_bars:
+        continue
+
+    # Compute the base y-level from the higher error bar of the two.
+    base_y = max(means[i] + stds[i], means[j] + stds[j])
+    line_y = base_y + base_margin
+
+    # Mark these bars as used.
+    used_bars.update([i, j])
+
+    # Draw horizontal significance line between the two bars.
+    fig.add_shape(
+        type="line",
+        x0=x_positions[i],
+        x1=x_positions[j],
+        y0=line_y,
+        y1=line_y,
+        line=dict(color="black", width=1),
+        xref="x",
+        yref="y",
+    )
+    # Place the significance annotation (stars) above the line.
+    fig.add_annotation(
+        x=(x_positions[i] + x_positions[j]) / 2,
+        y=line_y + base_margin * 0.2,
+        text=significance_label(p_value),
+        showarrow=False,
+        font=dict(size=12),
+    )
+
 
-# To avoid overlapping significance annotations when bars share a common index,
-# keep track of the highest annotation drawn for each bar.
-annotation_levels = {i: 0 for i in range(n_keys)}
-# Base margin added above the highest error bar to start the significance line.
-base_margin = max(means) * 0.05  # Adjust as needed
-
-# Loop over adjacent pairs in the sorted order and perform t-tests.
-for i in range(n_keys - 1):
-    group1 = data_sorted[sorted_keys[i]]
-    group2 = data_sorted[sorted_keys[i + 1]]
-    stat, p_value = ttest_ind(group1, group2)
-
-    # Only annotate if the difference is significant (p < 0.05)
-    if p_value < 0.05:
-        # Determine the base y-level from the taller of the two error bars.
-        base_y = max(means[i] + stds[i], means[i + 1] + stds[i + 1])
-        # Ensure new annotations are drawn above any previous ones.
-        current_level = max(annotation_levels[i], annotation_levels[i + 1])
-        line_y = base_y + base_margin + current_level
-
-        # Update annotation levels for these two bars.
-        annotation_levels[i] = line_y
-        annotation_levels[i + 1] = line_y
-
-        # Add a horizontal line between the two bars.
-        fig.add_shape(
-            type="line",
-            x0=x_positions[i],
-            x1=x_positions[i + 1],
-            y0=line_y,
-            y1=line_y,
-            line=dict(color="black", width=1),
-            xref="x",
-            yref="y",
-        )
-        # Add the significance stars annotation above the line.
-        fig.add_annotation(
-            x=(x_positions[i] + x_positions[i + 1]) / 2,
-            y=line_y + base_margin * 0.2,
-            text=significance_label(p_value),
-            showarrow=False,
-            font=dict(size=12),
-        )
 fig.update_layout(
     autosize=False,
     width=1920,

{autoforge-1.6.0 → autoforge-1.6.2}/src/autoforge/Loss/LossFunctions.py

@@ -50,46 +50,42 @@ def compute_loss(
     add_penalty_loss: float = 0.0,
 ) -> torch.Tensor:
     """
-    Combined MSE + Perceptual + penalty losses.
+    Combined MSE + Perceptual + penalty losses with patch-based smoothness.
     """
-    # MSE
+    # MSE Loss
     comp_mse = srgb_to_lab(comp)
-    # we slightly increase saturation of our target image to make the color matching more robust
-    # target = increase_saturation(target, 0.1)
     target_mse = srgb_to_lab(target)
     mse_loss = F.huber_loss(comp_mse, target_mse)
 
     if pixel_height_logits is not None:
+        # Existing neighbor-based smoothness loss:
         target_gray = target.mean(dim=2)  # shape becomes [H, W]
-
-        # Compute weights using the grayscale image
-        weight_x = torch.exp(
-            -torch.abs(target_gray[:, 1:] - target_gray[:, :-1])
-        )  # shape [H, W-1]
-        weight_y = torch.exp(
-            -torch.abs(target_gray[1:, :] - target_gray[:-1, :])
-        )  # shape [H-1, W]
-
+        weight_x = torch.exp(-torch.abs(target_gray[:, 1:] - target_gray[:, :-1]))
+        weight_y = torch.exp(-torch.abs(target_gray[1:, :] - target_gray[:-1, :]))
         weight_x = torch.clamp(weight_x, 0.5, 1.0)
         weight_y = torch.clamp(weight_y, 0.5, 1.0)
-
-        # Compute differences in the depth map
-        dx = torch.abs(
-            pixel_height_logits[:, 1:] - pixel_height_logits[:, :-1]
-        )  # shape [H, W-1]
-        dy = torch.abs(
-            pixel_height_logits[1:, :] - pixel_height_logits[:-1, :]
-        )  # shape [H-1, W]
-
-        # provide a target tensor (zeros) for the huber loss.
+        dx = torch.abs(pixel_height_logits[:, 1:] - pixel_height_logits[:, :-1])
+        dy = torch.abs(pixel_height_logits[1:, :] - pixel_height_logits[:-1, :])
         loss_dx = torch.mean(F.huber_loss(dx * weight_x, torch.zeros_like(dx)))
         loss_dy = torch.mean(F.huber_loss(dy * weight_y, torch.zeros_like(dy)))
-        # loss_dx = torch.mean(F.huber_loss(dx, torch.zeros_like(dx)))
-        # loss_dy = torch.mean(F.huber_loss(dy, torch.zeros_like(dy)))
+        smoothness_loss = (loss_dx + loss_dy) * (10 * add_penalty_loss)
 
-        smoothness_loss = (loss_dx + loss_dy) * (20 * add_penalty_loss)
-        # print(f"mse_loss: {mse_loss}, smoothness_loss: {smoothness_loss}")
-        total_loss = mse_loss + smoothness_loss
+        # Additional patch-based smoothness loss (using a 3x3 Laplacian):
+        laplacian_kernel = (
+            torch.tensor(
+                [[0, 1, 0], [1, -4, 1], [0, 1, 0]],
+                dtype=pixel_height_logits.dtype,
+                device=pixel_height_logits.device,
+            )
+            .unsqueeze(0)
+            .unsqueeze(0)
+        )
+        height_map = pixel_height_logits.unsqueeze(0).unsqueeze(0)
+        laplacian_output = F.conv2d(height_map, laplacian_kernel, padding=1)
+        patch_smooth_loss = (
+            F.huber_loss(laplacian_output, torch.zeros_like(laplacian_output)) * 10
+        )
+        total_loss = mse_loss + smoothness_loss + add_penalty_loss * patch_smooth_loss
     else:
         total_loss = mse_loss
     return total_loss

{autoforge-1.6.0 → autoforge-1.6.2}/src/autoforge/Modules/Optimizer.py

@@ -95,7 +95,7 @@ class FilamentOptimizer:
 
         # Tau schedule
         self.num_steps_done = 0
-        self.warmup_steps = args.iterations // 10
+        self.warmup_steps = min(args.iterations-1,args.warmup_fraction * args.iterations)
         self.decay_rate = (self.init_tau - self.final_tau) / (
             args.iterations - self.warmup_steps
         )
@@ -197,6 +197,18 @@ class FilamentOptimizer:
                 / (self.args.iterations // 5),
             )
 
+        #start_fraction = self.args.height_logits_learning_start_fraction
+        #full_fraction = self.args.height_logits_learning_full_fraction
+
+        # total_iters = self.args.iterations
+        # current_step = self.num_steps_done
+        # if current_step < start_fraction * total_iters:
+        #     scaling = 0.0
+        # elif current_step < full_fraction * total_iters:
+        #     scaling = (current_step - start_fraction * total_iters) / ((full_fraction - start_fraction) * total_iters)
+        # else:
+        #     scaling = 1.0
+
         loss = loss_fn(
             self.params,
             target=self.target,
@@ -208,10 +220,17 @@ class FilamentOptimizer:
             material_TDs=self.material_TDs,
             background=self.background,
             perception_loss_module=self.perception_loss_module,
-            add_penalty_loss=penalty_loss,
+            add_penalty_loss=1.0,
         )
 
         loss.backward()
+
+        # We scale the gradients for the height logits by a factor of 'scaling' to force no updates in the beginning
+        # and allow stronger and stronger updates as we go along.
+        if self.params["pixel_height_logits"].grad is not None:
+            self.params["pixel_height_logits"].grad.mul_(1e-8)
+
+
         self.optimizer.step()
 
         self.num_steps_done += 1

{autoforge-1.6.0 → autoforge-1.6.2}/src/autoforge/auto_forge.py

@@ -45,10 +45,38 @@ def main():
         "--iterations", type=int, default=5000, help="Number of optimization iterations"
     )
 
+    parser.add_argument(
+        "--warmup_fraction",
+        type=float,
+        default=1.0,
+        help="Fraction of iterations for keeping the tau at the initial value",
+    )
+    #
+    # parser.add_argument(
+    #     "--height_logits_learning_start_fraction", type=float, default=0.75, help="Fraction of iterations at which we start to learn the height map"
+    # )
+    #
+    # parser.add_argument(
+    #     "--height_logits_learning_full_fraction", type=float, default=1.0, help="Fraction of iterations at which the height map learning is fully enabled"
+    # )
+
+    parser.add_argument(
+        "--init_tau",
+        type=float,
+        default=1.0,
+        help="Initial tau value for Gumbel-Softmax",
+    )
+    parser.add_argument(
+        "--final_tau",
+        type=float,
+        default=0.01,
+        help="Final tau value for Gumbel-Softmax",
+    )
+
     parser.add_argument(
         "--learning_rate",
         type=float,
-        default=1e-2,
+        default=0.015,
         help="Learning rate for optimization",
     )
     parser.add_argument(
@@ -73,19 +101,6 @@ def main():
         "--background_color", type=str, default="#000000", help="Background color"
     )
 
-    parser.add_argument(
-        "--init_tau",
-        type=float,
-        default=1.0,
-        help="Initial tau value for Gumbel-Softmax",
-    )
-    parser.add_argument(
-        "--final_tau",
-        type=float,
-        default=0.01,
-        help="Final tau value for Gumbel-Softmax",
-    )
-
     parser.add_argument(
         "--visualize",
         action="store_true",
@@ -289,7 +304,7 @@ def main():
             random_seed=random_seed,
             num_threads=16,
             init_method="kmeans",
-            cluster_layers=300,
+            cluster_layers=18,
             lab_space=True,
         )