Skip to main content

A Python toolkit to convert between binary segmentation masks and YOLO labels

Project description

Segment Toolkit

A modern, robust, and premium Python package designed to bridge the gap between pixel-level binary/colored segmentation masks and YOLO labels. It provides a bidirectional pipeline for both standard Bounding Boxes (YOLO object detection) and Polygon Coordinates (YOLO instance segmentation format).

Equipped with a flexible image-mask transform pipeline, exception handling, extensive logging, a CLI, and a clean Python API.


Features

  • Bidirectional Conversions:
    • Bounding Boxes: Convert binary masks to YOLO format labels (supports standard axis-aligned or advanced minimum area rotated bounding boxes) and vice versa.
    • Polygon Segmentation: Convert binary or multi-class color-coded masks to YOLO instance segmentation polygon coordinates and vice versa.
  • Transform Pipeline (Compose, Resize, Normalize): Jointly apply spatial and pixel transforms (like resizing and normalizations) on both images and masks.
  • Robust Exception Handling: Prevents crashes on corrupted, missing, or empty files.
  • Dynamic Dataset Matching: Parse Ground Truth CSV or JSON files (supports multiple indicator formats) to automatically map filenames to diagnostic class IDs (e.g. standard ISIC classes).
  • YOLO Dataset Splitting: Automatically partitions images and label files into standard train and test structures and generates the data.yaml configuration file.
  • Visualization Overlays:
    • Overlay bounding boxes onto source images.
    • Blend masks onto source images.
    • Render colored segmentation polygons with class tags onto source images.
  • Dual Interface: Use as a command-line application (segment-toolkit) or import as a Python library (import segment_toolkit).

Installation

1. Standard Installation (via PyPI)

pip install segment-toolkit

2. Local Development Installation

# Clone the repository
git clone https://github.com/zkzkGamal/mask-to-yolo-toolkit.git
cd mask-to-yolo-toolkit

# Install in editable mode
pip install -e .

Command Line Interface (CLI)

The package installs a console script called segment-toolkit.

1. Bounding Box Conversions

Convert Masks to YOLO Bounding Boxes

# Single File
segment-toolkit mask-to-yolo \
  --image images/sample.jpg \
  --mask masks/sample.png \
  --output-txt labels/sample.txt \
  --class-id 4

# Batch Directory
segment-toolkit mask-to-yolo \
  --image-dir images/ \
  --mask-dir masks/ \
  --output-dir labels/ \
  --ground-truth GroundTruth.csv

Options:

  • --rotated: Use rotated minimum area rectangles (cv2.minAreaRect) instead of standard axis-aligned boxes.
  • --resize WIDTH HEIGHT: Set target size for resizing (default: 640 640).

Convert YOLO Bounding Boxes to Masks

# Single File
segment-toolkit yolo-to-mask \
  --label labels/sample.txt \
  --output-mask reconstructed/sample.png

# Batch Directory
segment-toolkit yolo-to-mask \
  --label-dir labels/ \
  --output-dir reconstructed/

Visualize Bounding Boxes

segment-toolkit visualize \
  --image images/sample.jpg \
  --label labels/sample.txt \
  --output visualization.png

2. Polygon Segmentation Conversions

Convert Masks to YOLO Polygon Labels

Supports standard black-and-white masks or multi-class colored masks (requires passing a --classes JSON config mapping colors to class names).

# Single binary mask
segment-toolkit mask-to-polygon \
  --image images/sample.jpg \
  --mask masks/sample.png \
  --output-txt labels/sample.txt \
  --class-id 1

# Batch directory of multi-class colored masks
segment-toolkit mask-to-polygon \
  --image-dir images/ \
  --mask-dir masks/ \
  --output-dir labels/ \
  --classes classes.json

Convert YOLO Polygons back to Masks

# Single file (creates color mask if classes.json is provided)
segment-toolkit polygon-to-mask \
  --label labels/sample.txt \
  --output-mask reconstructed/sample.png \
  --classes classes.json

