autoforge 1.2.1__tar.gz → 1.3.0__tar.gz

{autoforge-1.2.1 → autoforge-1.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: AutoForge
-Version: 1.2.1
+Version: 1.3.0
 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
@@ -9,6 +9,7 @@ Description-Content-Type: text/markdown
 Requires-Dist: numpy>=1.19.0
 Requires-Dist: torch>=2.6.0
 Requires-Dist: torchvision>=0.21.0
+Requires-Dist: transformers>=4.48.1
 Requires-Dist: opencv-python>=4.2.0.32
 Requires-Dist: matplotlib>=3.3.0
 Requires-Dist: configargparse>=1.2.3
@@ -92,28 +93,46 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
 - `--config`: *(Optional)* Path to a configuration file with the settings.
 - `--input_image`: **(Required)** Path to the input image.
 - `--csv_file`: **(Required)** Path to the CSV file containing material data. The CSV should include columns for the brand, name, color (hex code), and TD values.
-- `--output_folder`: **(Required)** Folder where output files will be saved.
+- `--output_folder`: Folder where output files will be saved. (default: `./outputs`).
 - `--iterations`: Number of optimization iterations (default: 5000).
-- `--best_loss_iterations`: Percentage of optimization iterations after which we start to record the best loss (default: 4000).
+- `--best_loss_iterations`: Percentage of optimization iterations after which we start to record the best loss (default: 4000). \
   **Note:** We currently do this as we increase a penalty loss for color and color swaps as the time goes on. Without this the penalties would have no impact.
 - `--learning_rate`: Learning rate for the optimizer (default: 1e-2).
 - `--layer_height`: Layer thickness in millimeters (default: 0.04).
-- `--max_layers`: Maximum number of layers (default: 75). 
+- `--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).  
+- `--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).
+- `--background_color`: Background color in hexadecimal format (default: `#000000` aka Black). \
   **Note:** The solver currently assumes that you have a solid color in the background, which means a color with a TD value of 4 or less (if you have a background height of 0.4)
 - `--output_size`: Maximum dimension for target image (default: 1024).
-- `--solver_size`: Maximum dimension for solver (fast) image (default: 128).
+- `--solver_size`: Maximum dimension for solver (fast) image (default: 128). \
   **Note:** We solve on a smaller size as this is many times faster, but also a bit less accurate. Increase if you need more accuracy.
 - `--decay`: Final tau value for the Gumbel-Softmax formulation (default: 0.01).
 - `--visualize`: Flag to enable live visualization of the composite image during optimization.
-- `--perform_pruning`: Perform pruning after optimization (default: True).
-- `--pruning_max_colors`: Max number of colors allowed after pruning (default: 10).
-- `--pruning_max_swaps`: Max number of swaps allowed after pruning (default: 20).
+- `--perform_pruning`: Perform pruning after optimization (default: True). \
+  **Note:** This is highly recommended even if you don't have a color/color swap limit, as it actually increases the quality of the output.
+- `--pruning_max_colors`: Max number of colors allowed after pruning (default: 100).
+- `--pruning_max_swaps`: Max number of swaps allowed after pruning (default: 100).
 - `--random_seed`: Random seed for reproducibility (default: 0 (disabled) ).
 
+### Experimental Deph Anything V2 parameters
+I got a request to add a initializing function that takes the original deph of the image into account when initializing. \
+In theory this should give you a nice background/foreground separation, but in practice it can degrade the quality of the output. \
+In addition you need the VRAM to run the depth model. \
+I currently don't recommend it as the output is not as good as without it, but I will leave it in for now. \
+
+- `--use_depth_anything`: Use a depth anything v2 model to initialize the height map (default: False).
+  **Note:** This will give you a nice background/foreground separation, but will could degrade quality in some cases.
+  **Note:** In addition you need the VRAM to run the depth model.
+
+- `--depth_strength`: Weight for blending even spacing with the cluster’s average depth when using depth initialization. (default: 0.25)
+- `--depth_threshold`: Threshold for splitting two distinct colors based on depth. (default: 0.05)
+- `--min_cluster_value`: Minimum normalized value for the lowest cluster (to avoid pure black). (default: 0.1)
+- `--w_depth`: Weight for depth difference in ordering. (default: 0.5)
+- `--w_lum`: Weight for luminance difference in ordering. (default: 1.0)
+- `--order_blend`: Blending factor between original luminance ordering (0) and depth-informed ordering (1). (default: 0.1)
+
 
 ## Outputs
 

{autoforge-1.2.1 → autoforge-1.3.0}/README.md

@@ -73,28 +73,46 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
 - `--config`: *(Optional)* Path to a configuration file with the settings.
 - `--input_image`: **(Required)** Path to the input image.
 - `--csv_file`: **(Required)** Path to the CSV file containing material data. The CSV should include columns for the brand, name, color (hex code), and TD values.
-- `--output_folder`: **(Required)** Folder where output files will be saved.
+- `--output_folder`: Folder where output files will be saved. (default: `./outputs`).
 - `--iterations`: Number of optimization iterations (default: 5000).
-- `--best_loss_iterations`: Percentage of optimization iterations after which we start to record the best loss (default: 4000).
+- `--best_loss_iterations`: Percentage of optimization iterations after which we start to record the best loss (default: 4000). \
   **Note:** We currently do this as we increase a penalty loss for color and color swaps as the time goes on. Without this the penalties would have no impact.
 - `--learning_rate`: Learning rate for the optimizer (default: 1e-2).
 - `--layer_height`: Layer thickness in millimeters (default: 0.04).
