165 lines
6.6 KiB
Python
165 lines
6.6 KiB
Python
"""Generate a synthetic YOLO-seg dataset for MOSAIC detection.
|
||
|
||
Takes a folder of CLEAN (uncensored) images — anime frames work best for the
|
||
anime domain — and produces censored copies with random mosaic regions plus
|
||
matching YOLO segmentation labels (class 0 = mosaic). Some outputs are left
|
||
clean (negatives / background) so the model learns what is *not* mosaic.
|
||
|
||
The model only needs to recognise mosaic *texture*, so random placement is fine
|
||
(we detect already-applied mosaic anywhere, not "where to censor").
|
||
|
||
Output layout (Ultralytics format):
|
||
<out>/images/train/*.jpg <out>/labels/train/*.txt
|
||
<out>/images/val/*.jpg <out>/labels/val/*.txt
|
||
<out>/data.yaml
|
||
|
||
Usage:
|
||
python scripts/training/gen_mosaic_dataset.py --input clean_frames --output dataset_mosaic
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import random
|
||
from pathlib import Path
|
||
|
||
import cv2
|
||
import numpy as np
|
||
|
||
IMG_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
|
||
|
||
|
||
# --- unicode-safe IO (self-contained, no hvideotool import needed) ------------
|
||
def imread(path: Path) -> "np.ndarray | None":
|
||
data = np.fromfile(str(path), dtype=np.uint8)
|
||
if data.size == 0:
|
||
return None
|
||
return cv2.imdecode(data, cv2.IMREAD_COLOR)
|
||
|
||
|
||
def imwrite(path: Path, img: np.ndarray, quality: int = 92) -> None:
|
||
ok, buf = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, quality])
|
||
if ok:
|
||
buf.tofile(str(path))
|
||
|
||
|
||
# --- mosaic + polygon ---------------------------------------------------------
|
||
def pixelate_region(img: np.ndarray, poly: np.ndarray, tile: int) -> None:
|
||
"""Apply mosaic inside the polygon (in place). Clamps to image bounds."""
|
||
H, W = img.shape[:2]
|
||
x, y, w, h = cv2.boundingRect(poly)
|
||
x, y = max(0, x), max(0, y)
|
||
x2, y2 = min(x + w, W), min(y + h, H)
|
||
w, h = x2 - x, y2 - y
|
||
if w < 1 or h < 1:
|
||
return
|
||
roi = img[y:y2, x:x2]
|
||
small = cv2.resize(roi, (max(1, w // tile), max(1, h // tile)), interpolation=cv2.INTER_LINEAR)
|
||
mosaic = cv2.resize(small, (w, h), interpolation=cv2.INTER_NEAREST)
|
||
mask = np.zeros((h, w), np.uint8)
|
||
cv2.fillPoly(mask, [poly - [x, y]], 255)
|
||
roi[mask > 0] = mosaic[mask > 0]
|
||
|
||
|
||
def make_region(W: int, H: int, area_min: float, area_max: float, shape: str) -> np.ndarray:
|
||
"""Return an Nx2 int polygon for a random mosaic region within the image."""
|
||
area = random.uniform(area_min, area_max) * W * H
|
||
aspect = random.uniform(0.5, 2.0)
|
||
w = int(min(W * 0.9, max(24, (area * aspect) ** 0.5)))
|
||
h = int(min(H * 0.9, max(24, area / max(1, w))))
|
||
x = random.randint(0, max(0, W - w))
|
||
y = random.randint(0, max(0, H - h))
|
||
if shape == "ellipse":
|
||
cx, cy = x + w // 2, y + h // 2
|
||
pts = cv2.ellipse2Poly((cx, cy), (w // 2, h // 2), random.randint(0, 180), 0, 360, 20)
|
||
pts[:, 0] = np.clip(pts[:, 0], 0, W - 1)
|
||
pts[:, 1] = np.clip(pts[:, 1], 0, H - 1)
|
||
return pts.astype(np.int32)
|
||
return np.array([[x, y], [x + w, y], [x + w, y + h], [x, y + h]], np.int32)
|
||
|
||
|
||
def poly_to_label(poly: np.ndarray, W: int, H: int) -> str:
|
||
coords = []
|
||
for px, py in poly:
|
||
coords.append(f"{np.clip(px / W, 0, 1):.6f}")
|
||
coords.append(f"{np.clip(py / H, 0, 1):.6f}")
|
||
return "0 " + " ".join(coords)
|
||
|
||
|
||
def main() -> None:
|
||
ap = argparse.ArgumentParser(description="Synthetic mosaic YOLO-seg dataset generator")
|
||
ap.add_argument("--input", required=True, help="folder of clean (uncensored) images")
|
||
ap.add_argument("--output", required=True, help="output dataset folder")
|
||
ap.add_argument("--variants", type=int, default=3, help="augmented copies per source image")
|
||
ap.add_argument("--val-split", type=float, default=0.15)
|
||
ap.add_argument("--neg-frac", type=float, default=0.2, help="fraction of outputs left clean")
|
||
ap.add_argument("--min-regions", type=int, default=1)
|
||
ap.add_argument("--max-regions", type=int, default=3)
|
||
ap.add_argument("--tile-min", type=int, default=6)
|
||
ap.add_argument("--tile-max", type=int, default=22)
|
||
ap.add_argument("--area-min", type=float, default=0.02)
|
||
ap.add_argument("--area-max", type=float, default=0.22)
|
||
ap.add_argument("--max-dim", type=int, default=1280, help="downscale clean images larger than this")
|
||
ap.add_argument("--shapes", default="rect,ellipse")
|
||
ap.add_argument("--seed", type=int, default=0)
|
||
args = ap.parse_args()
|
||
|
||
random.seed(args.seed)
|
||
np.random.seed(args.seed)
|
||
shapes = [s.strip() for s in args.shapes.split(",") if s.strip()]
|
||
|
||
sources = sorted(p for p in Path(args.input).rglob("*") if p.suffix.lower() in IMG_EXTS)
|
||
if not sources:
|
||
raise SystemExit(f"Не найдено изображений в {args.input}")
|
||
random.shuffle(sources)
|
||
n_val = max(1, int(len(sources) * args.val_split))
|
||
val_set = set(sources[:n_val])
|
||
|
||
out = Path(args.output)
|
||
for split in ("train", "val"):
|
||
(out / "images" / split).mkdir(parents=True, exist_ok=True)
|
||
(out / "labels" / split).mkdir(parents=True, exist_ok=True)
|
||
|
||
counts = {"train": 0, "val": 0, "neg": 0, "pos": 0}
|
||
for src in sources:
|
||
img0 = imread(src)
|
||
if img0 is None:
|
||
continue
|
||
H0, W0 = img0.shape[:2]
|
||
scale = min(1.0, args.max_dim / max(H0, W0))
|
||
if scale < 1.0:
|
||
img0 = cv2.resize(img0, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
|
||
H, W = img0.shape[:2]
|
||
split = "val" if src in val_set else "train"
|
||
|
||
for v in range(args.variants):
|
||
img = img0.copy()
|
||
lines: list[str] = []
|
||
if random.random() >= args.neg_frac:
|
||
for _ in range(random.randint(args.min_regions, args.max_regions)):
|
||
shape = random.choice(shapes)
|
||
poly = make_region(W, H, args.area_min, args.area_max, shape)
|
||
tile = random.randint(args.tile_min, args.tile_max)
|
||
pixelate_region(img, poly, tile)
|
||
lines.append(poly_to_label(poly, W, H))
|
||
stem = f"{src.stem}_{v:02d}"
|
||
imwrite(out / "images" / split / f"{stem}.jpg", img)
|
||
(out / "labels" / split / f"{stem}.txt").write_text("\n".join(lines), encoding="utf-8")
|
||
counts[split] += 1
|
||
counts["neg" if not lines else "pos"] += 1
|
||
|
||
(out / "data.yaml").write_text(
|
||
f"path: {out.resolve().as_posix()}\n"
|
||
"train: images/train\n"
|
||
"val: images/val\n"
|
||
"names:\n 0: mosaic\n",
|
||
encoding="utf-8",
|
||
)
|
||
print(f"Готово: train={counts['train']} val={counts['val']} "
|
||
f"(с мозаикой={counts['pos']}, чистых={counts['neg']})")
|
||
print(f"data.yaml: {(out / 'data.yaml').resolve()}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|