rgb-to-segmentation 0.0.2__tar.gz

rgb_to_segmentation-0.0.2/.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+local/
+.vscode/
+lightning_logs/

rgb_to_segmentation-0.0.2/PKG-INFO

@@ -0,0 +1,142 @@
+Metadata-Version: 2.4
+Name: rgb-to-segmentation
+Version: 0.0.2
+Summary: Tools for processing and cleaning segmentation images using palette mapping and neural networks
+Author: Alex Senden
+Maintainer: Alex Senden
+License: MIT
+Keywords: image-processing,machine-learning,pytorch,segmentation
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Topic :: Scientific/Engineering :: Image Processing
+Requires-Python: >=3.8
+Requires-Dist: numpy>=1.21.0
+Requires-Dist: pillow>=8.0.0
+Requires-Dist: pytorch-lightning>=1.5.0
+Requires-Dist: scipy>=1.7.0
+Requires-Dist: torch>=1.9.0
+Requires-Dist: torchvision>=0.10.0
+Description-Content-Type: text/markdown
+
+# RGB to Segmentation
+
+A Python package for processing and cleaning segmentation images. This package provides tools to convert RGB images to segmentation masks using palette-based color mapping and neural network-based refinement.
+
+## Features
+
+- **Palette-based Cleaning**: Clean noisy segmentation images by mapping pixels to the nearest colors in a predefined palette, with optional morphological operations to refine boundaries.
+- **Neural Network Refinement**: Use a trained pixelwise classifier to refine segmentation masks using PyTorch Lightning.
+- **Command-Line Interface**: Unified CLI for cleaning with method selection, plus separate training command.
+- **Programmatic API**: Direct access to cleaning and training functions for integration into other workflows.
+
+## Installation
+
+Install from PyPI:
+
+```bash
+pip install rgb-to-segmentation
+```
+
+Or install from source:
+
+```bash
+git clone https://github.com/alexsenden/rgb-to-segmentation.git
+cd rgb-to-segmentation
+pip install .
+```
+
+## Usage
+
+### Cleaning Noisy Segmentation Images
+
+Use the `segment-clean` command to clean segmentation images using various methods:
+
+#### Palette-based cleaning:
+
+```bash
+segment-clean --method palette --input_dir /path/to/input --output_dir /path/to/output --colour_map "0,0,0;255,0,0;0,255,0" --output_type rgb
+```
+
+#### Neural network-based cleaning:
+
+```bash
+segment-clean --method nn --input_dir /path/to/input --output_dir /path/to/output --model_path /path/to/model.ckpt --colour_map "0,0,0;255,0,0;0,255,0" --output_type index
+```
+
+You can also provide colours via file with `--colour_map_file /path/to/colours.txt` (one `r,g,b` per line). The CLI parses colours and constructs the palette/colour map internally, mirroring the Python API which accepts parsed structures (NumPy array for palette, dictionary for colour map).
+
+Options:
+
+- `--method`: Cleaning method ('palette' or 'nn')
+- `--input_dir`: Path to input directory containing images
+- `--output_dir`: Directory where cleaned images will be written
+- `--inplace`: Overwrite input images in place
+- `--exts`: Comma-separated list of allowed image extensions
+- `--name_filter`: Only process files whose name contains this substring
+- `--output_type`: Output format ('rgb' or 'index')
+
+For palette method:
+- `--colour_map`: Semicolon-separated list of RGB triples
+- `--colour_map_file`: Path to a file listing RGB triples
+- `--morph_kernel_size`: Size of morphological kernel for boundary cleaning
+
+For nn method:
+- `--model_path`: Path to trained model file
+- `--colour_map`: Semicolon-separated list of RGB triples
+- `--colour_map_file`: Path to a file listing RGB triples
+
+### Training the Neural Network Model
+
+Train a pixelwise classifier to refine segmentation masks:
+
+```bash
+segment-train --image_dir /path/to/noisy_images --label_dir /path/to/labels --output_dir /path/to/model_output --colour_map "0,0,0;255,0,0;0,255,0"
+```
+
+Options:
+- `--image_dir`: Path to directory containing noisy images
+- `--label_dir`: Path to directory containing target RGB labels
+- `--output_dir`: Directory where model weights will be saved
+- `--colour_map`: Semicolon-separated list of RGB triples
+- `--colour_map_file`: Path to a file listing RGB triples
+- `--model_type`: The type of model to train (default: pixelwise)
+
+Note that one label image may have multiple corresponding noisy masks. Labels are matched to noisy masks whose filenames contain the label file basename (pre-extension name, i.e. `my_image.png` -> `my_image`).
+
+## API
+
+You can also use the package programmatically:
+
+```python
+import numpy as np
+from rgb_to_segmentation import clean, nn, train, utils
+
+# Palette cleaning
+colours = utils.parse_colours_from_string("0,0,0;255,0,0;0,255,0")
+palette = np.asarray(colours, dtype=np.uint8)
+clean.clean_segmentation(input_dir="/path/to/input", output_dir="/path/to/output", palette=palette, output_type="index")
+
+# NN inference
+colours = utils.parse_colours_from_string("0,0,0;255,0,0;0,255,0")
+colour_map = {i: rgb for i, rgb in enumerate(colours)}
+nn.run_inference(input_dir="/path/to/input", output_dir="/path/to/output", model_path="/path/to/model.ckpt", colour_map=colour_map, output_type="rgb")
+
+# Train model
+colours = utils.parse_colours_from_string("0,0,0;255,0,0;0,255,0")
+colour_map = {i: rgb for i, rgb in enumerate(colours)}
+train.train_model(image_dir="/path/to/images", label_dir="/path/to/labels", output_dir="/path/to/output", colour_map=colour_map)
+```
+
+## Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request.
+
+## License
+
+This project is licensed under the MIT License - see the LICENSE file for details.

