Spaces:

Divya499
/

Automated-Pavement-Distress-System

Sleeping

File size: 9,118 Bytes
"""
STEP 3 - Verify Dataset + Train / Resume YOLOv11n or Train YOLOv11s
=====================================================================
HOW TO USE:
  Verify only      : python train.py --verify
  Fresh train n    : python train.py --train
  Resume n 50 more : python train.py --resume
  Train s          : python train.py --trains
"""

import argparse
import yaml
from pathlib import Path
from collections import Counter


DATA_YAML   = "data/combined/data.yaml"
OUTPUT_DIR  = "runs/pothole_v1"
WEIGHTS_N   = "yolo11n.pt"
WEIGHTS_S   = "yolo11s.pt"

# ⚠️ Yeh path tera last.pt ka hai — check kar sahi hai ya nahi
RESUME_CKPT = r"C:\Users\HP\runs\detect\runs\pothole_v1\yolo11n_pothole\weights\last.pt"


# -------------------------------------------------------
# VERIFY
# -------------------------------------------------------

def verify_dataset():
    print("=" * 55)
    print("  Dataset Verification")
    print("=" * 55)

    yaml_path = Path(DATA_YAML)
    if not yaml_path.exists():
        print(f"\n❌ data.yaml not found at: {DATA_YAML}")
        return False

    with open(yaml_path, "r") as f:
        data = yaml.safe_load(f)

    dataset_root = Path(data["path"])
    class_names  = data["names"]

    print(f"\n  Dataset root : {dataset_root}")
    print(f"  Classes      : {class_names}")
    print(f"  Num classes  : {data['nc']}\n")

    all_ok = True

    for split in ["train", "val", "test"]:
        img_dir = dataset_root / split / "images"
        lbl_dir = dataset_root / split / "labels"

        img_count = len(list(img_dir.glob("*"))) if img_dir.exists() else 0
        lbl_count = len(list(lbl_dir.glob("*.txt"))) if lbl_dir.exists() else 0

        class_counts = Counter()
        empty_labels = 0
        if lbl_dir.exists():
            for lbl_file in lbl_dir.glob("*.txt"):
                content = lbl_file.read_text().strip()
                if not content:
                    empty_labels += 1
                    continue
                for line in content.splitlines():
                    parts = line.strip().split()
                    if parts:
                        class_counts[int(parts[0])] += 1

        status = "✅" if img_count > 0 and lbl_count > 0 else "❌"
        if img_count == 0:
            all_ok = False

        print(f"  {status} {split.upper():6s} | Images: {img_count:5d} | "
              f"Labels: {lbl_count:5d} | Annotations: {sum(class_counts.values()):6d} | "
              f"Empty labels: {empty_labels}")

    print()
    if all_ok:
        print("  ✅ Dataset looks good! Ready to train.")
    else:
        print("  ❌ Issues found. Re-run merge_datasets.py")

    print("=" * 55)
    return all_ok


# -------------------------------------------------------
# RESUME YOLOv11n (50 more epochs → 51 to 100)
# -------------------------------------------------------

def resume_model():
    print("\n" + "=" * 55)
    print("  Resuming YOLOv11n — Epoch 51 → 100")
    print("=" * 55)

    try:
        from ultralytics import YOLO
    except ImportError:
        print("\n❌ ultralytics not installed!")
        return

    ckpt = Path(RESUME_CKPT)
    if not ckpt.exists():
        print(f"\n❌ Checkpoint nahi mila: {RESUME_CKPT}")
        print("\n   Fix: RESUME_CKPT path check karo is file mein (line 20)")
        print("   last.pt yahan hoga:")
        print("   C:\\Users\\HP\\runs\\detect\\runs\\pothole_v1\\yolo11n_pothole\\weights\\last.pt")
        return

    print(f"\n  Checkpoint: {RESUME_CKPT}")
    print("  last.pt mein resume state nahi tha — explicitly settings pass kar rahe hain\n")

    yaml_path = Path(DATA_YAML)
    if not yaml_path.exists():
        print(f"\n❌ data.yaml nahi mila: {DATA_YAML}")
        return

    model = YOLO(str(ckpt))

    model.train(
        data=str(yaml_path.resolve()),
        epochs=100,        # 100 tak chalega (jo baki hain woh)
        imgsz=416,
        batch=8,
        device=0,
        amp=False,
        workers=2,
        cache=False,
        patience=15,
        save=True,
        project=OUTPUT_DIR,
        name="yolo11n_pothole",
        exist_ok=True,
        hsv_h=0.015,
        hsv_s=0.7,
        hsv_v=0.4,
        flipud=0.1,
        fliplr=0.5,
        mosaic=1.0,
        mixup=0.1,
    )

    print("\n" + "=" * 55)
    print("  ✅ RESUME COMPLETE!")
    print("=" * 55)
    _print_results(OUTPUT_DIR, "yolo11n_pothole")