-- `--max_layers`: Maximum number of layers (default: 75). 
+- `--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).  
+- `--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).
+- `--background_color`: Background color in hexadecimal format (default: `#000000` aka Black). \
   **Note:** The solver currently assumes that you have a solid color in the background, which means a color with a TD value of 4 or less (if you have a background height of 0.4)
 - `--output_size`: Maximum dimension for target image (default: 1024).
-- `--solver_size`: Maximum dimension for solver (fast) image (default: 128).
+- `--solver_size`: Maximum dimension for solver (fast) image (default: 128). \
   **Note:** We solve on a smaller size as this is many times faster, but also a bit less accurate. Increase if you need more accuracy.
 - `--decay`: Final tau value for the Gumbel-Softmax formulation (default: 0.01).
 - `--visualize`: Flag to enable live visualization of the composite image during optimization.
-- `--perform_pruning`: Perform pruning after optimization (default: True).
-- `--pruning_max_colors`: Max number of colors allowed after pruning (default: 10).
-- `--pruning_max_swaps`: Max number of swaps allowed after pruning (default: 20).
+- `--perform_pruning`: Perform pruning after optimization (default: True). \
+  **Note:** This is highly recommended even if you don't have a color/color swap limit, as it actually increases the quality of the output.
+- `--pruning_max_colors`: Max number of colors allowed after pruning (default: 100).
+- `--pruning_max_swaps`: Max number of swaps allowed after pruning (default: 100).
 - `--random_seed`: Random seed for reproducibility (default: 0 (disabled) ).
 
+### Experimental Deph Anything V2 parameters
+I got a request to add a initializing function that takes the original deph of the image into account when initializing. \
+In theory this should give you a nice background/foreground separation, but in practice it can degrade the quality of the output. \
+In addition you need the VRAM to run the depth model. \
+I currently don't recommend it as the output is not as good as without it, but I will leave it in for now. \
+
+- `--use_depth_anything`: Use a depth anything v2 model to initialize the height map (default: False).
+  **Note:** This will give you a nice background/foreground separation, but will could degrade quality in some cases.
+  **Note:** In addition you need the VRAM to run the depth model.
+
+- `--depth_strength`: Weight for blending even spacing with the cluster’s average depth when using depth initialization. (default: 0.25)
+- `--depth_threshold`: Threshold for splitting two distinct colors based on depth. (default: 0.05)
+- `--min_cluster_value`: Minimum normalized value for the lowest cluster (to avoid pure black). (default: 0.1)
+- `--w_depth`: Weight for depth difference in ordering. (default: 0.5)
+- `--w_lum`: Weight for luminance difference in ordering. (default: 1.0)
+- `--order_blend`: Blending factor between original luminance ordering (0) and depth-informed ordering (1). (default: 0.1)
+
 
 ## Outputs
 