rgb_to_segmentation-0.0.2/README.md

@@ -0,0 +1,116 @@
+# RGB to Segmentation
+
+A Python package for processing and cleaning segmentation images. This package provides tools to convert RGB images to segmentation masks using palette-based color mapping and neural network-based refinement.
+
+## Features
+
+- **Palette-based Cleaning**: Clean noisy segmentation images by mapping pixels to the nearest colors in a predefined palette, with optional morphological operations to refine boundaries.
+- **Neural Network Refinement**: Use a trained pixelwise classifier to refine segmentation masks using PyTorch Lightning.
+- **Command-Line Interface**: Unified CLI for cleaning with method selection, plus separate training command.
+- **Programmatic API**: Direct access to cleaning and training functions for integration into other workflows.
+
+## Installation
+
+Install from PyPI:
+
+```bash
+pip install rgb-to-segmentation
+```
+
+Or install from source:
+
+```bash
+git clone https://github.com/alexsenden/rgb-to-segmentation.git
+cd rgb-to-segmentation
+pip install .
+```
+
+## Usage
+
+### Cleaning Noisy Segmentation Images
+
+Use the `segment-clean` command to clean segmentation images using various methods:
+
+#### Palette-based cleaning:
+
+```bash
+segment-clean --method palette --input_dir /path/to/input --output_dir /path/to/output --colour_map "0,0,0;255,0,0;0,255,0" --output_type rgb
+```
+
+#### Neural network-based cleaning:
+
+```bash
+segment-clean --method nn --input_dir /path/to/input --output_dir /path/to/output --model_path /path/to/model.ckpt --colour_map "0,0,0;255,0,0;0,255,0" --output_type index
+```
+
+You can also provide colours via file with `--colour_map_file /path/to/colours.txt` (one `r,g,b` per line). The CLI parses colours and constructs the palette/colour map internally, mirroring the Python API which accepts parsed structures (NumPy array for palette, dictionary for colour map).
+
+Options:
+
+- `--method`: Cleaning method ('palette' or 'nn')
+- `--input_dir`: Path to input directory containing images
+- `--output_dir`: Directory where cleaned images will be written
+- `--inplace`: Overwrite input images in place
+- `--exts`: Comma-separated list of allowed image extensions
+- `--name_filter`: Only process files whose name contains this substring
+- `--output_type`: Output format ('rgb' or 'index')
+
+For palette method:
+- `--colour_map`: Semicolon-separated list of RGB triples
+- `--colour_map_file`: Path to a file listing RGB triples
+- `--morph_kernel_size`: Size of morphological kernel for boundary cleaning
+
+For nn method:
+- `--model_path`: Path to trained model file
+- `--colour_map`: Semicolon-separated list of RGB triples
+- `--colour_map_file`: Path to a file listing RGB triples
+
+### Training the Neural Network Model
+
+Train a pixelwise classifier to refine segmentation masks:
+
+```bash
+segment-train --image_dir /path/to/noisy_images --label_dir /path/to/labels --output_dir /path/to/model_output --colour_map "0,0,0;255,0,0;0,255,0"
+```
+
+Options:
+- `--image_dir`: Path to directory containing noisy images
+- `--label_dir`: Path to directory containing target RGB labels
+- `--output_dir`: Directory where model weights will be saved
+- `--colour_map`: Semicolon-separated list of RGB triples
+- `--colour_map_file`: Path to a file listing RGB triples
+- `--model_type`: The type of model to train (default: pixelwise)
+
+Note that one label image may have multiple corresponding noisy masks. Labels are matched to noisy masks whose filenames contain the label file basename (pre-extension name, i.e. `my_image.png` -> `my_image`).
+
+## API
+
+You can also use the package programmatically:
+
+```python
+import numpy as np
+from rgb_to_segmentation import clean, nn, train, utils
+
+# Palette cleaning
+colours = utils.parse_colours_from_string("0,0,0;255,0,0;0,255,0")
+palette = np.asarray(colours, dtype=np.uint8)
+clean.clean_segmentation(input_dir="/path/to/input", output_dir="/path/to/output", palette=palette, output_type="index")
+
+# NN inference
+colours = utils.parse_colours_from_string("0,0,0;255,0,0;0,255,0")
+colour_map = {i: rgb for i, rgb in enumerate(colours)}
+nn.run_inference(input_dir="/path/to/input", output_dir="/path/to/output", model_path="/path/to/model.ckpt", colour_map=colour_map, output_type="rgb")
+
+# Train model
+colours = utils.parse_colours_from_string("0,0,0;255,0,0;0,255,0")
+colour_map = {i: rgb for i, rgb in enumerate(colours)}
+train.train_model(image_dir="/path/to/images", label_dir="/path/to/labels", output_dir="/path/to/output", colour_map=colour_map)
+```
+
+## Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request.
+
+## License
+
+This project is licensed under the MIT License - see the LICENSE file for details.

