autoforge 1.1.3__tar.gz → 1.2.1__tar.gz

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

@@ -1,14 +1,14 @@
 Metadata-Version: 2.2
 Name: AutoForge
-Version: 1.1.3
+Version: 1.2.1
 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
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 Requires-Dist: numpy>=1.19.0
-Requires-Dist: jax>=0.3.10
-Requires-Dist: jaxlib>=0.3.10
+Requires-Dist: torch>=2.6.0
+Requires-Dist: torchvision>=0.21.0
 Requires-Dist: opencv-python>=4.2.0.32
 Requires-Dist: matplotlib>=3.3.0
 Requires-Dist: configargparse>=1.2.3
@@ -16,9 +16,6 @@ Requires-Dist: tqdm>=4.50.0
 Requires-Dist: pandas>=1.1.0
 Requires-Dist: scikit-image
 Requires-Dist: scikit-learn
-Requires-Dist: optax>=0.1.3
-Provides-Extra: gpu
-Requires-Dist: jax[cuda12]; extra == "gpu"
 
 # AutoForge
 
@@ -27,7 +24,10 @@ AutoForge is a Python tool for generating 3D printed layered models from an inpu
 **TLDR:** It uses a picture to generate a 3D layer image that you can print with a 3d printer. Similar to [Hueforge](https://shop.thehueforge.com/), but without the manual work (and without the artistic control).
 
 
+## Important Information
 
+We recently switched from JAX to PyTorch, to allow for a more streamlined development process. \
+This can have some unforeseen consequences, so please report any bugs you find. 
 
 ## Example
 All examples use only the 13 BambuLab Basic filaments, currently available in Hueforge.
@@ -59,7 +59,7 @@ 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.
-- **Learned Optimization**: Optimizes per-pixel height and per-layer material assignments using JAX and Optax.
+- **Learned Optimization**: Optimizes per-pixel height and per-layer material assignments using PyTorch.
 - **Gumbel Softmax Sampling**: Leverages the Gumbel softmax method to decide material assignments for each layer.
 - **STL File Generation**: Exports an ASCII STL file based on the optimized height map.
 - **Swap Instructions**: Generates clear swap instructions for changing materials during printing.
@@ -74,17 +74,6 @@ To install AutoForge, simply install the current version from PyPI:
    pip install autoforge
 ```
 
-The optimizer is built using JAX, which benefits from a CUDA-compatible GPU for optimal performance.  \
-Without a dedicated GPU the process can take significantly longer (up to 10-20x slower).  \
-If you have a GPU, you can install the GPU version of autoforge by running:\
-
-```bash
-   pip install autoforge[gpu]
-```
-
-Currently, jax gpu support is only available under linux. \
-If you use Windows, please install WSL2 to use the GPU version of jax.
-
 ## Usage
 
 The script is run from the command line and accepts several arguments. Below is an example command:
@@ -104,18 +93,27 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
 - `--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.
-- `--iterations`: Number of optimization iterations (default: 20000).
-- `--learning_rate`: Learning rate for the optimizer (default: 5e-3).
+- `--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).
+  **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). 
   **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: `#8e9089` aka Bambulab Grey).
-- `--max_size`: Maximum dimension (width or height) for the resized target image (default: 512).
+- `--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).
+  **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).
-- `--loss`: Loss function to use. Choices are `mse`, `perceptual`, or `perceptual_l1` (default: `mse`).
 - `--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).
+- `--random_seed`: Random seed for reproducibility (default: 0 (disabled) ).
+
 
 ## Outputs
 
@@ -141,6 +139,9 @@ Although I would love to be it fully color compatible with hueforge, I don't thi
 ## License
 
 AutoForge © 2025 by Hendric Voss is licensed under [CC BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/).
+The software is provided as-is and comes with no warranty or guarantee of support.
+
+The 
 
 ## Acknowledgements
 
@@ -150,7 +151,7 @@ Without it, this project would not have been possible.
 
 AutoForge makes use of several open source libraries:
 
