Snake Detection

1import cv2
2from ultralytics import YOLO
3import argparse
4
5def load_model(weights_path):
6    """Load the YOLOv8 model with custom weights."""
7    model = YOLO(weights_path)
8    return model
9
10def detect_snake(model, source, conf=0.5):
11    """Run snake detection on the given source."""
12    results = model.predict(
13        source=source,
14        conf=conf,
15        save=True,
16        show=True
17    )
18    return results
19
20def process_results(results):
21    """Process and display detection results."""
22    for result in results:
23        boxes = result.boxes
24        for box in boxes:
25            cls = int(box.cls[0])
26            conf = float(box.conf[0])
27            label = result.names[cls]
28            print(f"Detected: {label} ({conf:.2f})")
29    return results
30
31if __name__ == "__main__":
32    parser = argparse.ArgumentParser()
33    parser.add_argument("--weights", default="best.pt")
34    parser.add_argument("--source", required=True)
35    parser.add_argument("--conf", type=float, default=0.5)
36    args = parser.parse_args()
37
38    model = load_model(args.weights)
39    results = detect_snake(model, args.source, args.conf)
40    process_results(results)