{autoforge-1.2.1 → autoforge-1.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "AutoForge"
-version="1.2.1"
+version="1.3.0"
 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"
@@ -14,6 +14,7 @@ dependencies = [
   "numpy>=1.19.0",
   "torch>=2.6.0",
   "torchvision>=0.21.0",
+  "transformers>=4.48.1",
   "opencv-python>=4.2.0.32",
   "matplotlib>=3.3.0",
   "configargparse>=1.2.3",

{autoforge-1.2.1 → autoforge-1.3.0}/src/AutoForge.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: AutoForge
-Version: 1.2.1
+Version: 1.3.0
 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
@@ -9,6 +9,7 @@ Description-Content-Type: text/markdown
 Requires-Dist: numpy>=1.19.0
 Requires-Dist: torch>=2.6.0
 Requires-Dist: torchvision>=0.21.0
+Requires-Dist: transformers>=4.48.1
 Requires-Dist: opencv-python>=4.2.0.32
 Requires-Dist: matplotlib>=3.3.0
 Requires-Dist: configargparse>=1.2.3
@@ -92,28 +93,46 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
 - `--config`: *(Optional)* Path to a configuration file with the settings.
 - `--input_image`: **(Required)** Path to the input image.
 - `--csv_file`: **(Required)** Path to the CSV file containing material data. The CSV should include columns for the brand, name, color (hex code), and TD values.
-- `--output_folder`: **(Required)** Folder where output files will be saved.
+- `--output_folder`: Folder where output files will be saved. (default: `./outputs`).
 - `--iterations`: Number of optimization iterations (default: 5000).
-- `--best_loss_iterations`: Percentage of optimization iterations after which we start to record the best loss (default: 4000).
+- `--best_loss_iterations`: Percentage of optimization iterations after which we start to record the best loss (default: 4000). \
   **Note:** We currently do this as we increase a penalty loss for color and color swaps as the time goes on. Without this the penalties would have no impact.
 - `--learning_rate`: Learning rate for the optimizer (default: 1e-2).
 - `--layer_height`: Layer thickness in millimeters (default: 0.04).
-- `--max_layers`: Maximum number of layers (default: 75). 
+- `--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).  
+- `--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).
+- `--background_color`: Background color in hexadecimal format (default: `#000000` aka Black). \
   **Note:** The solver currently assumes that you have a solid color in the background, which means a color with a TD value of 4 or less (if you have a background height of 0.4)
 - `--output_size`: Maximum dimension for target image (default: 1024).
-- `--solver_size`: Maximum dimension for solver (fast) image (default: 128).
+- `--solver_size`: Maximum dimension for solver (fast) image (default: 128). \
   **Note:** We solve on a smaller size as this is many times faster, but also a bit less accurate. Increase if you need more accuracy.
 - `--decay`: Final tau value for the Gumbel-Softmax formulation (default: 0.01).
 - `--visualize`: Flag to enable live visualization of the composite image during optimization.
-- `--perform_pruning`: Perform pruning after optimization (default: True).
-- `--pruning_max_colors`: Max number of colors allowed after pruning (default: 10).
-- `--pruning_max_swaps`: Max number of swaps allowed after pruning (default: 20).
+- `--perform_pruning`: Perform pruning after optimization (default: True). \
+  **Note:** This is highly recommended even if you don't have a color/color swap limit, as it actually increases the quality of the output.
+- `--pruning_max_colors`: Max number of colors allowed after pruning (default: 100).
+- `--pruning_max_swaps`: Max number of swaps allowed after pruning (default: 100).
 - `--random_seed`: Random seed for reproducibility (default: 0 (disabled) ).
 
+### Experimental Deph Anything V2 parameters
+I got a request to add a initializing function that takes the original deph of the image into account when initializing. \
+In theory this should give you a nice background/foreground separation, but in practice it can degrade the quality of the output. \
+In addition you need the VRAM to run the depth model. \
+I currently don't recommend it as the output is not as good as without it, but I will leave it in for now. \
+
+- `--use_depth_anything`: Use a depth anything v2 model to initialize the height map (default: False).
+  **Note:** This will give you a nice background/foreground separation, but will could degrade quality in some cases.
+  **Note:** In addition you need the VRAM to run the depth model.
+
+- `--depth_strength`: Weight for blending even spacing with the cluster’s average depth when using depth initialization. (default: 0.25)
+- `--depth_threshold`: Threshold for splitting two distinct colors based on depth. (default: 0.05)
+- `--min_cluster_value`: Minimum normalized value for the lowest cluster (to avoid pure black). (default: 0.1)
+- `--w_depth`: Weight for depth difference in ordering. (default: 0.5)
+- `--w_lum`: Weight for luminance difference in ordering. (default: 1.0)
+- `--order_blend`: Blending factor between original luminance ordering (0) and depth-informed ordering (1). (default: 0.1)
+
 
 ## Outputs
 

{autoforge-1.2.1 → autoforge-1.3.0}/src/AutoForge.egg-info/requires.txt

@@ -1,6 +1,7 @@
 numpy>=1.19.0
 torch>=2.6.0
 torchvision>=0.21.0
+transformers>=4.48.1
 opencv-python>=4.2.0.32
 matplotlib>=3.3.0
 configargparse>=1.2.3

{autoforge-1.2.1 → autoforge-1.3.0}/src/autoforge/Helper/OptimizerHelper.py

@@ -5,7 +5,12 @@ import torch
 import torch.nn.functional as F
 from skimage.color import rgb2lab
 from sklearn.cluster import KMeans
+from sklearn.exceptions import ConvergenceWarning
+from sklearn.utils._testing import ignore_warnings
 from tqdm import tqdm
+from transformers import pipeline
+from PIL import Image
+from sklearn.metrics import silhouette_score
 
 
 @torch.jit.script
@@ -223,6 +228,265 @@ def initialize_pixel_height_logits(target):
     return pixel_height_logits
 
 
+@ignore_warnings(category=ConvergenceWarning)
+def init_height_map_depth_color_adjusted(
+    target,
+    max_layers,
+    eps=1e-6,
+    random_seed=None,
+    depth_strength=0.25,
+    depth_threshold=0.2,
+    min_cluster_value=0.1,
+    w_depth=0.5,
+    w_lum=0.5,
+    order_blend=0.1,
+):
+    """
+    Initialize pixel height logits by combining depth and color information while allowing a blend
+    between the original luminance-based ordering and a depth-informed ordering.
+
+    Steps:
+      1. Obtain a normalized depth map using Depth Anything v2.
+      2. Determine the optimal number of color clusters (between 2 and max_layers) via silhouette score.
+      3. Cluster the image colors and (if needed) split clusters with large depth spreads.
+      4. For each final cluster, compute its average depth and average luminance.
+      5. Compute two orderings:
+            - ordering_orig: Sorted purely by average luminance (approximating the original code).
+            - ordering_depth: A TSP-inspired ordering using a weighted distance based on depth and luminance.
+      6. For each cluster, blend its rank (normalized position) between the two orderings using order_blend.
+      7. Based on the blended ordering, assign an even spacing value from min_cluster_value to 1.
+      8. Finally, blend the even spacing with the cluster's average depth using depth_strength and
+         convert the result to logits via an inverse sigmoid transform.
+
+    Args:
+        target (np.ndarray): Input image of shape (H, W, 3) in [0, 255].
+        max_layers (int): Maximum number of clusters to consider.
+        eps (float): Small constant to avoid division by zero.
+        random_seed (int): Random seed for reproducibility.
+        depth_strength (float): Weight (0 to 1) for blending even spacing with the cluster's average depth.
+        depth_threshold (float): If a cluster’s depth spread exceeds this value, it is split.
+        min_cluster_value (float): Minimum normalized value for the lowest cluster.
+        w_depth (float): Weight for depth difference in ordering_depth.
+        w_lum (float): Weight for luminance difference in ordering_depth.
+        order_blend (float): Slider (0 to 1) blending original luminance ordering (0) and depth-informed ordering (1).
+
+    Returns:
+        np.ndarray: Pixel height logits (H, W).
+    """
+
+    # ---------------------------
+    # Step 1: Obtain normalized depth map using Depth Anything v2
+    # ---------------------------
+    target_uint8 = target.astype(np.uint8)
+    image_pil = Image.fromarray(target_uint8)
+    pipe = pipeline(
+        task="depth-estimation", model="depth-anything/Depth-Anything-V2-Small-hf"
+    )
+    depth_result = pipe(image_pil)
+    depth_map = depth_result["depth"]
+    if hasattr(depth_map, "convert"):
+        depth_map = np.array(depth_map)
+    depth_map = depth_map.astype(np.float32)
+    depth_min = depth_map.min()
+    depth_max = depth_map.max()
+    depth_norm = (depth_map - depth_min) / (depth_max - depth_min + eps)
+
+    # ---------------------------
+    # Step 2: Find optimal number of clusters (n in [2, max_layers]) for color clustering
+    # ---------------------------
+    H, W, _ = target.shape
+    pixels = target.reshape(-1, 3).astype(np.float32)
+
+    def find_best_n_clusters(pixels, max_layers, random_seed):
+        sample_size = 1000
+        if pixels.shape[0] > sample_size:
+            indices = np.random.choice(pixels.shape[0], sample_size, replace=False)
+            sample_pixels = pixels[indices]
+        else:
+            sample_pixels = pixels
+        best_n = None
+        best_score = -np.inf
+        for n in range(2, max_layers + 1):
+            kmeans_temp = KMeans(n_clusters=n, random_state=random_seed)
+            labels_temp = kmeans_temp.fit_predict(sample_pixels)
+            score = silhouette_score(sample_pixels, labels_temp)
+            if score > best_score:
+                best_score = score
+                best_n = n
+        return best_n
+
+    optimal_n = find_best_n_clusters(pixels, max_layers, random_seed)
+    print(f"Optimal number of clusters: {optimal_n}")
+    # ---------------------------
+    # Step 3: Perform color clustering on the full image
+    # ---------------------------
+    kmeans = KMeans(n_clusters=optimal_n, random_state=random_seed).fit(pixels)
+    labels = kmeans.labels_.reshape(H, W)
+
+    # ---------------------------
+    # Step 4: Adjust clusters based on depth (split clusters with high depth spread)
+    # ---------------------------
+    final_labels = np.copy(labels)
+    new_cluster_id = 0
+    cluster_info = {}  # Mapping: final_cluster_id -> avg_depth
+    unique_labels = np.unique(labels)
+    for orig_label in unique_labels:
+        mask = labels == orig_label
+        cluster_depths = depth_norm[mask]
+        avg_depth = np.mean(cluster_depths)
+        depth_range = cluster_depths.max() - cluster_depths.min()
+        if depth_range > depth_threshold:
+            # Split this cluster into 2 subclusters based on depth values.
+            depth_values = cluster_depths.reshape(-1, 1)
+            k_split = 2
+            kmeans_split = KMeans(n_clusters=k_split, random_state=random_seed)
+            split_labels = kmeans_split.fit_predict(depth_values)
+            indices = np.argwhere(mask)
+            for split in range(k_split):
+                sub_mask = split_labels == split
+                inds = indices[sub_mask]
+                if inds.size == 0:
+                    continue
+                for i, j in inds:
+                    final_labels[i, j] = new_cluster_id
+                sub_avg_depth = np.mean(depth_norm[mask][split_labels == split])
+                cluster_info[new_cluster_id] = sub_avg_depth
+                new_cluster_id += 1
+        else:
+            indices = np.argwhere(mask)
+            for i, j in indices:
+                final_labels[i, j] = new_cluster_id
+            cluster_info[new_cluster_id] = avg_depth
+            new_cluster_id += 1
+
+    num_final_clusters = new_cluster_id
+
+    # ---------------------------
+    # Step 5: Compute average luminance for each final cluster (using standard weights)
+    # ---------------------------
+    cluster_colors = {}
+    for cid in range(num_final_clusters):
+        mask = final_labels == cid
+        if np.sum(mask) == 0:
+            continue
+        avg_color = np.mean(
+            target.reshape(-1, 3)[final_labels.reshape(-1) == cid], axis=0
+        )
+        lum = (
+            0.299 * avg_color[0] + 0.587 * avg_color[1] + 0.114 * avg_color[2]
+        ) / 255.0
+        cluster_colors[cid] = lum
+
+    # ---------------------------
+    # Step 6: Build cluster feature list: (cid, avg_depth, avg_luminance)
+    # ---------------------------
+    cluster_features = []
+    for cid in range(num_final_clusters):
+        avg_depth = cluster_info[cid]
+        avg_lum = cluster_colors.get(cid, 0.5)
+        cluster_features.append((cid, avg_depth, avg_lum))
+
+    # ---------------------------
+    # Step 7: Compute depth-informed ordering (TSP-inspired, using w_depth and w_lum)
+    # ---------------------------
+    def distance(feat1, feat2):
+        return w_depth * abs(feat1[1] - feat2[1]) + w_lum * abs(feat1[2] - feat2[2])
+
+    # Greedy nearest-neighbor ordering starting from cluster with lowest avg_depth
+    start_idx = np.argmin([feat[1] for feat in cluster_features])
+    unvisited = cluster_features.copy()
+    ordering_depth = []
+    current = unvisited.pop(start_idx)
+    ordering_depth.append(current)
+    while unvisited:
+        next_idx = np.argmin([distance(current, candidate) for candidate in unvisited])
+        current = unvisited.pop(next_idx)
+        ordering_depth.append(current)
+
+    # 2-opt refinement for ordering_depth
+    def total_distance(ordering):
+        return sum(
+            distance(ordering[i], ordering[i + 1]) for i in range(len(ordering) - 1)
+        )
+
+    improved = True
+    best_order_depth = ordering_depth
+    best_dist = total_distance(ordering_depth)
+    while improved:
+        improved = False
+        for i in range(1, len(best_order_depth) - 1):
+            for j in range(i + 1, len(best_order_depth)):
+                new_order = (
+                    best_order_depth[:i]
+                    + best_order_depth[i : j + 1][::-1]
+                    + best_order_depth[j + 1 :]
+                )
+                new_dist = total_distance(new_order)
+                if new_dist < best_dist:
+                    best_order_depth = new_order
+                    best_dist = new_dist
+                    improved = True
+        ordering_depth = best_order_depth
+
+    # ---------------------------
+    # Step 8: Compute original (luminance-based) ordering: simply sort by avg_lum (darkest first)
+    # ---------------------------
+    ordering_orig = sorted(cluster_features, key=lambda x: x[2])
+
+    # ---------------------------
+    # Step 9: Blend the two orderings via their rank positions using order_blend
+    # ---------------------------
+    # Map each cluster id to its rank in each ordering.
+    rank_orig = {feat[0]: idx for idx, feat in enumerate(ordering_orig)}
+    rank_depth = {feat[0]: idx for idx, feat in enumerate(ordering_depth)}
+    # Normalize ranks to [0, 1]
+    norm_rank_orig = {
+        cid: rank_orig[cid] / (len(ordering_orig) - 1) if len(ordering_orig) > 1 else 0
+        for cid in rank_orig
+    }
+    norm_rank_depth = {
+        cid: rank_depth[cid] / (len(ordering_depth) - 1)
+        if len(ordering_depth) > 1
+        else 0
+        for cid in rank_depth
+    }
+
+    # Compute blended rank for each cluster
+    blended_ranks = {}
+    for cid in norm_rank_orig:
+        blended_ranks[cid] = (1 - order_blend) * norm_rank_orig[
+            cid
+        ] + order_blend * norm_rank_depth[cid]
+
+    # Final ordering: sort clusters by blended rank (ascending)
+    final_order = sorted(cluster_features, key=lambda x: blended_ranks[x[0]])
+
+    # ---------------------------
+    # Step 10: Assign new normalized values to clusters
+    # Even spacing from min_cluster_value to 1 based on final ordering
+    even_spacing = np.linspace(min_cluster_value, 1, num_final_clusters)
+    final_mapping = {}
+    for rank, (cid, avg_depth, avg_lum) in enumerate(final_order):
+        # Blend even spacing with the cluster's average depth using depth_strength.
+        # (When depth_strength=0, purely even spacing; when 1, purely avg_depth.)
+        blended_value = (1 - depth_strength) * even_spacing[
+            rank
+        ] + depth_strength * avg_depth
+        blended_value = np.clip(blended_value, min_cluster_value, 1)
+        final_mapping[cid] = blended_value
+
+    # ---------------------------
+    # Step 11: Create new normalized label image and convert to logits.
+    # ---------------------------
+    new_labels = np.vectorize(lambda x: final_mapping[x])(final_labels).astype(
+        np.float32
+    )
+    if new_labels.max() > 0:
+        new_labels = new_labels / new_labels.max()
+    pixel_height_logits = np.log((new_labels + eps) / (1 - new_labels + eps))
+    return pixel_height_logits
+
+
 def init_height_map(target, max_layers, h, eps=1e-6, random_seed=None):
     """
     Initialize pixel height logits based on the luminance of the target image.

{autoforge-1.2.1 → autoforge-1.3.0}/src/autoforge/Helper/PruningHelper.py

@@ -1,7 +1,7 @@
 import torch
 from tqdm import tqdm
 
-from autoforge.Helper.OptimizerHelper import composite_image
+from autoforge.Helper.OptimizerHelper import composite_image, discretize_solution
 from autoforge.Loss.LossFunctions import compute_loss
 
 
@@ -66,7 +66,7 @@ def prune_num_colors(
                 material_TDs,
                 background,
                 mode="discrete",
-                rng_seed = rng_seed
+                rng_seed=rng_seed,
             )
 
         # For the loss, we pass the actual dg_test so that compute_loss() knows it’s a discrete assignment
@@ -86,16 +86,16 @@ def prune_num_colors(
     distinct_mats = torch.unique(best_dg)
     tbar = tqdm(
         range(100),
-        desc=f"Pruning colors down to {max_colors_allowed}. Current colors: {len(distinct_mats)}, Loss: {best_loss:.4f}",
+        desc=f"Pruning colors down to {max_colors_allowed} or until loss increases. Current colors: {len(distinct_mats)}, Loss: {best_loss:.4f}",
     )
-    while len(distinct_mats) > max_colors_allowed:
+    while True:
         found_merge = False
         best_merge_loss = None
         best_merge_dg_candidate = None
 
         distinct_mats = torch.unique(best_dg)
         tbar.set_description(
-            f"Pruning colors down to {max_colors_allowed}. Current colors: {len(distinct_mats)}, Loss: {best_loss:.4f}",
+            f"Pruning colors down to {max_colors_allowed} or until loss increases. Current colors: {len(distinct_mats)}, Loss: {best_loss:.4f}",
         )
         tbar.update(1)
 
@@ -111,7 +111,7 @@ def prune_num_colors(
                     best_merge_dg_candidate = dg_test
 
         if best_merge_dg_candidate is not None:
-            if best_merge_loss < best_loss or True:
+            if best_merge_loss < best_loss or len(distinct_mats) > max_colors_allowed:
                 best_dg = best_merge_dg_candidate
                 best_loss = best_merge_loss
                 found_merge = True
@@ -159,7 +159,7 @@ def prune_num_swaps(
                 material_TDs,
                 background,
                 mode="discrete",
-                rng_seed=rng_seed
+                rng_seed=rng_seed,
             )
         return compute_loss(
             material_assignment=dg_test,
@@ -175,17 +175,15 @@ def prune_num_swaps(
     best_loss = get_image_loss(best_dg)
     tbar = tqdm(
         range(100),
-        desc=f"Pruning swaps down to {max_swaps_allowed}. Current swaps: {len(find_color_bands(best_dg)) - 1}, Loss = {best_loss:.4f}",
+        desc=f"Pruning swaps down to {max_swaps_allowed} or until loss increases. Current swaps: {len(find_color_bands(best_dg)) - 1}, Loss = {best_loss:.4f}",
     )
     while True:
         bands = find_color_bands(best_dg)
         num_swaps = len(bands) - 1
         tbar.set_description(
-            f"Pruning swaps down to {max_swaps_allowed}. Current swaps: {num_swaps}, Loss = {best_loss:.4f}",
+            f"Pruning swaps down to {max_swaps_allowed} or until loss increases. Current swaps: {num_swaps}, Loss = {best_loss:.4f}",
         )
         tbar.update(1)
-        if num_swaps <= max_swaps_allowed:
-            break
 
         best_merge_loss = None
         best_merge_dg_candidate = None
@@ -211,7 +209,9 @@ def prune_num_swaps(
                 best_merge_loss = candidate_loss
                 best_merge_dg_candidate = candidate_dg
 
-        if best_merge_dg_candidate is not None:
+        if best_merge_dg_candidate is not None and (
+            best_merge_loss < best_loss or num_swaps > max_swaps_allowed
+        ):
             best_dg = best_merge_dg_candidate
             best_loss = best_merge_loss
         else:
@@ -234,7 +234,6 @@ def prune_colors_and_swaps(
     tau_for_comp: float,
     rng_seed: int,
     perception_loss_module: torch.nn.Module,
-
 ) -> torch.Tensor:
     """
     First limit # colors, then limit # swaps.
@@ -318,3 +317,183 @@ def merge_bands(
     else:
         dg_new[band_a[0] : band_a[1] + 1] = c_b
     return dg_new
+
+
+def remove_layer_from_solution(
+    params, layer_to_remove, final_tau, h, current_max_layers, rng_seed
+):
+    """
+    Remove one layer from the solution.
+
+    Args:
+        params (dict): Current parameters with keys "global_logits" and "pixel_height_logits".
+        layer_to_remove (int): Candidate layer index to remove.
+        final_tau (float): Final tau value used in discretization/compositing.
+        h (float): Layer height.
+        current_max_layers (int): Current total number of layers.
+        rng_seed (int): Seed used in discretization/compositing.
+
+    Returns:
+        new_params (dict): New parameters with the candidate layer removed.
+        new_max_layers (int): Updated number of layers.
+    """
+    # Get the current discrete height image (used to decide which pixels need adjusting)
+    _, disc_height = discretize_solution(
+        params, final_tau, h, current_max_layers, rng_seed
+    )
+
+    # Remove the candidate layer from the global (color) assignment.
+    new_global_logits = torch.cat(
+        [
+            params["global_logits"][:layer_to_remove],
+            params["global_logits"][layer_to_remove + 1 :],
+        ],
+        dim=0,
+    )
+    new_max_layers = current_max_layers - 1
+
+    # Compute current effective height: height = (current_max_layers * h) * sigmoid(pixel_height_logits)
+    current_height = (
+        current_max_layers * h * torch.sigmoid(params["pixel_height_logits"])
+    )
+
+    # For pixels where the discrete height is at or above the removed layer, subtract h.
+    new_height = current_height.clone()
+    mask = disc_height >= layer_to_remove  # <-- Changed from > to >= here
+    new_height[mask] = new_height[mask] - h
+
+    # Invert the sigmoid mapping:
+    # We need new_pixel_height_logits such that:
+    #    sigmoid(new_pixel_height_logits) = new_height / (new_max_layers * h)
+    eps = 1e-6
+    new_ratio = torch.clamp(new_height / (new_max_layers * h), eps, 1 - eps)
+    new_pixel_height_logits = torch.log(new_ratio) - torch.log(1 - new_ratio)
+
+    new_params = {
+        "global_logits": new_global_logits,
+        "pixel_height_logits": new_pixel_height_logits,
+    }
+    return new_params, new_max_layers
+
+
+def prune_redundant_layers(
+    params,
+    final_tau,
+    h,
+    target,
+    material_colors,
+    material_TDs,
+    background,
+    rng_seed,
+    perception_loss_module,
+    tolerance=1e-3,
+):
+    """
+    Iteratively check each layer and remove it if the removal does not worsen the loss by more than `tolerance`.
+
+    Args:
+        params (dict): Dictionary with keys "global_logits" and "pixel_height_logits".
+        final_tau (float): Final tau (decay) value for discretization.
+        h (float): Layer height.
+        target (torch.Tensor): Target image tensor.
+        material_colors (torch.Tensor): Tensor of material colors.
+        material_TDs (torch.Tensor): Tensor of material transmission parameters.
+        background (torch.Tensor): Background color tensor.
+        rng_seed (int): Random seed for discretization.
+        perception_loss_module (torch.nn.Module): Perceptual loss module.
+        tolerance (float): Acceptable increase in loss.
+
+    Returns:
+        current_params (dict): Updated parameters with redundant layers removed.
+        best_loss (float): Loss corresponding to the pruned solution.
+        current_max_layers (int): New number of layers after pruning.
+    """
+    current_params = params
+    current_max_layers = current_params["global_logits"].shape[0]
+
+    # Compute baseline composite and loss.
+    comp = composite_image(
+        current_params["pixel_height_logits"],
+        current_params["global_logits"],
+        final_tau,
+        final_tau,
+        h,
+        current_max_layers,
+        material_colors,
+        material_TDs,
+        background,
+        mode="discrete",
+        rng_seed=rng_seed,
+    )
+    disc_global, _ = discretize_solution(
+        current_params, final_tau, h, current_max_layers, rng_seed
+    )
+    best_loss = compute_loss(
+        disc_global,
+        comp,
+        target,
+        perception_loss_module,
+        final_tau,
+        material_colors.shape[0],
+        add_penalty_loss=False,
+    ).item()
+
+    improvement = True
+    # Optional: use a progress bar to indicate the number of layers removed.
+    from tqdm import tqdm
+
+    tbar = tqdm(
+        desc=f"Pruning redundant layers: Loss {best_loss:.4f}",
+        total=current_max_layers - 1,
+    )
+    removed_layers = 0
+
+    # Continue until no layer can be removed.
+    while improvement and current_max_layers > 1:
+        tbar.update(1)
+        improvement = False
+        # Try each layer as a candidate for removal.
+        for layer in range(current_max_layers):
+            candidate_params, candidate_max_layers = remove_layer_from_solution(
+                current_params, layer, final_tau, h, current_max_layers, rng_seed
+            )
+            candidate_comp = composite_image(
+                candidate_params["pixel_height_logits"],
+                candidate_params["global_logits"],
+                final_tau,
+                final_tau,
+                h,
+                candidate_max_layers,
+                material_colors,
+                material_TDs,
+                background,
+                mode="discrete",
+                rng_seed=rng_seed,
+            )
+            candidate_disc_global, _ = discretize_solution(
+                candidate_params, final_tau, h, candidate_max_layers, rng_seed
+            )
+            candidate_loss = compute_loss(
+                candidate_disc_global,
+                candidate_comp,
+                target,
+                perception_loss_module,
+                final_tau,
+                material_colors.shape[0],
+                add_penalty_loss=False,
+            ).item()
+
+            # Accept the removal if loss increase is within tolerance.
+            if candidate_loss <= best_loss + tolerance:
+                removed_layers += 1
+                tbar.set_description(
+                    f"Pruning redundant layers: Loss {candidate_loss:.4f}, Removed {removed_layers}, new max layers {candidate_max_layers}"
+                )
+                current_params = candidate_params
+                current_max_layers = candidate_max_layers
+                best_loss = candidate_loss
+                improvement = True
+                break  # Restart search after a successful removal.
+        # End for each candidate.
+    tbar.close()
+    return current_params, best_loss, current_max_layers

{autoforge-1.2.1 → autoforge-1.3.0}/src/autoforge/Modules/Optimizer.py

@@ -108,7 +108,7 @@ class FilamentOptimizer:
                 "We only record the best discrete function at the end to allow for penalty loss to function correctly."
             )
             plt.ion()
-            self.fig, self.ax = plt.subplots(1, 3, figsize=(14, 6))
+            self.fig, self.ax = plt.subplots(1, 4, figsize=(14, 6))
             self.target_im_ax = self.ax[0].imshow(
                 np.array(self.target.cpu(), dtype=np.uint8)
             )
@@ -117,10 +117,16 @@ class FilamentOptimizer:
                 np.zeros((self.H, self.W, 3), dtype=np.uint8)
             )
             self.ax[1].set_title("Current Composite")
