train_utility/scripts/create_yolo26_smoke_datasets.py

"""Create tiny deterministic datasets for all YOLO tasks supported by the WebUI."""

from __future__ import annotations

import argparse
import math
import shutil
from pathlib import Path

from PIL import Image, ImageDraw


IMAGE_SIZE = 128
TRAIN_IMAGES = 6
VAL_IMAGES = 2


def image_geometry(index: int) -> tuple[int, tuple[int, int, int, int]]:
    class_id = index % 2
    offset = (index % 3) * 5
    box = (26 + offset, 29, 93 + offset, 98)
    return class_id, box


def make_image(path: Path, index: int, *, rotated: bool = False) -> None:
    class_id, box = image_geometry(index)
    colors = ((225, 72, 72), (55, 145, 225))
    image = Image.new("RGB", (IMAGE_SIZE, IMAGE_SIZE), (238, 241, 245))
    draw = ImageDraw.Draw(image)
    if rotated:
        cx, cy = 64 + (index % 3) * 3, 64
        half_w, half_h = 37, 23
        angle = math.radians(15 if class_id == 0 else -15)
        points = []
        for x, y in ((-half_w, -half_h), (half_w, -half_h), (half_w, half_h), (-half_w, half_h)):
            points.append(
                (
                    cx + x * math.cos(angle) - y * math.sin(angle),
                    cy + x * math.sin(angle) + y * math.cos(angle),
                )
            )
        draw.polygon(points, fill=colors[class_id], outline=(25, 25, 25), width=2)
    else:
        draw.rectangle(box, fill=colors[class_id], outline=(25, 25, 25), width=2)
    path.parent.mkdir(parents=True, exist_ok=True)
    image.save(path)


def normalized_box(box: tuple[int, int, int, int]) -> tuple[float, float, float, float]:
    left, top, right, bottom = box
    return (
        (left + right) / 2 / IMAGE_SIZE,
        (top + bottom) / 2 / IMAGE_SIZE,
        (right - left) / IMAGE_SIZE,
        (bottom - top) / IMAGE_SIZE,
    )


def write_yaml(root: Path, task: str, extra: str = "") -> None:
    yaml_text = (
        f"path: {root.resolve()}\n"
        "train: images/train\n"
        "val: images/val\n"
        "names:\n"
        "  0: red_shape\n"
        "  1: blue_shape\n"
        f"{extra}"
    )
    (root / f"{task}.yaml").write_text(yaml_text, encoding="utf-8")


def create_detection_style(base: Path, task: str) -> None:
    root = base / task
    for split, count in (("train", TRAIN_IMAGES), ("val", VAL_IMAGES)):
        for index in range(count):
            sample = index if split == "train" else index + TRAIN_IMAGES
            image_path = root / "images" / split / f"sample_{sample:02d}.png"
            label_path = root / "labels" / split / f"sample_{sample:02d}.txt"
            make_image(image_path, sample, rotated=task == "obb")
            class_id, box = image_geometry(sample)
            cx, cy, width, height = normalized_box(box)

            if task == "detect":
                label = f"{class_id} {cx:.6f} {cy:.6f} {width:.6f} {height:.6f}\n"
            elif task == "segment":
                left, top, right, bottom = (value / IMAGE_SIZE for value in box)
                label = (
                    f"{class_id} {left:.6f} {top:.6f} {right:.6f} {top:.6f} "
                    f"{right:.6f} {bottom:.6f} {left:.6f} {bottom:.6f}\n"
                )
            elif task == "pose":
                class_id = 0
                points = (
                    (cx, cy - height * 0.25),
                    (cx - width * 0.25, cy),
                    (cx + width * 0.25, cy),
                    (cx, cy + height * 0.25),
                )
                keypoints = " ".join(f"{x:.6f} {y:.6f} 2" for x, y in points)
                label = f"{class_id} {cx:.6f} {cy:.6f} {width:.6f} {height:.6f} {keypoints}\n"
            elif task == "obb":
                angle = math.radians(15 if class_id == 0 else -15)
                center_x, center_y = 64 + (sample % 3) * 3, 64
                half_w, half_h = 37, 23
                points = []
                for x, y in ((-half_w, -half_h), (half_w, -half_h), (half_w, half_h), (-half_w, half_h)):
                    px = center_x + x * math.cos(angle) - y * math.sin(angle)
                    py = center_y + x * math.sin(angle) + y * math.cos(angle)
                    points.extend((px / IMAGE_SIZE, py / IMAGE_SIZE))
                label = f"{class_id} " + " ".join(f"{value:.6f}" for value in points) + "\n"
            else:
                raise ValueError(f"Unsupported task: {task}")

            label_path.parent.mkdir(parents=True, exist_ok=True)
            label_path.write_text(label, encoding="utf-8")

    if task == "pose":
        write_yaml(
            root,
            task,
            extra=(
                "kpt_shape: [4, 3]\n"
                "flip_idx: [0, 2, 1, 3]\n"
            ),
        )
    else:
        write_yaml(root, task)


def create_classification(base: Path) -> None:
    root = base / "classify"
    for split, count in (("train", TRAIN_IMAGES), ("val", 4)):
        for index in range(count):
            sample = index if split == "train" else index + TRAIN_IMAGES
            class_id = sample % 2
            make_image(root / split / ("red_shape" if class_id == 0 else "blue_shape") / f"sample_{sample:02d}.png", sample)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, default=Path("datasets/yolo26_smoke"))
    parser.add_argument("--force", action="store_true")
    args = parser.parse_args()

    if args.output.exists():
        if not args.force:
            raise SystemExit(f"Dataset already exists: {args.output}; use --force to recreate it")
        shutil.rmtree(args.output)

    for task in ("detect", "segment", "pose", "obb"):
        create_detection_style(args.output, task)
    create_classification(args.output)
    print(args.output.resolve())


if __name__ == "__main__":
    main()