Deploy 3a26184

server/train_plate_chars.py

@@ -0,0 +1,55 @@
+"""
+Trains the Saudi plate character detector (YOLOv8) and writes
+plate_chars.pt. Used by .github/workflows/train-plate-chars.yml (CPU runner,
+so a time budget keeps it inside GitHub's 6-hour job limit) and runnable by
+hand:  python server/train_plate_chars.py <folder with data.yaml> [hours]
+"""
+import glob
+import os
+import shutil
+import sys
+
+import yaml
+from ultralytics import YOLO
+
+
+def fix_yaml(root):
+    yp = glob.glob(os.path.join(root, "**", "data.yaml"), recursive=True)[0]
+    base = os.path.dirname(os.path.abspath(yp))
+    d = yaml.safe_load(open(yp))
+    d["path"] = base
+    for k, cands in (("train", ["train"]), ("val", ["valid", "val"]), ("test", ["test"])):
+        for c in cands:
+            p = os.path.join(base, c, "images")
+            if os.path.isdir(p):
+                d[k] = p
+                break
+        else:
+            d.pop(k, None)
+    d.setdefault("val", d["train"])
+    yaml.safe_dump(d, open(yp, "w"), allow_unicode=True)
+    print(f"classes ({d['nc']}): {d['names']}", flush=True)
+    return yp
+
+
+def main():
+    root = sys.argv[1] if len(sys.argv) > 1 else "data"
+    hours = float(sys.argv[2]) if len(sys.argv) > 2 else 4.5
+    yp = fix_yaml(root)
+    import torch
+    gpu = torch.cuda.is_available()
+    model = YOLO("yolov8s.pt" if gpu else "yolov8n.pt")
+    model.train(data=yp, imgsz=640 if gpu else 512, epochs=80 if gpu else 40, patience=15,
+                batch=16, workers=4, time=hours,          # stop (and keep best) within the budget
+                fliplr=0.0,                               # never mirror characters
+                degrees=8, perspective=0.0005, hsv_v=0.5, name="plate_chars", exist_ok=True)
+    best = sorted(glob.glob("runs/detect/plate_chars*/weights/best.pt"), key=os.path.getmtime)[-1]
+    r = YOLO(best).val(data=yp, imgsz=640 if gpu else 512)
+    print(f"mAP50 = {r.box.map50:.3f}", flush=True)
+    shutil.copy(best, "plate_chars.pt")
+    with open("train_summary.txt", "w") as f:
+        f.write(f"mAP50={r.box.map50:.3f}\nclasses={YOLO(best).names}\n")
+
+
+if __name__ == "__main__":
+    main()