diff --git a/fixtracks/utils/merger.py b/fixtracks/utils/merger.py index 2dfc02c..a67bd58 100644 --- a/fixtracks/utils/merger.py +++ b/fixtracks/utils/merger.py @@ -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