# Batch directory
segment-toolkit polygon-to-mask \
  --label-dir labels/ \
  --output-dir reconstructed/

Visualize Polygon Segmentation

segment-toolkit visualize-polygon \
  --image images/sample.jpg \
  --label labels/sample.txt \
  --output polygon_overlay.png \
  --classes classes.json

3. Dataset Utilities

Split Dataset

Partitions image-label pairs into YOLO-compliant subfolders (train/, test/) and generates data.yaml:

segment-toolkit split \
  --images images/ \
  --labels labels/ \
  --output dataset/ \
  --ratio 0.8

Validation Overlays

To verify the pipeline, check out the visualization overlays generated by running the test suite under the scripts/ folder:

Bounding Box Overlay (ISIC Skin Lesion)

ISIC Bounding Box Validation Overlay

Bounding Box Overlay (Plant Disease)

Plant Leaf Bounding Box Validation Overlay

Polygon Segmentation Overlay (ISIC Skin Lesion)

Polygon Segmentation Validation Overlay

Reconstructed Mask Overlay (Multi-class Color Mask)

Multi-class Colored Mask Reconstruction


Python API

Programmatically build custom preprocessing and conversion pipelines:

from segment_toolkit import MaskToYoloConverter, YoloToMaskConverter
from segment_toolkit import MaskToPolygonConverter, PolygonToMaskConverter
from segment_toolkit.transforms import Compose, Resize, Normalize

# 1. Image & Mask Joint Transforms Pipeline
transform_pipeline = Compose([
    Resize((640, 640)),
    Normalize()
])

# 2. Bounding Box Converter
bbox_converter = MaskToYoloConverter(target_size=(640, 640), bbox_type="rotated")
bbox_converter.convert_single(
    image_path="images/sample.jpg",
    mask_path="masks/sample.png",
    output_txt_path="labels/sample.txt",
    class_id=0
)

# 3. Multi-class Colored Mask to Polygon Converter
classes = [
    ((255, 0, 0), "lesion_red"),
    ((0, 255, 0), "lesion_green")
]
poly_converter = MaskToPolygonConverter(target_size=(640, 640))
poly_converter.convert_single(
    image_path="images/sample.jpg",
    mask_path="masks/sample_colored.png",
    output_txt_path="labels/sample_poly.txt",
    classes=classes
)

Ground Truth JSON Format Example

To feed multi-class configurations to CLI commands using --classes, prepare a JSON file listing color lists and class names:

[
  [[255, 0, 0], "lesion_red"],
  [[0, 255, 0], "lesion_green"]
]

License

This project is licensed under the MIT License - see the LICENSE file for details.


Author

Zakria Gamal

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

segment_toolkit-1.2.0.tar.gz (22.6 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

segment_toolkit-1.2.0-py3-none-any.whl (28.1 kB view details)

Uploaded Python 3

File details

Details for the file segment_toolkit-1.2.0.tar.gz.

File metadata

  • Download URL: segment_toolkit-1.2.0.tar.gz
  • Upload date:
  • Size: 22.6 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.10.12

File hashes

Hashes for segment_toolkit-1.2.0.tar.gz
Algorithm Hash digest
SHA256 4687637f7905ebaccfce33ffd299240ff954e6d1a8d4118798aa9e6421afacd5
MD5 dac4ead6d10a6304084c402c2e9f76eb
BLAKE2b-256 b0d19b73dff4be64592605aaceb54031d652c1eadd9fc2821603b18c01d3b2a2

See more details on using hashes here.

File details

Details for the file segment_toolkit-1.2.0-py3-none-any.whl.

File metadata

File hashes

Hashes for segment_toolkit-1.2.0-py3-none-any.whl
Algorithm Hash digest
SHA256 a5919ca1f4131f9757bc659ea9cc207c030c095133217f67f6e7be7e9cd8f72e
MD5 dc366870a61b44a33937f630886f6411
BLAKE2b-256 2a4b0055fe8513dc10e51b9ea6cfcdc3805a359d75c5b9650f90b5d6bd4e33e1

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page