-
Notifications
You must be signed in to change notification settings - Fork 7
/
predict2D.py
executable file
·125 lines (99 loc) · 4.42 KB
/
predict2D.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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
"""
JARVIS-MoCap (https://jarvis-mocap.github.io/jarvis-docs)
Copyright (c) 2022 Timo Hueser.
https://github.com/JARVIS-MoCap/JARVIS-HybridNet
Licensed under GNU Lesser General Public License v2.1
"""
import os
import csv
import itertools
import numpy as np
import torch
import cv2
from tqdm import tqdm
import time
from ruamel.yaml import YAML
from jarvis.prediction.jarvis2D import JarvisPredictor2D
from jarvis.config.project_manager import ProjectManager
def create_info_file(params):
with open(os.path.join(params.output_dir, 'info.yaml'), 'w') as file:
yaml=YAML()
yaml.dump({'recording_path': params.recording_path,
'frame_start': params.frame_start,
'number_frames': params.number_frames}, file)
def predict2D(params):
project = ProjectManager()
if not project.load(params.project_name):
print (f'{CLIColors.FAIL}Could not load project: {project_name}! '
f'Aborting....{CLIColors.ENDC}')
return
cfg = project.cfg
params.output_dir = os.path.join(project.parent_dir,
cfg.PROJECTS_ROOT_PATH, params.project_name,
'predictions', 'predictions2D',
f'Predictions_2D_{time.strftime("%Y%m%d-%H%M%S")}')
os.makedirs(params.output_dir, exist_ok = True)
create_info_file(params)
jarvisPredictor = JarvisPredictor2D(cfg, params.weights_center_detect,
params.weights_keypoint_detect, params.trt_mode)
recording_paths = []
multiple_videos = False
if os.path.isfile(params.recording_path):
recording_paths.append(params.recording_path)
elif os.path.exists(params.recording_path):
print ("DIRECTORY!")
multiple_videos = True
recording_paths = [os.path.join(params.recording_path, file) for file in os.listdir(params.recording_path)]
print (recording_paths)
for recording_path in recording_paths:
cap = cv2.VideoCapture(recording_path)
cap.set(1,params.frame_start)
img_size = [int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))]
csv_filename = 'data2D.csv'
if multiple_videos:
csv_filename = f'{recording_path.split(os.sep)[-1].split(".")[0]}_{csv_filename}'
csvfile = open(os.path.join(params.output_dir, csv_filename), 'w',
newline='')
writer = csv.writer(csvfile, delimiter=',',
quotechar='"', quoting=csv.QUOTE_MINIMAL)
#if keypoint names are defined, add header to csvs
if (len(cfg.KEYPOINT_NAMES) == cfg.KEYPOINTDETECT.NUM_JOINTS):
create_header(writer, cfg)
assert params.frame_start < cap.get(cv2.CAP_PROP_FRAME_COUNT), \
"frame_start bigger than total framecount!"
if (params.number_frames == -1):
params.number_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) \
- params.frame_start
else:
assert params.frame_start+params.number_frames \
<= cap.get(cv2.CAP_PROP_FRAME_COUNT), \
"make sure your selected segment is not longer that the " \
"total video!"
for frame_num in tqdm(range(params.number_frames)):
ret, img_orig = cap.read()
img = torch.from_numpy(
img_orig).cuda().float().permute(2,0,1)[[2, 1, 0]]/255.
points2D, confidences = jarvisPredictor(img.unsqueeze(0))
if points2D != None:
points2D = points2D.cpu().numpy()
confidences = confidences.cpu().numpy()
row = []
for i,point in enumerate(points2D):
row = row + point.tolist() + [confidences[i]]
writer.writerow(row)
else:
row = []
for i in range(cfg.KEYPOINTDETECT.NUM_JOINTS*3):
row = row + ['NaN']
writer.writerow(row)
if params.progress_bar != None:
params.progress_bar.progress(float(frame_num+1)
/ float(params.number_frames))
cap.release()
def create_header(writer, cfg):
joints = list(itertools.chain.from_iterable(itertools.repeat(x, 3)
for x in cfg.KEYPOINT_NAMES))
coords = ['x','y','confidence']*len(cfg.KEYPOINT_NAMES)
writer.writerow(joints)
writer.writerow(coords)