-- [JAX](https://github.com/google/jax)
+- [PyTorch](https://pytorch.org/)
 - [Optax](https://github.com/deepmind/optax)
 - [OpenCV](https://opencv.org/)
 - [Matplotlib](https://matplotlib.org/)

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

@@ -5,7 +5,10 @@ AutoForge is a Python tool for generating 3D printed layered models from an inpu
 **TLDR:** It uses a picture to generate a 3D layer image that you can print with a 3d printer. Similar to [Hueforge](https://shop.thehueforge.com/), but without the manual work (and without the artistic control).
 
 
+## Important Information
 
+We recently switched from JAX to PyTorch, to allow for a more streamlined development process. \
+This can have some unforeseen consequences, so please report any bugs you find. 
 
 ## Example
 All examples use only the 13 BambuLab Basic filaments, currently available in Hueforge.
@@ -37,7 +40,7 @@ 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.
-- **Learned Optimization**: Optimizes per-pixel height and per-layer material assignments using JAX and Optax.
+- **Learned Optimization**: Optimizes per-pixel height and per-layer material assignments using PyTorch.
 - **Gumbel Softmax Sampling**: Leverages the Gumbel softmax method to decide material assignments for each layer.
 - **STL File Generation**: Exports an ASCII STL file based on the optimized height map.
 - **Swap Instructions**: Generates clear swap instructions for changing materials during printing.
@@ -52,17 +55,6 @@ To install AutoForge, simply install the current version from PyPI:
    pip install autoforge
 ```
 
-The optimizer is built using JAX, which benefits from a CUDA-compatible GPU for optimal performance.  \
-Without a dedicated GPU the process can take significantly longer (up to 10-20x slower).  \
-If you have a GPU, you can install the GPU version of autoforge by running:\
-
-```bash
-   pip install autoforge[gpu]
-```
-
-Currently, jax gpu support is only available under linux. \
-If you use Windows, please install WSL2 to use the GPU version of jax.
-
 ## Usage
 
 The script is run from the command line and accepts several arguments. Below is an example command:
@@ -82,18 +74,27 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
 - `--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.
-- `--iterations`: Number of optimization iterations (default: 20000).
-- `--learning_rate`: Learning rate for the optimizer (default: 5e-3).
+- `--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).
+  **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). 
   **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: `#8e9089` aka Bambulab Grey).
-- `--max_size`: Maximum dimension (width or height) for the resized target image (default: 512).
+- `--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).
+  **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).
-- `--loss`: Loss function to use. Choices are `mse`, `perceptual`, or `perceptual_l1` (default: `mse`).
 - `--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).
+- `--random_seed`: Random seed for reproducibility (default: 0 (disabled) ).
+
 
 ## Outputs
 
@@ -119,6 +120,9 @@ Although I would love to be it fully color compatible with hueforge, I don't thi
 ## License
 
 AutoForge © 2025 by Hendric Voss is licensed under [CC BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/).
+The software is provided as-is and comes with no warranty or guarantee of support.
+
+The 
 
 ## Acknowledgements
 
@@ -128,7 +132,7 @@ Without it, this project would not have been possible.
 
 AutoForge makes use of several open source libraries:
 
-- [JAX](https://github.com/google/jax)
+- [PyTorch](https://pytorch.org/)
 - [Optax](https://github.com/deepmind/optax)
 - [OpenCV](https://opencv.org/)
 - [Matplotlib](https://matplotlib.org/)

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "AutoForge"
-version="1.1.3"
+version="1.2.1"
 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"
@@ -12,25 +12,19 @@ requires-python = ">=3.9"
 license = { text = "CC BY-NC-SA 4.0" }
 dependencies = [
   "numpy>=1.19.0",
-  "jax>=0.3.10",
-  "jaxlib>=0.3.10",
+  "torch>=2.6.0",
+  "torchvision>=0.21.0",
   "opencv-python>=4.2.0.32",
   "matplotlib>=3.3.0",
   "configargparse>=1.2.3",
   "tqdm>=4.50.0",
   "pandas>=1.1.0",
   "scikit-image",
-  "scikit-learn",
-  "optax>=0.1.3"
+  "scikit-learn"
 ]
 
 [tool.setuptools]
 package-dir = {"" = "src"}
 
-[project.optional-dependencies]
-gpu = [
-  "jax[cuda12]",
-]
-
 [project.scripts]
 autoforge = "autoforge.auto_forge:main"

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

@@ -1,14 +1,14 @@
 Metadata-Version: 2.2
 Name: AutoForge
-Version: 1.1.3
+Version: 1.2.1
 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
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 Requires-Dist: numpy>=1.19.0
-Requires-Dist: jax>=0.3.10
-Requires-Dist: jaxlib>=0.3.10
+Requires-Dist: torch>=2.6.0
+Requires-Dist: torchvision>=0.21.0
 Requires-Dist: opencv-python>=4.2.0.32
 Requires-Dist: matplotlib>=3.3.0
 Requires-Dist: configargparse>=1.2.3
@@ -16,9 +16,6 @@ Requires-Dist: tqdm>=4.50.0
 Requires-Dist: pandas>=1.1.0
 Requires-Dist: scikit-image
 Requires-Dist: scikit-learn
-Requires-Dist: optax>=0.1.3
-Provides-Extra: gpu
-Requires-Dist: jax[cuda12]; extra == "gpu"
 
 # AutoForge
 
@@ -27,7 +24,10 @@ AutoForge is a Python tool for generating 3D printed layered models from an inpu
 **TLDR:** It uses a picture to generate a 3D layer image that you can print with a 3d printer. Similar to [Hueforge](https://shop.thehueforge.com/), but without the manual work (and without the artistic control).
 
 
+## Important Information
 
+We recently switched from JAX to PyTorch, to allow for a more streamlined development process. \
+This can have some unforeseen consequences, so please report any bugs you find. 
 
 ## Example
 All examples use only the 13 BambuLab Basic filaments, currently available in Hueforge.
@@ -59,7 +59,7 @@ 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.
-- **Learned Optimization**: Optimizes per-pixel height and per-layer material assignments using JAX and Optax.
+- **Learned Optimization**: Optimizes per-pixel height and per-layer material assignments using PyTorch.
 - **Gumbel Softmax Sampling**: Leverages the Gumbel softmax method to decide material assignments for each layer.
 - **STL File Generation**: Exports an ASCII STL file based on the optimized height map.
 - **Swap Instructions**: Generates clear swap instructions for changing materials during printing.
@@ -74,17 +74,6 @@ To install AutoForge, simply install the current version from PyPI:
    pip install autoforge
 ```
 
-The optimizer is built using JAX, which benefits from a CUDA-compatible GPU for optimal performance.  \
-Without a dedicated GPU the process can take significantly longer (up to 10-20x slower).  \
-If you have a GPU, you can install the GPU version of autoforge by running:\
-
-```bash
-   pip install autoforge[gpu]
-```
-
-Currently, jax gpu support is only available under linux. \
-If you use Windows, please install WSL2 to use the GPU version of jax.
-
 ## Usage
 
 The script is run from the command line and accepts several arguments. Below is an example command:
@@ -104,18 +93,27 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
 - `--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.
-- `--iterations`: Number of optimization iterations (default: 20000).
-- `--learning_rate`: Learning rate for the optimizer (default: 5e-3).
+- `--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).
+  **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). 
   **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: `#8e9089` aka Bambulab Grey).
-- `--max_size`: Maximum dimension (width or height) for the resized target image (default: 512).
+- `--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).
+  **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).
-- `--loss`: Loss function to use. Choices are `mse`, `perceptual`, or `perceptual_l1` (default: `mse`).
 - `--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).
+- `--random_seed`: Random seed for reproducibility (default: 0 (disabled) ).
+
 
 ## Outputs
 
@@ -141,6 +139,9 @@ Although I would love to be it fully color compatible with hueforge, I don't thi
 ## License
 
 AutoForge © 2025 by Hendric Voss is licensed under [CC BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/).
+The software is provided as-is and comes with no warranty or guarantee of support.
+
+The 
 
 ## Acknowledgements
 
@@ -150,7 +151,7 @@ Without it, this project would not have been possible.
 
 AutoForge makes use of several open source libraries:
 
-- [JAX](https://github.com/google/jax)
+- [PyTorch](https://pytorch.org/)
 - [Optax](https://github.com/deepmind/optax)
 - [OpenCV](https://opencv.org/)
 - [Matplotlib](https://matplotlib.org/)

{autoforge-1.1.3 → autoforge-1.2.1}/src/AutoForge.egg-info/SOURCES.txt

@@ -12,4 +12,15 @@ src/autoforge/auto_forge.py
 src/autoforge/auto_forge_gradio.py
 src/autoforge/color_fit_bayesian.py
 src/autoforge/filamentcolors_library.py
-src/autoforge/helper_functions.py
+src/autoforge/old_jax_files.py
+src/autoforge/Helper/FilamentHelper.py
+src/autoforge/Helper/ImageHelper.py
+src/autoforge/Helper/OptimizerHelper.py
+src/autoforge/Helper/OutputHelper.py
+src/autoforge/Helper/PruningHelper.py
+src/autoforge/Helper/__init__.py
+src/autoforge/Loss/LossFunctions.py
+src/autoforge/Loss/PerceptionLoss.py
+src/autoforge/Loss/__init__.py
+src/autoforge/Modules/Optimizer.py
+src/autoforge/Modules/__init__.py

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

@@ -1,6 +1,6 @@
 numpy>=1.19.0
-jax>=0.3.10
-jaxlib>=0.3.10
+torch>=2.6.0
+torchvision>=0.21.0
 opencv-python>=4.2.0.32
 matplotlib>=3.3.0
 configargparse>=1.2.3
@@ -8,7 +8,3 @@ tqdm>=4.50.0
 pandas>=1.1.0
 scikit-image
 scikit-learn
-optax>=0.1.3
-
-[gpu]
-jax[cuda12]

autoforge-1.2.1/src/autoforge/Helper/FilamentHelper.py

@@ -0,0 +1,127 @@
+import numpy as np
+import pandas as pd
+import torch
+
+
+def load_materials(csv_filename):
+    """
+    Load material data from a CSV file.
+
+    Args:
+        csv_filename (str): Path to the hueforge CSV file containing material data.
+
+    Returns:
+        tuple: A tuple containing:
+            - material_colors (jnp.ndarray): Array of material colors in float64.
+            - material_TDs (jnp.ndarray): Array of material transmission/opacity parameters in float64.
+            - material_names (list): List of material names.
+            - colors_list (list): List of color hex strings.
+    """
+    df = pd.read_csv(csv_filename)
+    material_names = [
+        brand + " - " + name
+        for brand, name in zip(df["Brand"].tolist(), df[" Name"].tolist())
+    ]
+    material_TDs = (df[" TD"].astype(float)).to_numpy()
+    colors_list = df[" Color"].tolist()
+    # Use float64 for material colors.
+    material_colors = np.array(
+        [hex_to_rgb(color) for color in colors_list], dtype=np.float64
+    )
+    material_TDs = np.array(material_TDs, dtype=np.float64)
+    return material_colors, material_TDs, material_names, colors_list
+
+
+def load_materials_data(csv_filename):
+    """
+    Load the full material data from the CSV file.
+
+    Args:
+        csv_filename (str): Path to the CSV file containing material data.
+
+    Returns:
+        list: A list of dictionaries (one per material) with keys such as
+              "Brand", "Type", "Color", "Name", "TD", "Owned", and "Uuid".
+    """
+    df = pd.read_csv(csv_filename)
+    # Use a consistent key naming. For example, convert 'TD' to 'Transmissivity' and 'Uuid' to 'uuid'
+    records = df.to_dict(orient="records")
+    return records
+
+
+def hex_to_rgb(hex_str):
+    """
+    Convert a hex color string to a normalized RGB list.
+
+    Args:
+        hex_str (str): The hex color string (e.g., '#RRGGBB').
+
+    Returns:
+        list: A list of three floats representing the RGB values normalized to [0, 1].
+    """
+    hex_str = hex_str.lstrip("#")
+    return [int(hex_str[i : i + 2], 16) / 255.0 for i in (0, 2, 4)]
+
+
+def extract_colors_from_swatches(swatch_data):
+    # we keep only data with transmission distance
+    swatch_data = [swatch for swatch in swatch_data if swatch["td"]]
+
+    # For now we load it and convert it in the same way as the hueforge csv files
+    out = {}
+    for swatch in swatch_data:
+        brand = swatch["manufacturer"]["name"]
+        name = swatch["color_name"]
+        color = swatch["hex_color"]
+        td = swatch["td"]
+        out[(brand, name)] = (color, td)
+
+    # convert to the same format as the hueforge csv files
+    material_names = [brand + " - " + name for (brand, name) in out.keys()]
+    material_colors = np.array(
+        [hex_to_rgb("#" + color) for color, _ in out.values()], dtype=np.float64
+    )
+    material_TDs = np.array([td for _, td in out.values()], dtype=np.float64)
+    colors_list = [color for color, _ in out.values()]
+
+    return material_colors, material_TDs, material_names, colors_list
+
+
+def swatch_data_to_table(swatch_data):
+    """
+    Converts swatch JSON data into a table (list of dicts) with columns:
+    "Brand", "Name", "Transmission Distance", "Hex Color".
+    """
+    table = []
+    for swatch in swatch_data:
+        if not swatch["td"]:
+            continue
+        brand = swatch["manufacturer"]["name"]
+        name = swatch["color_name"]
+        hex_color = swatch["hex_color"]
+        td = swatch["td"]
+        table.append(
+            {
+                "Brand": brand,
+                "Name": name,
+                "Transmission Distance": td,
+                "Hex Color": f"#{hex_color}",
+            }
+        )
+    return table
+
+
+def count_distinct_colors(dg: torch.Tensor) -> int:
+    """
+    Count how many distinct color/material IDs appear in dg.
+    """
+    unique_mats = torch.unique(dg)
+    return len(unique_mats)
+
+
+def count_swaps(dg: torch.Tensor) -> int:
+    """
+    Count how many color changes (swaps) occur between adjacent layers.
+    """
+    # A 'swap' is whenever dg[i] != dg[i+1].
+    return int((dg[:-1] != dg[1:]).sum().item())

autoforge-1.2.1/src/autoforge/Helper/ImageHelper.py

@@ -0,0 +1,13 @@
+import cv2
+
+
+def resize_image(img, max_size):
+    h_img, w_img, _ = img.shape
+    if w_img >= h_img:
+        new_w = max_size
+        new_h = int(max_size * h_img / w_img)
+    else:
+        new_h = max_size
+        new_w = int(max_size * w_img / h_img)
+    img_out = cv2.resize(img, (new_w, new_h))
+    return img_out