-            self.best_comp_ax = self.ax[2].imshow(
+
+            self.depth_map_ax = self.ax[2].imshow(
+                np.zeros((self.H, self.W), dtype=np.uint8), cmap="viridis"
+            )
+            self.ax[2].set_title("Current Depth Map")
+
+            self.best_comp_ax = self.ax[3].imshow(
                 np.zeros((self.H, self.W, 3), dtype=np.uint8)
             )
-            self.ax[2].set_title("Best Discrete Composite")
+            self.ax[3].set_title("Best Discrete Composite")
             plt.pause(0.1)
 
     def _get_tau(self):
@@ -232,6 +238,26 @@ class FilamentOptimizer:
             )
             self.best_comp_ax.set_data(best_comp_np)
 
+        # Update the depth map correctly.
+        with torch.no_grad():
+            height_map = (self.max_layers * self.h) * torch.sigmoid(
+                self.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.depth_map_ax.set_data(height_map_uint8)
+        # Update the color limits so that the colormap correctly maps the range 0–255.
+        self.depth_map_ax.set_clim(0, 255)
+
         self.fig.suptitle(
             f"Step {self.num_steps_done}, Tau: {tau_g:.4f}, Loss: {self.loss:.4f}, Best Discrete Loss: {self.best_discrete_loss:.4f}"
         )

{autoforge-1.2.1 → autoforge-1.3.0}/src/autoforge/auto_forge.py

@@ -12,12 +12,14 @@ from autoforge.Helper.OptimizerHelper import (
     init_height_map,
     composite_image,
     discretize_solution,
+    init_height_map_depth_color_adjusted,
 )
 from autoforge.Helper.OutputHelper import (
     generate_stl,
     generate_swap_instructions,
     generate_project_file,
 )
+from autoforge.Helper.PruningHelper import prune_redundant_layers
 from autoforge.Loss.PerceptionLoss import MultiLayerVGGPerceptualLoss
 from autoforge.Modules.Optimizer import FilamentOptimizer
 
@@ -102,13 +104,13 @@ def main():
     parser.add_argument(
         "--pruning_max_colors",
         type=int,
-        default=10,
+        default=100,
         help="Max number of colors allowed after pruning",
     )
     parser.add_argument(
         "--pruning_max_swaps",
         type=int,
-        default=20,
+        default=100,
         help="Max number of swaps allowed after pruning",
     )
 
@@ -119,6 +121,48 @@ def main():
         help="Specify the random seed, or use 0 for automatic generation",
     )
 