# -------------------------------------------------------
# FRESH TRAIN YOLOv11n
# -------------------------------------------------------

def train_model():
    print("\n" + "=" * 55)
    print("  Training YOLOv11n — GTX 1650 (4GB VRAM) Config")
    print("=" * 55)

    try:
        from ultralytics import YOLO
    except ImportError:
        print("\n❌ ultralytics not installed!")
        return

    yaml_path = Path(DATA_YAML)
    if not yaml_path.exists():
        print(f"\n❌ data.yaml not found: {DATA_YAML}")
        return

    print("\n  Loading YOLOv11n pretrained weights...")
    model = YOLO(WEIGHTS_N)

    print("\n  Starting training:")
    print("   imgsz=416  batch=8  amp=False  epochs=50\n")

    model.train(
        data=str(yaml_path.resolve()),
        epochs=50,
        imgsz=416,
        batch=8,
        device=0,
        amp=False,
        workers=2,
        cache=False,
        patience=15,
        save=True,
        project=OUTPUT_DIR,
        name="yolo11n_pothole",
        exist_ok=True,
        hsv_h=0.015,
        hsv_s=0.7,
        hsv_v=0.4,
        flipud=0.1,
        fliplr=0.5,
        mosaic=1.0,
        mixup=0.1,
    )

    print("\n" + "=" * 55)
    print("  ✅ TRAINING COMPLETE!")
    print("=" * 55)
    _print_results(OUTPUT_DIR, "yolo11n_pothole")


# -------------------------------------------------------
# TRAIN YOLOv11s
# -------------------------------------------------------

def train_model_s():
    print("\n" + "=" * 55)
    print("  Training YOLOv11s — GTX 1650 (4GB VRAM) Config")
    print("=" * 55)

    try:
        from ultralytics import YOLO
    except ImportError:
        print("\n❌ ultralytics not installed!")
        return

    yaml_path = Path(DATA_YAML)
    if not yaml_path.exists():
        print(f"\n❌ data.yaml not found: {DATA_YAML}")
        return

    if not Path(WEIGHTS_S).exists():
        print("\n  yolo11s.pt nahi mili!")
        print("  Download karo:")
        print("  https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11s.pt")
        print("  Aur project folder mein rakh do\n")

    print("\n  Loading YOLOv11s pretrained weights...")
    model = YOLO(WEIGHTS_S)

    print("\n  Starting YOLOv11s training:")
    print("   imgsz=512  batch=8  amp=False  epochs=80\n")

    model.train(
        data=str(yaml_path.resolve()),
        epochs=80,
        imgsz=512,        # v11n se bada — better accuracy
        batch=8,
        device=0,
        amp=False,
        workers=2,
        cache=False,
        patience=15,
        save=True,
        project=OUTPUT_DIR,
        name="yolo11s_pothole",  # alag folder — v11n se mix nahi hoga
        exist_ok=True,
        hsv_h=0.015,
        hsv_s=0.7,
        hsv_v=0.4,
        flipud=0.1,
        fliplr=0.5,
        mosaic=1.0,
        mixup=0.1,
        optimizer='AdamW',
lr0=0.001,
    )

    print("\n" + "=" * 55)
    print("  ✅ YOLOv11s TRAINING COMPLETE!")
    print("=" * 55)
    _print_results(OUTPUT_DIR, "yolo11s_pothole")


# -------------------------------------------------------
# HELPER
# -------------------------------------------------------

def _print_results(project, name):
    base = Path(project) / name
    best = base / "weights" / "best.pt"
    results_png = base / "results.png"

    print(f"\n  Best weights : {best}")
    print(f"  Results      : {results_png}")
    print(f"\n  Inference:")
    print(f"  yolo detect predict model={best} source=your_image.jpg conf=0.25")
    print("=" * 55)


# -------------------------------------------------------
# ENTRY POINT
# -------------------------------------------------------

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--verify", action="store_true", help="Verify merged dataset")
    parser.add_argument("--train",  action="store_true", help="Fresh train YOLOv11n")
    parser.add_argument("--resume", action="store_true", help="Resume YOLOv11n 50 more epochs")
    parser.add_argument("--trains", action="store_true", help="Train YOLOv11s")
    args = parser.parse_args()

    if not any([args.verify, args.train, args.resume, args.trains]):
        print("Usage:")
        print("  python train.py --verify   → dataset verify karo")
        print("  python train.py --train    → YOLOv11n fresh train")
        print("  python train.py --resume   → YOLOv11n resume (50 more epochs)")
        print("  python train.py --trains   → YOLOv11s train karo")
    else:
        if args.verify:
            verify_dataset()
        if args.train:
            train_model()
        if args.resume:
            resume_model()
        if args.trains:
            train_model_s()