rgb_to_segmentation-0.0.2/pyproject.toml

@@ -0,0 +1,48 @@
+[build-system]
+requires = ["hatchling", "hatch-vcs"]
+build-backend = "hatchling.build"
+
+[project]
+name = "rgb-to-segmentation"
+description = "Tools for processing and cleaning segmentation images using palette mapping and neural networks"
+readme = "README.md"
+license = { text = "MIT" }
+requires-python = ">=3.8"
+authors = [{ name = "Alex Senden" }]
+maintainers = [{ name = "Alex Senden" }]
+keywords = ["segmentation", "image-processing", "pytorch", "machine-learning"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Topic :: Scientific/Engineering :: Image Processing",
+]
+dependencies = [
+    "torch>=1.9.0",
+    "torchvision>=0.10.0",
+    "pytorch-lightning>=1.5.0",
+    "numpy>=1.21.0",
+    "scipy>=1.7.0",
+    "Pillow>=8.0.0",
+]
+
+# Version is automatically provided by hatch-vcs
+dynamic = ["version"]
+
+[project.scripts]
+segment-clean = "rgb_to_segmentation.cli:main_clean"
+segment-train = "rgb_to_segmentation.cli:main_train"
+
+[tool.hatch.version]
+source = "vcs"
+
+[tool.hatch.build.targets.sdist]
+include = ["rgb_to_segmentation/**", "README.md"]
+
+[tool.hatch.build.targets.wheel]
+include = ["rgb_to_segmentation/**"]

rgb_to_segmentation-0.0.2/rgb_to_segmentation/clean.py

