[merger] speed up numpy conversion

This commit is contained in:
Jan Grewe 2025-01-28 09:32:58 +01:00
parent a98770610b
commit 456bd16a59

View File

@ -56,28 +56,23 @@ class Merger(QRunnable):
numpy.ndarray
2D array, Coordinates of the bounding box for each detection. Shape: (num_detections, 4) x1, y1, x2, y2
"""
x = lambda kp: np.array([float(kpp) for kpp in kp[1:-1].split(",")])
logging.info("Converting to numpy ...")
# from IPython import embed
# embed()
dimensions = 2
key_columns = [c for c in df.columns if "key_" in c]
box_columns = [c for c in df.columns if "box_" in c]
num_frames = len(df)
num_keypoints = len(key_columns)
dimensions = 2
keypoints = np.empty((num_frames, num_keypoints, dimensions))
visibility = np.empty((num_frames, num_keypoints))
boxcoordinates = np.empty((num_frames, 4))
for i, row in df.iterrows():
for j, k in enumerate(key_columns):
key_data = row[k]
l = list(map(float, list(key_data[1:-1].split(","))))
keypoints[i, j, :] = l
for j, b in enumerate(box_columns):
boxcoordinates[i, j] = row[b]
if isinstance(row["visible"], str):
vis = list(map(float, row["visible"][1:-1].split()))
visibility[i, :] = vis
else:
visibility[i, :] = row["visible"]
keypoints = np.empty((len(df), len(key_columns), dimensions))
visibility = np.empty((len(df), len(key_columns)))
boxcoordinates = np.stack([df[b].values for b in box_columns]).astype(np.float32).T
temp_v = df.visible.values
temp_k = np.stack([df[k].values for k in key_columns]).T
for i, k in enumerate(temp_k):
temp = np.array([x(kp) for kp in k])
keypoints[i, :, :] = temp
visibility[i,:] = [float(val) for val in temp_v[i][1:-1].split()]
logging.info("Converting to numpy done!")
return keypoints, visibility, boxcoordinates