File size: 3,587 Bytes
d176ecd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | #!/usr/bin/env python3
"""Fine-tune YOLO11n (COCO-pretrained) as a single-class face detector on converted WIDER FACE.
Adds video-like degradation (motion blur + JPEG compression) on top of ultralytics' default augmentation,
implemented with plain cv2 so no albumentations / second OpenCV build is needed.
Usage:
python train.py # full run (40 epochs)
python train.py --fraction 0.02 --epochs 1 --name smoke # quick pipeline check
"""
import argparse
import random
from pathlib import Path
import cv2
import numpy as np
from ultralytics import YOLO
from ultralytics.data import augment
def _motion_blur(img, max_k=15):
k = random.choice(range(3, max_k + 1, 2))
kernel = np.zeros((k, k), np.float32)
kernel[k // 2, :] = 1.0
rot = cv2.getRotationMatrix2D((k / 2 - 0.5, k / 2 - 0.5), random.uniform(0, 180), 1.0)
kernel = cv2.warpAffine(kernel, rot, (k, k))
s = kernel.sum()
return img if s == 0 else cv2.filter2D(img, -1, kernel / s)
def _jpeg(img, qmin=20, qmax=90):
ok, enc = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, random.randint(qmin, qmax)])
return cv2.imdecode(enc, cv2.IMREAD_COLOR) if ok else img
def install_video_degradation(p_blur=0.3, p_jpeg=0.3):
"""Wrap ultralytics' Albumentations step (runs after mosaic/affine, train only) with cv2 degradations."""
orig = augment.Albumentations.__call__
def call(self, labels):
img = labels["img"]
if random.random() < p_blur:
img = _motion_blur(img)
if random.random() < p_jpeg:
img = _jpeg(img)
labels["img"] = img
return orig(self, labels)
augment.Albumentations.__call__ = call
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--data", default="dataset/data.yaml")
ap.add_argument("--model", default="yolo11n.pt")
ap.add_argument("--epochs", type=int, default=40)
ap.add_argument("--batch", type=int, default=64)
ap.add_argument("--imgsz", type=int, default=640)
ap.add_argument("--workers", type=int, default=16)
ap.add_argument("--fraction", type=float, default=1.0, help="fraction of train set (smoke tests)")
ap.add_argument("--cache", default="disk", help="ram | disk | False")
ap.add_argument("--project", default=str(Path(__file__).resolve().parent / "runs"))
ap.add_argument("--name", default="face_yolo11n")
ap.add_argument("--no-degrade", action="store_true")
ap.add_argument("--resume", default=None,
help="path to a stopped run's last.pt; continues it and finishes at --epochs "
"(the LR schedule is recomputed for the new total)")
a = ap.parse_args()
if not a.no_degrade:
install_video_degradation()
if a.resume:
import torch
ck = torch.load(a.resume, map_location="cpu", weights_only=False)
ck["train_args"]["epochs"] = a.epochs
torch.save(ck, a.resume)
model = YOLO(a.resume)
model.train(resume=True)
else:
model = YOLO(a.model)
model.train(
data=a.data, epochs=a.epochs, imgsz=a.imgsz, batch=a.batch, workers=a.workers,
fraction=a.fraction, cache=False if a.cache == "False" else a.cache, amp=True,
close_mosaic=5, patience=0, project=a.project, name=a.name, exist_ok=True,
plots=False,
)
m = model.val(data=a.data, imgsz=a.imgsz, batch=a.batch)
print(f"RESULT mAP50={m.box.map50:.4f} mAP50-95={m.box.map:.4f} P={m.box.mp:.4f} R={m.box.mr:.4f}")
if __name__ == "__main__":
main()
|