"""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): /images/train/*.jpg /labels/train/*.txt /images/val/*.jpg /labels/val/*.txt /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()