37 lines
1.0 KiB
Python
37 lines
1.0 KiB
Python
from ultralytics import YOLO
|
|
import cv2
|
|
import os
|
|
import numpy as np
|
|
|
|
# Load a pretrained YOLO11n model
|
|
model = YOLO("yolo11m-2_uniform.onnx")
|
|
model_pose = YOLO("yolo11s-pose.pt")
|
|
|
|
video_path = "PF-071124-2.mp4"
|
|
cap = cv2.VideoCapture(video_path)
|
|
|
|
ret = True
|
|
frame_count = 0 # Initialize a frame counter
|
|
save_dir = "/results/" # Directory to save the crops
|
|
os.makedirs(save_dir, exist_ok=True) # Create the directory if it doesn't exist
|
|
|
|
while ret:
|
|
ret, frame = cap.read()
|
|
frame_count += 1 # Increment the frame counter
|
|
#frame2_, frame_ = np.empty(2)
|
|
if ret:
|
|
results = model.predict(frame, conf=0.5)
|
|
frame_ = results[0].plot()
|
|
pose_results = model_pose.predict(frame_, show_boxes=False, show_conf=False, show_labels=False)
|
|
results[0].keypoints = pose_results[0].keypoints
|
|
|
|
frame_ = results[0].plot()
|
|
|
|
|
|
cv2.imshow("frame", frame_)
|
|
|
|
if cv2.waitKey(25) & 0xFF == ord('q'):
|
|
break
|
|
|
|
cap.release()
|
|
cv2.destroyAllWindows() |