@@ -0,0 +1,219 @@
+import os
+from typing import List
+
+import numpy as np
+
+from PIL import Image
+from scipy import ndimage
+
+
+def nearest_palette_image(image_array: np.ndarray, palette: np.ndarray) -> np.ndarray:
+    """
+    Assign each pixel in `image_array` (H,W,3 uint8) to the nearest colour in `palette` (K,3 uint8).
+    Returns recoloured image array with same shape and dtype uint8.
+    """
+    if image_array.ndim != 3 or image_array.shape[2] != 3:
+        raise ValueError("image_array must have shape (H, W, 3)")
+
+    h, w, _ = image_array.shape
+    flat = image_array.reshape(-1, 3).astype(np.int64)
+    pal = palette.astype(np.int64)
+
+    # Compute squared distances between each pixel and each palette colour.
+    # distances shape: (N_pixels, K)
+    d = np.sum((flat[:, None, :] - pal[None, :, :]) ** 2, axis=2)
+
+    idx = np.argmin(d, axis=1)
+    new_flat = pal[idx]
+    new = new_flat.reshape(h, w, 3).astype(np.uint8)
+    return new
+
+
+def get_palette_for_image(
+    image_array: np.ndarray, full_palette: np.ndarray
+) -> np.ndarray:
+    """
+    Identify which colours from the full palette are present in the image,
+    and return only those colours.
+    """
+    h, w, _ = image_array.shape
+    flat_img = image_array.reshape(-1, 3).astype(np.int16)
+    pal = full_palette.astype(np.int16)
+
+    # For each pixel, find the nearest palette colour
+    d = np.sum((flat_img[:, None, :] - pal[None, :, :]) ** 2, axis=2)
+    idx = np.argmin(d, axis=1)
+
+    # Get unique indices that are actually used
+    unique_idx = np.unique(idx)
+
+    # Return only the palette colours that are used
+    return full_palette[unique_idx]
+
+
+def apply_morphological_clean(image_array: np.ndarray, kernel_size: int) -> np.ndarray:
+    """
+    Apply morphological closing (erosion followed by dilation) per class to clean up
+    class boundaries and remove noise.
+    """
+    if kernel_size <= 0:
+        return image_array
+
+    # Create morphological kernel
+    kernel = ndimage.generate_binary_structure(2, 2)
+
+    # Get unique colours that actually appear in the image
+    h, w, _ = image_array.shape
+    flat_img = image_array.reshape(-1, 3)
+    unique_colours = np.unique(flat_img, axis=0)
+
+    # Process each class separately to avoid blending
+    result = np.zeros_like(image_array)
+
+    for colour in unique_colours:
+        # Create binary mask for this class
+        mask = np.all(image_array == colour, axis=-1)
+
+        # Apply closing: erosion then dilation
+        for _ in range(kernel_size):
+            mask = ndimage.binary_erosion(mask, structure=kernel)
+        for _ in range(kernel_size):
+            mask = ndimage.binary_dilation(mask, structure=kernel)
+
+        # Assign pixels back
+        result[mask] = colour
+
+    # Fill any remaining pixels (from eroded areas) with nearest colour from result
+    unfilled = ~np.any(result != 0, axis=-1)
+    if np.any(unfilled):
+        # For unfilled pixels, use nearest palette colour again or copy from nearby
+        result[unfilled] = image_array[unfilled]
+
+    return result
+
+
+def rgb_image_to_index(image_array: np.ndarray, palette: np.ndarray) -> np.ndarray:
+    """
+    Map each RGB pixel in `image_array` to the index of the matching colour in `palette`.
+    Assumes pixels take values from `palette`.
+    """
+    h, w, _ = image_array.shape
+    palette_list = [tuple(map(int, c)) for c in palette.tolist()]
+    lookup = {c: i for i, c in enumerate(palette_list)}
+    flat = image_array.reshape(-1, 3)
+    idx = np.array([lookup[tuple(map(int, px))] for px in flat], dtype=np.uint16)
+    return idx.reshape(h, w)
+
+
+def process_file(
+    input_path: str,
+    output_path: str,
+    palette: np.ndarray,
+    kernel_size: int,
+    output_type: str = "rgb",
+):
+    try:
+        img = Image.open(input_path).convert("RGB")
+    except Exception as e:
+        print(f"Skipping {input_path}: cannot open image ({e})")
+        return
+
+    arr = np.array(img, dtype=np.uint8)
+
+    # Reduce palette to only colours present in this image
+    reduced_palette = get_palette_for_image(arr, palette)
+
+    cleaned = nearest_palette_image(arr, reduced_palette)
+
+    # Apply morphological transformations if kernel_size > 0
+    if kernel_size > 0:
+        cleaned = apply_morphological_clean(cleaned, kernel_size)
+
+    # Ensure output directory exists
+    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
+
+    if output_type == "rgb":
+        Image.fromarray(cleaned).save(output_path)
+    elif output_type == "index":
+        index_mask = rgb_image_to_index(cleaned, reduced_palette)
+        Image.fromarray(index_mask.astype(np.uint8), mode="L").save(output_path)
+    else:
+        raise ValueError("output_type must be 'rgb' or 'index'")
+
+
+def process_directory(
+    input_dir: str,
+    output_dir: str,
+    palette: np.ndarray,
+    exts: List[str],
+    inplace: bool,
+    name_filter: str = "",
+    kernel_size: int = 0,
+    output_type: str = "rgb",
+):
+    exts = [e.lower().strip() for e in exts]
+    for root, dirs, files in os.walk(input_dir):
+        # Determine the corresponding output root
+        rel = os.path.relpath(root, input_dir)
+        out_root = os.path.join(output_dir, rel) if not inplace else root
+        os.makedirs(out_root, exist_ok=True)
+
+        for fname in files:
+            if not any(fname.lower().endswith(e) for e in exts):
+                continue
+            if name_filter and name_filter not in fname:
+                continue
+            in_path = os.path.join(root, fname)
+            out_path = os.path.join(out_root, fname)
+            process_file(in_path, out_path, palette, kernel_size, output_type)
+
+
+def clean_segmentation(
+    input_dir: str,
+    output_dir: str = None,
+    inplace: bool = False,
+    palette: np.ndarray = None,
+    exts: str = ".png,.jpg,.jpeg,.tiff,.bmp,.gif",
+    name_filter: str = "",
+    morph_kernel_size: int = 3,
+    output_type: str = "rgb",
+):
+    """
+    Clean segmentation images using palette-based color mapping.
+
+    Args:
+        input_dir (str): Path to input directory containing segmentation images.
+        output_dir (str, optional): Directory where cleaned images will be written. Required if not inplace.
+        inplace (bool): Overwrite input images in place.
+        palette (np.ndarray): Array of RGB triples (K, 3) uint8.
+        exts (str): Comma-separated list of allowed image extensions.
+        name_filter (str): Only process files whose name contains this substring.
+        morph_kernel_size (int): Size of morphological kernel for boundary cleaning.
+    """
+    if not inplace and output_dir is None:
+        raise ValueError("Either output_dir must be provided or inplace must be True")
+
+    if palette is None:
+        raise ValueError("palette must be provided")
+
+    exts_list = [e if e.startswith(".") else "." + e for e in exts.split(",")]
+
+    out_dir = output_dir if not inplace else input_dir
+
+    if not inplace:
+        os.makedirs(out_dir, exist_ok=True)
+
+    print(
+        f"Processing: input={input_dir} -> output={out_dir}, colours={len(palette)}, morph_kernel={morph_kernel_size}, output_type={output_type}"
+    )
+    process_directory(
+        input_dir,
+        out_dir,
+        palette,
+        exts_list,
+        inplace,
+        name_filter,
+        morph_kernel_size,
+        output_type,
+    )
+    print("Done.")

