initial commit
This commit is contained in:
parent
d8521ef8d3
commit
6fd7cbf07e
20
main.py
Normal file
20
main.py
Normal file
@ -0,0 +1,20 @@
|
||||
from ultralytics import YOLO
|
||||
|
||||
# Load a pretrained YOLO11n model
|
||||
model = YOLO("yolo11m-2_uniform.onnx")
|
||||
model_pose = YOLO("yolo11s-pose.pt")
|
||||
|
||||
# Define path to video file
|
||||
source = "PF-071124-2.mp4"
|
||||
|
||||
# Run inference on the source
|
||||
results = model(source, stream=True) # generator of Results objects
|
||||
|
||||
for result in results:
|
||||
boxes = result.boxes # Boxes object for bounding box outputs
|
||||
masks = result.masks # Masks object for segmentation masks outputs
|
||||
keypoints = result.keypoints # Keypoints object for pose outputs
|
||||
probs = result.probs # Probs object for classification outputs
|
||||
obb = result.obb # Oriented boxes object for OBB outputs
|
||||
result.show() # display to screen
|
||||
result.save(filename="result.jpg") # save to disk
|
||||
37
main2.py
Normal file
37
main2.py
Normal file
@ -0,0 +1,37 @@
|
||||
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()
|
||||
50
main3.py
Normal file
50
main3.py
Normal file
@ -0,0 +1,50 @@
|
||||
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()
|
||||
BIN
yolo11s-pose.pt
Normal file
BIN
yolo11s-pose.pt
Normal file
Binary file not shown.
Loading…
Reference in New Issue
Block a user