-
Notifications
You must be signed in to change notification settings - Fork 0
/
videoPredictor.py
95 lines (77 loc) · 3 KB
/
videoPredictor.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
import cv2
import pandas as pd
from videoProcessor import extract_keypoint, keyPointsToVideo
from utils import dataCols
import os
def runPrediction(model,
input_video_path,
output_video_path,
is_keypoint_file=False,
is_keypoint_video=False,
keypoint_bones=False):
try:
# Load video
video_path = os.path.join('local_uploads',input_video_path)
cap = cv2.VideoCapture(video_path)
# Define output video writer
output_path = os.path.join('local_processed',output_video_path)
fps = cap.get(cv2.CAP_PROP_FPS)
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
print('Frame Width, Frame Height --- ',frame_width,frame_height)
out = cv2.VideoWriter(
output_path,
cv2.VideoWriter_fourcc(*'mp4v'),
fps,
(frame_width, frame_height)
)
frame_no = 0
dataset_csv = []
# Process video frame by frame
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
# Perform inference
results = model(frame)
annotated_frame = results[0].plot()
# Write the frame to the output video
out.write(annotated_frame)
results_keypoint = results[0].keypoints.xy.cpu().numpy()
#results_keypoint = results[0].keypoints.xyn.cpu().numpy()
#Process Key Points from prediction results
for result_keypoint in results_keypoint:
if len(result_keypoint) == 17:
keypoint_list = extract_keypoint(result_keypoint)
keypoint_list.append(frame_no)
keypoint_list.append(frame_width)
keypoint_list.append(frame_height)
dataset_csv.append(keypoint_list)
frame_no += 1
if cv2.waitKey(1) & 0xFF == ord('q'):
break
# Release video capture and writer
cap.release()
out.release()
cv2.destroyAllWindows()
#Store the Keypoints as CSV
if is_keypoint_file:
df = pd.DataFrame(
dataset_csv,
columns=dataCols
)
if '\\' in output_path:
datafilename = output_path.split('\\')[-1].split('.')[0] + '.csv'
else:
datafilename = output_path.split('/')[-1].split('.')[0] + '.csv'
df.to_csv( os.path.join('local_data', datafilename), index=False)
print('Output Kepoints File Shape : ',df.shape)
#Store the Keypoints as Video
if is_keypoint_video:
keyPointsToVideo(
df=df,
plotBones=keypoint_bones,
output_file_name=output_video_path.split('.')[0]
)
except Exception as e:
raise(e)