+    parser.add_argument(
+        "--use_depth_anything_height_initialization",
+        action="store_true",
+        help="Use Depth Anything v2 for height image initialization",
+    )
+    parser.add_argument(
+        "--depth_strength",
+        type=float,
+        default=0.25,
+        help="Weight (0 to 1) for blending even spacing with the cluster's average depth",
+    )
+    parser.add_argument(
+        "--depth_threshold",
+        type=float,
+        default=0.05,
+        help="Threshold for splitting two distinct colors based on depth",
+    )
+    parser.add_argument(
+        "--min_cluster_value",
+        type=float,
+        default=0.1,
+        help="Minimum normalized value for the lowest cluster (to avoid pure black)",
+    )
+    parser.add_argument(
+        "--w_depth",
+        type=float,
+        default=0.5,
+        help="Weight for depth difference in ordering_depth",
+    )
+    parser.add_argument(
+        "--w_lum",
+        type=float,
+        default=1.0,
+        help="Weight for luminance difference in ordering_depth",
+    )
+    parser.add_argument(
+        "--order_blend",
+        type=float,
+        default=0.1,
+        help="Slider (0 to 1) blending original luminance ordering (0) and depth-informed ordering (1).",
+    )
+
     args = parser.parse_args()
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -163,9 +207,22 @@ def main():
     output_target = torch.tensor(output_img_np, dtype=torch.float32, device=device)
 
     # Initialize pixel_height_logits from the large (final) image
