File size: 2,958 Bytes
7cfb55e | 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 | """
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) |