rgb_to_segmentation-0.0.2/rgb_to_segmentation/cli.py

@@ -0,0 +1,179 @@
+import argparse
+
+import numpy as np
+
+from . import clean, nn, train, utils
+
+
+def main_clean():
+    parser = argparse.ArgumentParser(
+        description="Clean segmentation images using various methods."
+    )
+
+    parser.add_argument(
+        "--method",
+        type=str,
+        required=True,
+        choices=["palette", "nn"],
+        help="Cleaning method to use: 'palette' for color palette mapping, 'nn' for neural network.",
+    )
+    parser.add_argument(
+        "--input_dir",
+        type=str,
+        required=True,
+        help="Path to input directory containing images.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        required=False,
+        help="Directory where cleaned images will be written. Required if not inplace.",
+    )
+    parser.add_argument(
+        "--inplace",
+        action="store_true",
+        help="Overwrite input images in place.",
+    )
+    parser.add_argument(
+        "--exts",
+        type=str,
+        default=".png,.jpg,.jpeg,.tiff,.bmp,.gif",
+        help="Comma-separated list of allowed image extensions.",
+    )
+    parser.add_argument(
+        "--name_filter",
+        type=str,
+        default="",
+        help="Only process files whose name contains this substring.",
+    )
+
+    parser.add_argument(
+        "--output_type",
+        type=str,
+        choices=["rgb", "index"],
+        default="rgb",
+        help="Output format: 'rgb' colour image or 'index' mask.",
+    )
+
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument(
+        "--colour_map",
+        type=str,
+        help="Semicolon-separated list of RGB triples.",
+    )
+    group.add_argument(
+        "--colour_map_file",
+        type=str,
+        help="Path to a file listing RGB triples.",
+    )
+
+    # Palette-specific args
+    parser.add_argument(
+        "--morph_kernel_size",
+        type=int,
+        default=3,
+        help="Size of morphological kernel for palette method.",
+    )
+
+    # NN-specific args
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        help="Path to trained model for nn method.",
+    )
+
+    args = parser.parse_args()
+
+    if args.colour_map_file:
+        colours = utils.parse_colours_from_file(args.colour_map_file)
+    else:
+        colours = utils.parse_colours_from_string(args.colour_map)
+
+    if args.method == "palette":
+        palette = np.asarray(colours, dtype=np.uint8)
+
+        clean.clean_segmentation(
+            input_dir=args.input_dir,
+            output_dir=args.output_dir,
+            inplace=args.inplace,
+            palette=palette,
+            exts=args.exts,
+            name_filter=args.name_filter,
+            morph_kernel_size=args.morph_kernel_size,
+            output_type=args.output_type,
+        )
+
+    elif args.method == "nn":
+        if not args.model_path:
+            parser.error("--model_path required for nn method")
+
+        colour_map = {i: rgb for i, rgb in enumerate(colours)}
+        nn.run_inference(
+            input_dir=args.input_dir,
+            output_dir=args.output_dir,
+            inplace=args.inplace,
+            model_path=args.model_path,
+            colour_map=colour_map,
+            exts=args.exts,
+            name_filter=args.name_filter,
+            output_type=args.output_type,
+        )
+
+
+def main_train():
+    parser = argparse.ArgumentParser(
+        description="Train a neural network model for segmentation cleaning."
+    )
+
+    parser.add_argument(
+        "--image_dir",
+        type=str,
+        required=True,
+        help="Path to directory containing noisy images.",
+    )
+    parser.add_argument(
+        "--label_dir",
+        type=str,
+        required=True,
+        help="Path to directory containing target RGB labels.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        required=True,
+        help="Directory where model weights will be saved.",
+    )
+    parser.add_argument(
+        "--model_type",
+        type=str,
+        default="pixelwise",
+        help="The type of model to train.",
+    )
+
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument(
+        "--colour_map",
+        type=str,
+        help="Semicolon-separated list of RGB triples.",
+    )
+    group.add_argument(
+        "--colour_map_file",
+        type=str,
+        help="Path to a file listing RGB triples.",
+    )
+
+    args = parser.parse_args()
+
+    if args.colour_map_file:
+        colours = utils.parse_colours_from_file(args.colour_map_file)
+    else:
+        colours = utils.parse_colours_from_string(args.colour_map)
+    colour_map = {i: rgb for i, rgb in enumerate(colours)}
+
+    train.train_model(
+        image_dir=args.image_dir,
+        label_dir=args.label_dir,
+        output_dir=args.output_dir,
+        colour_map=colour_map,
+        model_type=args.model_type,
+    )

