Karan9630/PPE-DataSet

1from ultralytics import YOLO
2import torch
3import os
4
5def main():
6    device = 0 if torch.cuda.is_available() else "cpu"
7    print("Using device:", device)
8
9    model = YOLO("yolov8n.pt")
10
11    model.train(
12        data=r"E:\VISION-BASED-PPE-COMPLIANCE-MONITORING-SYSTEM\PPE Data\data.yml",
13        epochs=25,
14        imgsz=640,
15        batch=16,
16        device=device,
17        name="ppe_yolov8"
18    )
19
20if __name__ == "__main__":
21    main()
22
23