-    pixel_height_logits_init = init_height_map(
-        output_img_np, args.max_layers, args.layer_height, random_seed=random_seed
-    )
+    if args.use_depth_anything_height_initialization:
+        pixel_height_logits_init = init_height_map_depth_color_adjusted(
+            output_img_np,
+            args.max_layers,
+            depth_strength=args.depth_strength,
+            depth_threshold=args.depth_threshold,
+            min_cluster_value=args.min_cluster_value,
+            w_depth=args.w_depth,
+            w_lum=args.w_lum,
+            order_blend=args.order_blend,
+        )
+    else:
+        # Default initialization using your existing method
+        pixel_height_logits_init = init_height_map(
+            output_img_np, args.max_layers, args.layer_height, random_seed=random_seed
+        )
 
     # But we will solve at the smaller resolution => resize that init to solver size
     # Add a dummy channel so that cv2.resize can handle it
@@ -231,16 +288,27 @@ def main():
         "global_logits": disc_global,
         "pixel_height_logits": torch.from_numpy(pixel_height_logits_init).to(device),
     }
+
+    if args.perform_pruning:
+        print("Removing redundant layers from the discrete solution...")
+        params, pruned_loss, max_layers = prune_redundant_layers(
+            params,
+            final_tau=args.decay,
+            h=args.layer_height,
+            target=output_target,  # or whichever target image you use for final compositing
+            material_colors=material_colors,
+            material_TDs=material_TDs,
+            background=background,
+            rng_seed=optimizer.best_seed,  # using the best seed from your optimizer
+            perception_loss_module=perception_loss_module,
+            tolerance=1e-3,  # adjust as needed
+        )
+        args.max_layers = max_layers
+
     disc_global, disc_height_image = discretize_solution(
         params, args.decay, args.layer_height, args.max_layers
     )
-    # Overwrite the big resolution height logits into optimizer params
-    # If you want to do a final compositing or more sophisticated approach, do it here:
-    # Example: (but keep in mind the shapes might differ)
-    # _, disc_height_image = optimizer.get_discretized_solution(
-    #     best=True,
-    #     custom_height_logits=torch.from_numpy(pixel_height_logits_init).to(device),
-    # )
+
     print("Done. Saving outputs...")
     # Save Image
     comp_disc = composite_image(