diff --git a/main.py b/main.py new file mode 100644 index 0000000..518529c --- /dev/null +++ b/main.py @@ -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 \ No newline at end of file diff --git a/main2.py b/main2.py new file mode 100644 index 0000000..904f8fc --- /dev/null +++ b/main2.py @@ -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() \ No newline at end of file diff --git a/main3.py b/main3.py new file mode 100644 index 0000000..69f639b --- /dev/null +++ b/main3.py @@ -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() \ No newline at end of file diff --git a/yolo11s-pose.pt b/yolo11s-pose.pt new file mode 100644 index 0000000..997c0bb Binary files /dev/null and b/yolo11s-pose.pt differ