Добавлено описание и документация для HVideoTool, включая функционал, требования, установку и запуск приложения для обнаружения цензуры на изображениях.
This commit is contained in:
Leonid Pershin
@@ -0,0 +1,50 @@
|
||||
# Обучение детектора мозаики (аниме)
|
||||
|
||||
Готовых публичных моделей детекции мозаики для рисованного/аниме контента нет,
|
||||
поэтому обучаем свою на **синтетике**: берём чистые (без цензуры) кадры, случайно
|
||||
накладываем мозаику и получаем разметку автоматически.
|
||||
|
||||
## 1. Подготовить чистые изображения
|
||||
|
||||
Сложите кадры **без цензуры** в папку (чем разнообразнее, тем лучше; от ~500
|
||||
картинок для черновой модели, тысячи — для нормальной). Кадры можно нарезать из
|
||||
чистых видео — например, любым плеером или нашим извлечением (папка проекта
|
||||
`frames/`).
|
||||
|
||||
## 2. Сгенерировать датасет
|
||||
|
||||
```powershell
|
||||
.\.venv\Scripts\Activate.ps1
|
||||
python scripts\training\gen_mosaic_dataset.py --input C:\clean_frames --output dataset_mosaic --variants 3
|
||||
```
|
||||
|
||||
Получите `dataset_mosaic/` с `images/`, `labels/` и `data.yaml`
|
||||
(класс `0: mosaic`, формат сегментации YOLO). Часть кадров остаётся чистыми
|
||||
(негативы) — это снижает ложные срабатывания.
|
||||
|
||||
Полезные флаги: `--neg-frac 0.2`, `--tile-min/--tile-max` (размер плиток мозаики),
|
||||
`--area-min/--area-max` (доля площади под мозаику), `--shapes rect,ellipse`.
|
||||
|
||||
## 3. Обучить
|
||||
|
||||
```powershell
|
||||
pip install -e ".[yolo]"
|
||||
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 # CUDA
|
||||
python scripts\training\train_mosaic.py --data dataset_mosaic\data.yaml --epochs 100
|
||||
```
|
||||
|
||||
Веса появятся в `runs/segment/mosaic/weights/best.pt`.
|
||||
|
||||
## 4. Подключить в приложении
|
||||
|
||||
**Параметры → Детектор: YOLO → Обзор…** и выбрать `best.pt`, затем
|
||||
**Файл → Детектировать заново**. Класс `mosaic` маппится в `CensorType.MOSAIC`,
|
||||
маски превращаются в контуры обводки.
|
||||
|
||||
## Советы по качеству
|
||||
|
||||
- Разнообразьте источник (разные стили, сцены, освещение).
|
||||
- Варьируйте размер плиток и областей (флаги выше) — мозаика в реальности разная.
|
||||
- Доля негативов 15–25 % обычно хорошо снижает ложные срабатывания.
|
||||
- GPU NVIDIA сильно ускоряет; на CPU обучение очень медленное.
|
||||
- Дальше можно добавить классы `bar` (плашки) — генератор легко расширить.
|
||||
@@ -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()
|
||||
@@ -0,0 +1,52 @@
|
||||
"""Train a YOLO11-seg mosaic detector on a generated synthetic dataset.
|
||||
|
||||
Prereqs: pip install -e ".[yolo]" plus PyTorch (CUDA build recommended — see README).
|
||||
|
||||
Usage:
|
||||
python scripts/training/train_mosaic.py --data dataset_mosaic/data.yaml --epochs 100
|
||||
|
||||
The resulting weights (runs/segment/<name>/weights/best.pt) drop straight into the
|
||||
app: Параметры → Детектор «YOLO» → выбрать этот .pt. The class is named "mosaic",
|
||||
which YoloDetector maps to CensorType.MOSAIC.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
|
||||
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser(description="Train YOLO11-seg mosaic detector")
|
||||
ap.add_argument("--data", required=True, help="path to data.yaml from the generator")
|
||||
ap.add_argument("--model", default="yolo11n-seg.pt", help="base model (n/s/m...-seg)")
|
||||
ap.add_argument("--epochs", type=int, default=100)
|
||||
ap.add_argument("--imgsz", type=int, default=640)
|
||||
ap.add_argument("--batch", default="-1", help="batch size (-1 = auto)")
|
||||
ap.add_argument("--device", default=None, help="cuda / 0 / cpu (default: auto)")
|
||||
ap.add_argument("--name", default="mosaic", help="run name under the project dir")
|
||||
ap.add_argument("--project", default=None, help="output dir for runs (default: runs/segment)")
|
||||
args = ap.parse_args()
|
||||
|
||||
try:
|
||||
from ultralytics import YOLO
|
||||
except ImportError as exc: # pragma: no cover
|
||||
raise SystemExit('Не установлен ultralytics: pip install -e ".[yolo]"') from exc
|
||||
|
||||
batch = int(args.batch) if str(args.batch).lstrip("-").isdigit() else args.batch
|
||||
model = YOLO(args.model)
|
||||
results = model.train(
|
||||
data=args.data,
|
||||
epochs=args.epochs,
|
||||
imgsz=args.imgsz,
|
||||
batch=batch,
|
||||
device=args.device,
|
||||
name=args.name,
|
||||
project=args.project,
|
||||
)
|
||||
save_dir = getattr(results, "save_dir", "runs/segment/" + args.name)
|
||||
print(f"\nГотово. Веса: {save_dir}/weights/best.pt")
|
||||
print("Подключите их в приложении: Параметры → Детектор YOLO → выбрать best.pt")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user