predict.py

"""
predict.py — inference helper for yolov8s-rdd2022
Usage:
    python predict.py --source path/to/images --weights yolov8s_best.pt
"""
import argparse
from pathlib import Path
from ultralytics import YOLO

CLASS_NAMES = [
    "D00-LongitudinalCrack",
    "D10-TransverseCrack",
    "D20-AlligatorCrack",
    "D40-Pothole",
]

def run(source, weights, conf=0.25, iou=0.45, imgsz=640, device=0, save=True):
    model = YOLO(weights)
    results = model.predict(
        source=source,
        conf=conf,
        iou=iou,
        imgsz=imgsz,
        device=device,
        save=save,
        save_txt=True,
        project="runs/detect",
        name="rdd_predict",
    )

    pothole_count = 0
    total_boxes   = 0
    class_counts  = {name: 0 for name in CLASS_NAMES}

    for r in results:
        total_boxes += len(r.boxes)
        for cls in r.boxes.cls.tolist():
            class_counts[CLASS_NAMES[int(cls)]] += 1
            if int(cls) == 3:
                pothole_count += 1

    # ┌─────────────────────────────────────────────┐
    # │           DETECTION SUMMARY                 │
    # └─────────────────────────────────────────────┘
    print("\n" + "="*52)
    print("  📊  DETECTION SUMMARY")
    print("="*52)
    print(f"  Total detections     : {total_boxes}")
    for name, count in class_counts.items():
        bar = "█" * min(count, 30)
        print(f"  {name:<30} : {count:>4}  {bar}")
    print("-"*52)

    # ┌─────────────────────────────────────────────┐
    # │        🕳️  POTHOLE (D40) REPORT             │
    # └─────────────────────────────────────────────┘
    pct = pothole_count / max(total_boxes, 1) * 100
    print("\n  🕳️  POTHOLE (D40) REPORT")
    print("  " + "-"*48)
    print(f"  Pothole count        : {pothole_count}")
    print(f"  % of all detections  : {pct:.1f}%")
    if pct > 30:
        print("  ⚠️  HIGH pothole density — road surface critical")
    elif pct > 10:
        print("  ⚠️  MODERATE pothole presence")
    else:
        print("  ✅  LOW pothole presence")
    print("="*52 + "\n")

    return results


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--source",  required=True, help="image / folder / video path")
    ap.add_argument("--weights", default="yolov8s_best.pt")
    ap.add_argument("--conf",    type=float, default=0.25)
    ap.add_argument("--iou",     type=float, default=0.45)
    ap.add_argument("--device",  default="0")
    args = ap.parse_args()
    run(args.source, args.weights, args.conf, args.iou, device=args.device)