rgb_to_segmentation-0.0.2/rgb_to_segmentation/models/pixelwise_classifier.py

@@ -0,0 +1,57 @@
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class PixelwiseClassifier(pl.LightningModule):
+    step = 0
+    val_step = 0
+
+    def __init__(self, input_dim, hidden_dim, output_dim):
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.output_dim = output_dim
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, output_dim),
+        )
+        self.softmax = nn.Softmax(dim=-1)
+        self.loss_fn = nn.CrossEntropyLoss()
+
+    def forward(self, x: torch.Tensor):
+        x = self.net(x)
+        return self.softmax(x)
+
+    def training_step(self, batch, batch_idx):
+        sample, target = batch
+        probs = self(sample)
+        target = F.one_hot(target.squeeze(-1), num_classes=self.output_dim).to(
+            dtype=sample.dtype
+        )
+        loss = self.loss_fn(probs, target)
+        self.log("train_loss", loss)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        sample, target = batch
+        probs = self(sample)
+        target = F.one_hot(target.squeeze(-1), num_classes=self.output_dim).to(
+            dtype=sample.dtype
+        )
+        loss = self.loss_fn(probs, target)
+        self.log("val_loss", loss, prog_bar=True)
+        return loss
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters())
+
+    def image_to_batch(self, x: torch.Tensor):
+        if len(x.shape) == 3:
+            x = x.unsqueeze(0)
+        return x.permute(0, 2, 3, 1).reshape(-1, x.shape[1]).to(self.device)
+
+    def batch_to_image(self, batch: torch.Tensor, height: int, width: int):
+        return batch.reshape(1, height, width, 3).permute(0, 3, 1, 2)

rgb_to_segmentation-0.0.2/rgb_to_segmentation/nn.py

