45 lines
1.4 KiB
Python
45 lines
1.4 KiB
Python
import cv2
|
|
import numpy as np
|
|
from ultralytics import YOLO
|
|
import cvzone
|
|
|
|
device = 'cuda'
|
|
|
|
# Load the YOLO11 model
|
|
model = YOLO("yolo11m-pecan.pt")
|
|
|
|
# Open the video file (use video file or webcam, here using webcam)
|
|
cap = cv2.VideoCapture(0)
|
|
cy1=550
|
|
offset=30
|
|
idDict={}
|
|
pecanCount = 0
|
|
|
|
while True:
|
|
ret,frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
# Run YOLO11 tracking on the frame, persisting tracks between frames
|
|
results = model.track(frame, persist=True,classes=0,device = device)
|
|
|
|
# Check if there are any boxes in the results
|
|
if results[0].boxes is not None and results[0].boxes.id is not None:
|
|
# Get the boxes (x, y, w, h), class IDs, track IDs, and confidences
|
|
boxes = results[0].boxes.xyxy.int().cpu().tolist() # Bounding boxes
|
|
class_ids = results[0].boxes.cls.int().cpu().tolist() # Class IDs
|
|
track_ids = results[0].boxes.id.int().cpu().tolist() # Track IDs
|
|
confidences = results[0].boxes.conf.cpu().tolist() # Confidence score
|
|
|
|
for box, class_id, track_id, conf in zip(boxes, class_ids, track_ids, confidences):
|
|
x1, y1, x2, y2 = box
|
|
cy = int(y1+y2)//2
|
|
if cy<(cy1+offset) and cy>(cy1-offset) and track_id not in idDict.keys():
|
|
pecanCount += 1
|
|
idDict[track_id] = pecanCount
|
|
|
|
# Release the video capture object and close the display window
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|
|
|
|
print(pecanCount) |