autoforge 1.6.2__tar.gz → 1.6.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: AutoForge
-Version: 1.6.2
+Version: 1.6.4
 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
@@ -21,6 +21,8 @@ Requires-Dist: tensorboard
 Requires-Dist: pillow>=11.1.0
 Requires-Dist: shapely>=2.0.7
 Requires-Dist: trimesh>=4.6.4
+Requires-Dist: torchcam>=0.4.0
+Requires-Dist: joblib>=1.4.2
 
 # AutoForge
 
@@ -61,7 +63,7 @@ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at m
 
 ## [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:
+If you're new to Google Colab, follow these simple steps to run autoforge and get your output:
 
 1. **Open the Notebook:**  
    Click [here](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih) to open the notebook in Google Colab.

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

@@ -37,7 +37,7 @@ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at m
 
 ## [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:
+If you're new to Google Colab, follow these simple steps to run autoforge and get your output:
 
 1. **Open the Notebook:**  
    Click [here](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih) to open the notebook in Google Colab.

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "AutoForge"
-version="1.6.2"
+version="1.6.4"
 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"
@@ -25,7 +25,9 @@ dependencies = [
   "tensorboard",
   "pillow>=11.1.0",
   "shapely>=2.0.7",
-  "trimesh>=4.6.4"
+  "trimesh>=4.6.4",
+  "torchcam>=0.4.0",
+  "joblib>=1.4.2"
 ]
 
 [tool.setuptools]

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: AutoForge
-Version: 1.6.2
+Version: 1.6.4
 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
@@ -21,6 +21,8 @@ Requires-Dist: tensorboard
 Requires-Dist: pillow>=11.1.0
 Requires-Dist: shapely>=2.0.7
 Requires-Dist: trimesh>=4.6.4
+Requires-Dist: torchcam>=0.4.0
+Requires-Dist: joblib>=1.4.2
 
 # AutoForge
 
@@ -61,7 +63,7 @@ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at m
 
 ## [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:
+If you're new to Google Colab, follow these simple steps to run autoforge and get your output:
 
 1. **Open the Notebook:**  
    Click [here](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih) to open the notebook in Google Colab.

{autoforge-1.6.2 → autoforge-1.6.4}/src/AutoForge.egg-info/requires.txt

@@ -13,3 +13,5 @@ tensorboard
 pillow>=11.1.0
 shapely>=2.0.7
 trimesh>=4.6.4
+torchcam>=0.4.0
+joblib>=1.4.2

{autoforge-1.6.2 → autoforge-1.6.4}/src/autoforge/Helper/Heightmaps/ChristofidesHeightMap.py

@@ -1,15 +1,13 @@
-import concurrent.futures
 import os
 import random
-from concurrent.futures import ThreadPoolExecutor
 
 import numpy as np
 from PIL import Image
 from PIL.Image import Quantize
+from joblib import Parallel, delayed
 from scipy.spatial.distance import cdist
 from skimage.color import rgb2lab
 from sklearn.cluster import MiniBatchKMeans
-from tqdm import tqdm
 
 
 def build_distance_matrix(labs, nodes):
@@ -240,25 +238,28 @@ def create_mapping(final_ordering, labs, all_labels):
 
 def tsp_order_christofides_path(nodes, labs, bg, fg):
     """
-    nodes: list of cluster indices (must include bg and fg)
+    Ensure that the background and foreground nodes are always in the TSP cycle.
+    nodes: list of cluster indices (ideally including bg and fg)
     labs: Lab-space coordinates (indexed by cluster index)
     bg, fg: background and foreground cluster indices
-
     Returns an ordering (list of cluster indices) that contains all nodes,
     starts with bg and ends with fg.
     """
+    # Guarantee that bg and fg are included in nodes.
+    nodes = list(set(nodes) | {bg, fg})
+
     artificial = -1
-    augmented_nodes = nodes + [artificial]
     LARGE = 1e6
+    n = len(nodes)
 
-    # Precompute the distance matrix for original nodes.
+    # Precompute the distance matrix for the given nodes.
     D = build_distance_matrix(labs, nodes)