@@ -0,0 +1,110 @@
+import os
+
+import torch
+
+from PIL import Image
+from torchvision.io import read_image
+
+from .models.pixelwise_classifier import PixelwiseClassifier
+
+
+def load_model(model_path: str):
+    print(model_path)
+    model = PixelwiseClassifier.load_from_checkpoint(checkpoint_path=model_path)
+    model.eval()
+
+    return model
+
+
+def map_int_to_rgb(indexed_image: torch.Tensor, colour_map: dict):
+    h, w = indexed_image.shape
+    rgb_image = torch.zeros((3, h, w), dtype=torch.uint8)
+
+    for idx, rgb in colour_map.items():
+        mask = indexed_image == idx
+        for c in range(3):
+            rgb_image[c][mask] = rgb[c]
+
+    return rgb_image
+
+
+def run_inference(
+    input_dir: str,
+    output_dir: str,
+    inplace: bool = False,
+    model_path: str = None,
+    colour_map: dict = None,
+    exts: str = ".png,.jpg,.jpeg,.tiff,.bmp,.gif",
+    name_filter: str = "",
+    output_type: str = "rgb",
+):
+    """
+    Run neural network inference on segmentation images.
+
+    Args:
+        input_dir (str): Path to input directory containing images.
+        output_dir (str, optional): Directory where output images will be written.
+        inplace (bool): Overwrite input images in place.
+        model_path (str): Path to the trained model file.
+        colour_map (dict): Mapping from class indices to RGB tuples.
+        exts (str): Comma-separated list of allowed image extensions.
+        name_filter (str): Only process files whose name contains this substring.
+    """
+    if not inplace and output_dir is None:
+        raise ValueError("Either output_dir must be provided or inplace must be True")
+
+    if model_path is None:
+        raise ValueError("model_path must be provided")
+
+    if colour_map is None:
+        raise ValueError("colour_map must be provided")
+
+    model = load_model(model_path)
+
+    exts_list = [e.lower().strip() for e in exts.split(",")]
+    out_dir = output_dir if not inplace else input_dir
+
+    if not inplace:
+        os.makedirs(out_dir, exist_ok=True)
+
+    print(f"Running inference: input={input_dir} -> output={out_dir}, output_type={output_type}")
+
+    for root, dirs, files in os.walk(input_dir):
+        rel = os.path.relpath(root, input_dir)
+        out_root = os.path.join(out_dir, rel) if not inplace else root
+
+        os.makedirs(out_root, exist_ok=True)
+
+        for fname in files:
+            if not any(fname.lower().endswith(e) for e in exts_list):
+                continue
+            if name_filter and name_filter not in fname:
+                continue
+
+            in_path = os.path.join(root, fname)
+            out_path = os.path.join(out_root, fname)
+
+            try:
+                img = read_image(in_path, mode="RGB").float() / 127.5 - 1.0  # Normalize
+                h, w = img.shape[1], img.shape[2]
+
+                batch = model.image_to_batch(img)
+
+                with torch.no_grad():
+                    probs = model(batch)
+                    predicted = torch.argmax(probs, dim=-1).reshape(h, w)
+
+                if output_type == "rgb":
+                    rgb_image = map_int_to_rgb(predicted, colour_map)
+                    pil_image = Image.fromarray(rgb_image.permute(1, 2, 0).numpy())
+                    pil_image.save(out_path)
+                elif output_type == "index":
+                    pil_image = Image.fromarray(predicted.cpu().numpy().astype("uint8"), mode="L")
+                    pil_image.save(out_path)
+                else:
+                    raise ValueError("output_type must be 'rgb' or 'index'")
+
+            except Exception as e:
+                print(f"Skipping {in_path}: {e}")
+
+    print("Inference done.")

rgb_to_segmentation-0.0.2/rgb_to_segmentation/train.py

