Добавлено описание и документация для HVideoTool, включая функционал, требования, установку и запуск приложения для обнаружения цензуры на изображениях.
This commit is contained in:
Leonid Pershin
@@ -0,0 +1,164 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user