40 lines
1.1 KiB
Python
40 lines
1.1 KiB
Python
import glob
|
|
import pickle
|
|
from tqdm import tqdm
|
|
|
|
import torch
|
|
from ultralytics import YOLO
|
|
import cv2
|
|
import time
|
|
|
|
import utils
|
|
from calculate import normalize_pose, compare_poses_boolean
|
|
from draw import draw_new
|
|
from utils import find_closest
|
|
from video_methods import initialize_method
|
|
|
|
model = YOLO("yolo11x-pose.pt")
|
|
model.to(torch.device('cuda:0'))
|
|
|
|
startTime = 0
|
|
def main():
|
|
moves = []
|
|
|
|
for i in tqdm(sorted(glob.glob("video/camA/*.jpg"), key=lambda f: int(__import__("re").search(r"\d+", f).group()))):
|
|
data = i.replace(f'video/camA\\', "")
|
|
frame = cv2.imread(f"video/camA/{data}")
|
|
results = model(frame, verbose=False)
|
|
|
|
if len(results) != 0:
|
|
result = results[0]
|
|
kpts = result.keypoints.data[0] if len(result.keypoints.data) else None
|
|
|
|
if kpts is None:
|
|
continue
|
|
|
|
moves.append((int(data.replace(".jpg", "")), normalize_pose(result.keypoints.xy.cpu().numpy()[0]), result.keypoints.xy.cpu()[0]))
|
|
|
|
with open('moves_makarena_best.pkl', 'wb') as f: # 'wb' = write binary
|
|
pickle.dump(moves, f)
|
|
|
|
main() |