@@ -0,0 +1,183 @@
+import os
+
+import torch
+
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
+from torch.utils.data import DataLoader, Dataset, random_split
+from torchvision.io import read_image
+
+from .nn import PixelwiseClassifier
+
+
+VALIDATION_FRACTION = 0.2
+
+
+def map_colour_to_int(sample, colour_map):
+    _, H, W = sample.shape
+    image_array = torch.zeros((1, H, W))
+
+    for idx, rgb in colour_map.items():
+        mask = torch.all(sample == torch.tensor(rgb).view(3, 1, 1), dim=0)
+        image_array[0][mask] = idx
+
+    return image_array
+
+
+class SegMaskDataset(Dataset):
+    def __init__(self, paired_filenames, colour_map, model):
+        self.paired_filenames = paired_filenames
+        self.colour_map = colour_map
+        self.to_batch = model.image_to_batch
+
+    def __len__(self):
+        return len(self.paired_filenames)
+
+    def __getitem__(self, index):
+        sample_path = self.paired_filenames[index]["sample"]
+        target_path = self.paired_filenames[index]["target"]
+
+        sample = (
+            read_image(sample_path, mode="RGB") / 127.5
+        ) - 1.0  # Normalize to [-1, 1]
+        target = read_image(target_path, mode="RGB")
+        target = map_colour_to_int(target, self.colour_map)
+        sample = self.to_batch(sample)
+        target = self.to_batch(target).to(torch.long)
+
+        return (sample, target)
+
+
+def get_png_basenames(directory: str) -> list[str]:
+    return [
+        os.path.splitext(filename)[0]
+        for filename in os.listdir(directory)
+        if filename.endswith(".png")
+    ]
+
+
+def get_paired_filenames(
+    image_dir: str,
+    label_dir: str,
+    noisy_basenames: list[str],
+    label_basenames: list[str],
+    training_label_basenames: list[str],
+) -> list[dict]:
+    training_paired_filenames = []
+    val_paired_filenames = []
+
+    for noisy_image_name in noisy_basenames:
+        targets = [
+            target_file
+            for target_file in label_basenames
+            if target_file in noisy_image_name
+        ]
+
+        if len(targets) > 1:
+            raise Exception(
+                f"Multiple target files exist for noisy file {noisy_image_name}: {targets}."
+            )
+        elif len(targets) < 1:
+            print(
+                f"WARNING: No target found for noisy file {noisy_image_name}. Discarding."
+            )
+        else:
+            target_filename = targets[0]
+
+            if target_filename in training_label_basenames:
+                training_paired_filenames.append(
+                    {
+                        "sample": f"{image_dir}/{noisy_image_name}.png",
+                        "target": f"{label_dir}/{target_filename}.png",
+                    }
+                )
+            else:
+                val_paired_filenames.append(
+                    {
+                        "sample": f"{image_dir}/{noisy_image_name}.png",
+                        "target": f"{label_dir}/{target_filename}.png",
+                    }
+                )
+
+    return training_paired_filenames, val_paired_filenames
+
+
+def get_dataloaders(
+    image_dir, label_dir, colour_map, model
+) -> tuple[DataLoader, DataLoader]:
+    noisy_basenames = get_png_basenames(image_dir)
+    label_basenames = get_png_basenames(label_dir)
+
+    training_label_basenames, val_label_basenames = random_split(
+        label_basenames, [1 - VALIDATION_FRACTION, VALIDATION_FRACTION]
+    )
+
+    training_paired_filenames, val_paired_filenames = get_paired_filenames(
+        image_dir,
+        label_dir,
+        noisy_basenames,
+        label_basenames,
+        training_label_basenames,
+    )
+
+    train_dataset = SegMaskDataset(training_paired_filenames, colour_map, model)
+    val_dataset = SegMaskDataset(val_paired_filenames, colour_map, model)
+
+    train_dataloader = DataLoader(
+        dataset=train_dataset,
+        batch_size=1,
+        shuffle=True,
+        num_workers=2,
+    )
+    val_dataloader = DataLoader(
+        dataset=val_dataset,
+        batch_size=1,
+        shuffle=True,
+        num_workers=2,
+    )
+
+    return train_dataloader, val_dataloader
+
+
+def get_model(model_type, colour_map):
+    num_classes = len(colour_map.keys())
+
+    if model_type == "pixelwise":
+        return PixelwiseClassifier(input_dim=3, hidden_dim=32, output_dim=num_classes)
+
+    raise ValueError(f"Invalid model type: {model_type}")
+
+
+def train_model(
+    image_dir: str,
+    label_dir: str,
+    output_dir: str,
+    colour_map: dict,
+    model_type: str = "pixelwise",
+):
+    """
+    Train a neural network model for segmentation cleaning.
+
+    Args:
+        image_dir (str): Path to directory containing noisy images.
+        label_dir (str): Path to directory containing target RGB labels.
+        output_dir (str): Directory where model weights will be saved.
+        colour_map (dict): Mapping from class indices to RGB tuples.
+        model_type (str): The type of model to train.
+    """
+    model = get_model(model_type, colour_map)
+    train_dataloader, val_dataloader = get_dataloaders(
+        image_dir, label_dir, colour_map, model
+    )
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    early_stop = EarlyStopping(monitor="val_loss", mode="min", verbose=True)
+    checkpoint = ModelCheckpoint(
+        monitor="val_loss", mode="min", save_top_k=1, dirpath=output_dir
+    )
+    trainer = Trainer(max_epochs=100, callbacks=[early_stop, checkpoint])
+
+    trainer.fit(
+        model, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader
+    )

rgb_to_segmentation-0.0.2/rgb_to_segmentation/utils.py

@@ -0,0 +1,40 @@
+from typing import List, Tuple
+
+
+def parse_colours_from_string(colours_str: str) -> List[Tuple[int, int, int]]:
+    """Parse a semicolon-separated string of RGB triples into a list of tuples."""
+
+    parts = [p.strip() for p in colours_str.split(";") if p.strip()]
+    colours = []
+
+    for p in parts:
+        rgb = tuple(int(x) for x in p.split(","))
+
+        if len(rgb) != 3:
+            raise ValueError(f"Invalid colour triple: {p}")
+        colours.append(rgb)
+
+    return colours
+
+
+def parse_colours_from_file(path: str) -> List[Tuple[int, int, int]]:
+    """Parse a file with one RGB triple per line into a list of tuples."""
+
+    colours = []
+
+    with open(path, "r") as f:
+        for line in f:
+            line = line.strip()
+            if not line or line.startswith("#"):
+                continue
+
+            rgb = tuple(int(x) for x in line.split(","))
+            if len(rgb) != 3:
+                raise ValueError(f"Invalid colour triple in file {path}: {line}")
+
+            colours.append(rgb)
+
+    if not colours:
+        raise ValueError(f"No colours found in file: {path}")
+
+    return colours