| from pathlib import Path | |
| from ultralytics import YOLO | |
| DATASET_YAML = Path(__file__).parent / "dataset" / "dataset.yaml" | |
| RUNS_DIR = Path(__file__).parent / "runs" | |
| def train(): | |
| model = YOLO("runs/pothole/weights/best.pt") | |
| results = model.train( | |
| data=str(DATASET_YAML), | |
| epochs=100, | |
| imgsz=1280, | |
| batch=64, | |
| project=str(RUNS_DIR), | |
| name="pothole", | |
| exist_ok=True, | |
| # LR scheduling — cosine annealing with warmup | |
| lr0=0.001, | |
| lrf=0.01, | |
| warmup_epochs=3, | |
| warmup_momentum=0.8, | |
| cos_lr=True, | |
| # Regularization | |
| weight_decay=0.0005, | |
| dropout=0.0, | |
| patience=50, | |
| # Augmentation — tuned for road/pothole data | |
| mosaic=1.0, | |
| copy_paste=0.3, # paste potholes onto new backgrounds | |
| degrees=5.0, # slight rotation (roads tilt a bit) | |
| scale=0.5, # scale jitter | |
| fliplr=0.5, # horizontal flip | |
| flipud=0.0, # roads are always right-side up | |
| hsv_h=0.015, # hue jitter (lighting conditions) | |
| hsv_s=0.7, # saturation jitter (wet/dry roads) | |
| hsv_v=0.4, # brightness jitter (day/night) | |
| translate=0.1, | |
| perspective=0.0, # road camera is always flat | |
| mixup=0.1, # blend images for harder examples | |
| ) | |
| print(f"Training complete. Results saved to: {results.save_dir}") | |
| return results | |
| def main(): | |
| train() | |
| if __name__ == "__main__": | |
| main() | |