-    n = len(nodes)
 
-    # Build an augmented matrix of size (n+1)x(n+1).
+    # Build an augmented (n+1)x(n+1) matrix.
     aug_mat = np.zeros((n + 1, n + 1))
     aug_mat[:n, :n] = D
-    # For the artificial node, set cost = 0 for bg/fg, LARGE otherwise.
+
+    # For the artificial node, set cost = 0 for bg/fg and LARGE for others.
     for i, u in enumerate(nodes):
         if u in {bg, fg}:
             aug_mat[i, n] = 0.0
@@ -268,32 +269,33 @@ def tsp_order_christofides_path(nodes, labs, bg, fg):
             aug_mat[n, i] = LARGE
     aug_mat[n, n] = 0.0
 
-    # Convert the augmented matrix into a graph dictionary.
+    # Create augmented nodes list.
     aug_nodes = nodes + [artificial]
     graph = matrix_to_graph(aug_mat, aug_nodes)
 
-    # Run Christofides on the augmented graph.
+    # Run Christofides algorithm on the augmented graph.
     cycle = christofides_tsp(graph)
-    # Remove the artificial node.
+    # Remove the artificial node if present.
     cycle = [node for node in cycle if node != artificial]
 
+    # Ensure bg and fg are in the cycle.
+    if bg not in cycle:
+        cycle.insert(0, bg)
+    if fg not in cycle:
+        cycle.append(fg)
+
     # Rotate the cycle so that bg is first.
-    if bg in cycle:
+    if cycle[0] != bg:
         idx = cycle.index(bg)
         cycle = cycle[idx:] + cycle[:idx]
-    else:
-        raise ValueError("Background node missing from cycle")
 
     # Force fg to be the last element.
-    if fg not in cycle:
-        raise ValueError("Foreground node missing from cycle")
     if cycle[-1] != fg:
         cycle.remove(fg)
         cycle.append(fg)
 
-    # --- NEW: Check if reversing the internal order yields a lower ordering metric ---
+    # Optionally, check if reversing the internal order improves the ordering metric.
     if len(cycle) > 2:
-        # Reverse the ordering of the nodes between bg and fg.
         reversed_cycle = [cycle[0]] + cycle[1:-1][::-1] + [cycle[-1]]
         if compute_ordering_metric(reversed_cycle, labs) < compute_ordering_metric(
             cycle, labs
@@ -449,14 +451,7 @@ def init_height_map(
     nodes = unique_clusters
 
     # Get the ordering via TSP ordering function.
-    for i in range(50):
-        try:
-            final_ordering = tsp_order_christofides_path(nodes, labs, bg_cluster, fg_cluster)
-            break
-        except Exception:
-            print(f"Failed to find ordering, retrying {i+1}/10")
-            if i == 49:
-                raise ValueError("Failed to find ordering after 50 retries")
+    final_ordering = tsp_order_christofides_path(nodes, labs, bg_cluster, fg_cluster)
 
     # Optionally prune out outliers.
     final_ordering = prune_ordering(
@@ -474,7 +469,7 @@ def init_height_map(
 def run_init_threads(
     target,
     max_layers,
-    h,  # unused here but preserved for API compatibility
+    h,  # unused but preserved for API compatibility
     background_tuple,
     eps=1e-6,
     random_seed=None,
@@ -485,33 +480,33 @@ def run_init_threads(
 ):
     if random_seed is None:
         random_seed = np.random.randint(1e6)
-    exec = ThreadPoolExecutor(max_workers=os.cpu_count())
-    futures = []
-    for i in range(num_threads):
-        futures.append(
-            exec.submit(
-                init_height_map,
-                target,
-                max_layers,
-                h,
-                background_tuple,
-                eps,
-                random_seed + i,
-                init_method=init_method,
-                cluster_layers=cluster_layers,
-                lab_space=lab_space,
-            )
-        )
-    results = [
-        f.result()
-        for f in tqdm(
-            concurrent.futures.as_completed(futures),
-            desc="Running Height Map Initialization Threads",
-            total=num_threads,
+
+    tasks = [
+        delayed(init_height_map)(
+            target,
+            max_layers,
+            h,
+            background_tuple,
+            eps,
+            random_seed + i,
+            init_method=init_method,
+            cluster_layers=cluster_layers,
+            lab_space=lab_space,
         )
+        for i in range(num_threads)
     ]
-    # print(f"metric values: {[r[1] for r in results]}")
-    # get lowest ordering metric
-    best_result = min(results, key=lambda x: x[1])
 
+    # Execute tasks in parallel; adjust n_jobs to match your available cores
+    results = Parallel(n_jobs=os.cpu_count(), verbose=10)(tasks)
+
+    metrics = [r[1] for r in results]
+    mean_metric = np.mean(metrics)
+    std_metric = np.std(metrics)
+    min_metric = np.min(metrics)
+    max_metric = np.max(metrics)
+    print(
+        f"mean: {mean_metric}, std: {std_metric}, min: {min_metric}, max: {max_metric}"
+    )
+    print(f"Choosing best ordering with metric: {min_metric}")
+    best_result = min(results, key=lambda x: x[1])
     return best_result[0]

autoforge-1.6.4/src/autoforge/Loss/LossFunctions.py

@@ -0,0 +1,109 @@
+import torch
+import torch.nn.functional as F
+
+from autoforge.Helper.ImageHelper import srgb_to_lab
+from autoforge.Helper.OptimizerHelper import composite_image_cont
+
+
+def loss_fn(
+    params: dict,
+    target: torch.Tensor,
+    tau_height: float,
+    tau_global: float,
+    h: float,
+    max_layers: int,
+    material_colors: torch.Tensor,
+    material_TDs: torch.Tensor,
+    background: torch.Tensor,
+    add_penalty_loss: float = 0.0,
+    focus_map: torch.Tensor = None,
+    focus_strength: float = 10.0,
+) -> torch.Tensor:
+    """
+    Full forward pass for continuous assignment:
+    composite, then compute unified loss on (global_logits).
+    """
+    comp = composite_image_cont(
+        params["pixel_height_logits"],
+        params["global_logits"],
+        tau_height,
+        tau_global,
+        h,
+        max_layers,
+        material_colors,
+        material_TDs,
+        background,
+    )
+    return compute_loss(
+        comp=comp,
+        target=target,
+        pixel_height_logits=params["pixel_height_logits"],
+        tau_height=tau_height,
+        add_penalty_loss=add_penalty_loss,
+        focus_map=focus_map,
+        focus_strength=focus_strength,
+    )
+
+
+def compute_loss(
+    comp: torch.Tensor,
+    target: torch.Tensor,
+    pixel_height_logits: torch.Tensor = None,
+    tau_height: float = 1.0,
+    add_penalty_loss: float = 0.0,
+    focus_map: torch.Tensor = None,
+    focus_strength: float = 10.0,
+) -> torch.Tensor:
+    comp_lab = srgb_to_lab(comp)
+    target_lab = srgb_to_lab(target)
+    loss1 = F.huber_loss(comp_lab, target_lab, delta=0.1)
+
+    return loss1
+
+    # if focus_map is not None and focus_strength > 0.0:
+    #     # Expand focus_map to [H, W, 1] to match the color channels.
+    #     focus_map_exp = focus_map.unsqueeze(-1)
+    #     base_loss = F.huber_loss(comp_mse, target_mse, reduction="none")
+    #     # we need to sum up to a maximum of one, so normalize based on strength
+    #     min_strength = 1 / focus_strength
+    #     max_strength = 1 - min_strength
+    #     # focus map is normalized between 0 and 1
+    #     normalized_focus = min_strength + focus_map_exp * max_strength
+    #     weighted_loss = base_loss * normalized_focus
+    #
+    #     mse_loss = weighted_loss.mean()
+    # else:
+    #
+
+    # loss_color = global_color_loss(comp_lab, target_lab)
+
+    # if pixel_height_logits is not None:
+    #     # Existing neighbor-based smoothness loss:
+    #     target_gray = target.mean(dim=2)  # shape becomes [H, 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)
+    #     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)))
+    #     smoothness_loss = (loss_dx + loss_dy) * add_penalty_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)
+    #     )
+    #     total_loss = mse_loss + smoothness_loss + add_penalty_loss * patch_smooth_loss
+    # else:

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

@@ -5,6 +5,12 @@ import matplotlib.pyplot as plt
 import numpy as np
 import torch
 from torch.utils.tensorboard import SummaryWriter
+from torchcam.methods import SmoothGradCAMpp
+from torchcam.utils import overlay_mask
+from torchvision import models
+from torchvision.models import ResNet18_Weights
+from torchvision.transforms._functional_tensor import resize, normalize
+from torchvision.transforms.functional import to_pil_image
 from tqdm import tqdm
 
 from autoforge.Helper.CAdamW import CAdamW
@@ -95,7 +101,9 @@ class FilamentOptimizer:
 
         # Tau schedule
         self.num_steps_done = 0
-        self.warmup_steps = min(args.iterations-1,args.warmup_fraction * args.iterations)
+        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
         )
@@ -106,6 +114,49 @@ class FilamentOptimizer:
             lr=self.learning_rate,
         )
 
+        # Setup the focus map component
+        # Load a state-of-the-art ResNet50 model with pretrained weights.
+        cam_model = models.resnet18(weights=ResNet18_Weights.DEFAULT).to(
+            self.target.device
+        )
+        cam_model.eval()
+        # Use ScoreCAM from torchcam on the 'layer4' block.
+        cam_extractor = SmoothGradCAMpp(cam_model, target_layer="layer4")
+        # Preprocess the target image for ResNet50.
+        target_preprocessed = normalize(
+            resize(self.target.permute(2, 0, 1), (224, 224)) / 255.0,
+            [0.485, 0.456, 0.406],
+            [0.229, 0.224, 0.225],
+        )
+        focus_maps = []
+        for _ in tqdm(range(1), desc="Creating loss function attention map..."):
+            out = cam_model(target_preprocessed.unsqueeze(0))
+            # ScoreCAM returns a dict with keys corresponding to layers.
+            focus_map = cam_extractor(out.squeeze(0).argmax().item(), out)[0]
+            # Normalize the focus map to [0, 1]
+            focus_map = (focus_map - focus_map.min()) / (
+                focus_map.max() - focus_map.min() + 1e-8
+            )
+            focus_maps.append(focus_map)
+        # get maximum over all maps
+        self.focus_map = torch.stack(focus_maps)
+        self.focus_map = self.focus_map.max(dim=0)[0]
+        # normalize again
+        self.focus_map = (self.focus_map - self.focus_map.min()) / (
+            self.focus_map.max() - self.focus_map.min() + 1e-8
+        )
+
+        # resize back to target size, current dim is [1,h,w]
+        self.focus_map = (
+            torch.nn.functional.interpolate(
+                self.focus_map.unsqueeze(0),
+                (self.H, self.W),
+                mode="bilinear",
+            )
+            .squeeze(0)
+            .squeeze(0)
+        )
+
         # Setup best discrete solution tracking
         self.best_discrete_loss = float("inf")
         self.best_params = None
@@ -137,17 +188,24 @@ class FilamentOptimizer:
             self.depth_map_ax = self.ax[1, 0].imshow(
                 np.zeros((self.H, self.W), dtype=np.uint8), cmap="viridis"
             )
-            self.ax[1, 0].set_title("Current Depth Map")
+            self.ax[1, 0].set_title("Current Height Map")
 
             self.diff_depth_map_ax = self.ax[1, 1].imshow(
                 np.zeros((self.H, self.W), dtype=np.uint8), cmap="viridis"
             )
-            self.ax[1, 1].set_title("Depth Map Difference")
+            self.ax[1, 1].set_title("Height Map Changes")
 
-            self.disc_depth_map_ax = self.ax[1, 2].imshow(
-                np.zeros((self.H, self.W), dtype=np.uint8), cmap="viridis"
+            target_vis = self.target.cpu().detach().numpy()
+            focus_vis = self.focus_map.cpu().detach().numpy()
+
+            overlay = overlay_mask(
+                to_pil_image(target_vis),
+                to_pil_image(focus_vis, mode="F"),
+                alpha=0.5,
             )
-            self.ax[1, 2].set_title("Discrete Depth Map")
+
+            self.focus_map_ax = self.ax[1, 2].imshow(overlay)
+            self.ax[1, 2].set_title("Image Focus Map (currently unused)")
 
             # Compute and store the initial height map for later difference computation.
             with torch.no_grad():
@@ -188,17 +246,8 @@ class FilamentOptimizer:
 
         tau_height, tau_global = self._get_tau()
 
-        if self.num_steps_done < self.args.iterations // 10:
-            penalty_loss = 0
-        else:
-            penalty_loss = min(
-                1.0,
-                (self.num_steps_done - self.args.iterations // 10)
-                / (self.args.iterations // 5),
-            )
-
-        #start_fraction = self.args.height_logits_learning_start_fraction
-        #full_fraction = self.args.height_logits_learning_full_fraction
+        # 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
@@ -219,18 +268,17 @@ class FilamentOptimizer:
             material_colors=self.material_colors,
             material_TDs=self.material_TDs,
             background=self.background,
-            perception_loss_module=self.perception_loss_module,
-            add_penalty_loss=1.0,
+            add_penalty_loss=10.0,
+            focus_map=self.focus_map,
+            focus_strength=0.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.
+        # We scale the gradients for the height logits by a factor to only allow updates for very strong (wrong layer/color) gradients.
         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
@@ -349,24 +397,6 @@ class FilamentOptimizer:
             )
             self.best_comp_ax.set_data(best_comp_np)
 
-            with torch.no_grad():
-                height_map = (self.max_layers * self.h) * torch.sigmoid(
-                    self.best_params["pixel_height_logits"]
-                )
-                height_map = height_map.cpu().detach().numpy()
-
-            # Normalize safely, checking for a constant image.
-            if np.allclose(height_map.max(), height_map.min()):
-                height_map_norm = np.zeros_like(height_map)
-            else:
-                height_map_norm = (height_map - height_map.min()) / (
-                    height_map.max() - height_map.min()
-                )
-
-            height_map_uint8 = (height_map_norm * 255).astype(np.uint8)
-            self.disc_depth_map_ax.set_data(height_map_uint8)
-            self.disc_depth_map_ax.set_clim(0, 255)
-
         # Update the depth map correctly.
         with torch.no_grad():
             height_map = (self.max_layers * self.h) * torch.sigmoid(
@@ -489,6 +519,7 @@ class FilamentOptimizer:
         max_swaps_allowed: int,
         min_layers_allowed: int,
         max_layers_allowed: int,
+        search_seed: bool = True,
     ):
         # Now run pruning
         from autoforge.Helper.PruningHelper import (
@@ -509,19 +540,22 @@ class FilamentOptimizer:
             self.final_tau,
             None,
         )
-        self.rng_seed_search(self.best_discrete_loss, 100, autoset_seed=True)
+        if search_seed:
+            self.rng_seed_search(self.best_discrete_loss, 100, autoset_seed=True)
         prune_num_swaps(
             self,
             max_swaps_allowed,
             self.final_tau,
             None,
         )
-        self.rng_seed_search(self.best_discrete_loss, 100, autoset_seed=True)
+        if search_seed:
+            self.rng_seed_search(self.best_discrete_loss, 100, autoset_seed=True)
 
         # prune_fireflies(self)
 
         prune_redundant_layers(self, None, min_layers_allowed, max_layers_allowed)
-        self.rng_seed_search(self.best_discrete_loss, 100, autoset_seed=True)
+        if search_seed:
+            self.rng_seed_search(self.best_discrete_loss, 100, autoset_seed=True)
 
     def _maybe_update_best_discrete(self):
         """

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

@@ -133,6 +133,12 @@ def main():
         default=True,
         help="Perform pruning after optimization",
     )
+    parser.add_argument(
+        "--pruning_rounds",
+        type=int,
+        default=2,
+        help="Number of pruning rounds to perform",
+    )
     parser.add_argument(
         "--pruning_max_colors",
         type=int,
@@ -211,6 +217,18 @@ def main():
     parser.add_argument(
         "--tensorboard", action="store_true", help="Enable TensorBoard logging"
     )
+    parser.add_argument(
+        "--num_init_rounds",
+        type=int,
+        default=64,
+        help="Number of rounds to choose the starting height map from.",
+    )
+    parser.add_argument(
+        "--num_init_cluster_layers",
+        type=int,
+        default=18,
+        help="Number of layers to cluster the image into.",
+    )
 
     args = parser.parse_args()
 
@@ -302,9 +320,9 @@ def main():
             args.layer_height,
             bgr_tuple,
             random_seed=random_seed,
-            num_threads=16,
+            num_threads=args.num_init_rounds,
             init_method="kmeans",
-            cluster_layers=18,
+            cluster_layers=args.num_init_cluster_layers,
             lab_space=True,
         )
 
@@ -360,15 +378,20 @@ def main():
     )
 
     if args.perform_pruning:
-        optimizer.prune(
-            max_colors_allowed=args.pruning_max_colors,
-            max_swaps_allowed=args.pruning_max_swaps,
-            min_layers_allowed=args.min_layers,
-            max_layers_allowed=args.pruning_max_layer,
-        )
-        optimizer.log_to_tensorboard(
-            interval=1, namespace="post_opt", step=(post_opt_step := post_opt_step + 1)
-        )
+        for i in range(args.pruning_rounds):
+            print(f"Pruning round {i + 1}/{args.pruning_rounds}:")
+            optimizer.prune(
+                max_colors_allowed=args.pruning_max_colors,
+                max_swaps_allowed=args.pruning_max_swaps,
+                min_layers_allowed=args.min_layers,
+                max_layers_allowed=args.pruning_max_layer,
+                search_seed=True if i == 0 else False,
+            )
+            optimizer.log_to_tensorboard(
+                interval=1,
+                namespace="post_opt",
+                step=(post_opt_step := post_opt_step + 1),
+            )
 
     disc_global, disc_height_image = optimizer.get_discretized_solution(best=True)
 

autoforge-1.6.2/src/autoforge/Loss/LossFunctions.py

@@ -1,91 +0,0 @@
-import torch
-import torch.nn.functional as F
-
-from autoforge.Helper.ImageHelper import srgb_to_lab
-from autoforge.Helper.OptimizerHelper import composite_image_cont
-
-
-def loss_fn(
-    params: dict,
-    target: torch.Tensor,
-    tau_height: float,
-    tau_global: float,
-    h: float,
-    max_layers: int,
-    material_colors: torch.Tensor,
-    material_TDs: torch.Tensor,
-    background: torch.Tensor,
-    perception_loss_module: torch.nn.Module,
-    add_penalty_loss: float = 0.0,
-) -> torch.Tensor:
-    """
-    Full forward pass for continuous assignment:
-    composite, then compute unified loss on (global_logits).
-    """
-    comp = composite_image_cont(
-        params["pixel_height_logits"],
-        params["global_logits"],
-        tau_height,
-        tau_global,
-        h,
-        max_layers,
-        material_colors,
-        material_TDs,
-        background,
-    )
-    return compute_loss(
-        comp=comp,
-        target=target,
-        pixel_height_logits=params["pixel_height_logits"],
-        tau_height=tau_height,
-        add_penalty_loss=add_penalty_loss,
-    )
-
-
-def compute_loss(
-    comp: torch.Tensor,
-    target: torch.Tensor,
-    pixel_height_logits: torch.Tensor = None,
-    tau_height: float = 1.0,
-    add_penalty_loss: float = 0.0,
-) -> torch.Tensor:
-    """
-    Combined MSE + Perceptual + penalty losses with patch-based smoothness.
-    """
-    # MSE Loss
-    comp_mse = srgb_to_lab(comp)
-    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]
-        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)
-        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)))
-        smoothness_loss = (loss_dx + loss_dy) * (10 * add_penalty_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