50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
from ultralytics import YOLO
|
|
import cv2
|
|
import os
|
|
import numpy as np
|
|
|
|
# Load the object detection model
|
|
det_model = YOLO("yolo11m-2_uniform.onnx") # General object detection model
|
|
|
|
# Load the pose estimation model
|
|
pose_model = YOLO("yolo11s-pose.pt") # Pose estimation model
|
|
|
|
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:
|
|
# Run object detection
|
|
det_results = det_model.predict(frame, conf=0.5)
|
|
|
|
# Filter detections for persons (class 0)
|
|
person_detections = [det for det in det_results if det.names[0] == 'crew']
|
|
|
|
#for i, det in enumerate(person_detections[0].boxes.xyxy):
|
|
# Extract bounding box coordinates
|
|
# x1, y1, x2, y2 = map(int, det[:4])
|
|
# crop = frame[y1:y2, x1:x2]
|
|
|
|
# Run pose estimation on detected persons
|
|
for i, person in person_detections:
|
|
#how to return the tensor to posemodel???
|
|
x1, y1, x2, y2 = map(int, person[i].boxes.xyxy[:4])
|
|
person_image = frame[y1:y2, x1:x2] # Crop the person from the image
|
|
|
|
pose_results = pose_model(person_image)
|
|
frame_ = pose_results[0].plot()
|
|
cv2.imshow("frame", frame_)
|
|
|
|
if cv2.waitKey(25) & 0xFF == ord('q'):
|
|
break
|
|
|
|
cap.release()
|
|
cv2.destroyAllWindows() |