[trackingdata] fixes and support for selections when getting data of columns

This commit is contained in:
Jan Grewe 2025-02-21 16:19:24 +01:00
parent f09c78adb5
commit af5dbc7dfc

View File

@ -5,6 +5,7 @@ import pandas as pd
from PySide6.QtCore import QObject from PySide6.QtCore import QObject
class TrackingData(QObject): class TrackingData(QObject):
def __init__(self, parent=None): def __init__(self, parent=None):
super().__init__(parent) super().__init__(parent)
@ -58,9 +59,14 @@ class TrackingData(QObject):
self._start = start self._start = start
self._stop = stop self._stop = stop
self._selection_column = col self._selection_column = col
self._indices = np.where((self._data[col] >= self._start) & (self._data[col] < self._stop))[0] col_indices = np.where((self._data[col] >= self._start) & (self._data[col] < self._stop))[0]
self._indices = self["index"][col_indices]
def selectedData(self, col): if len(col_indices) < 1:
logging.warning("TrackingData: Selection range is empty!")
def selectedData(self, col:str):
if col not in self.columns:
logging.error("TrackingData:selectedData: Invalid column name! %s", col)
return self[col][self._indices] return self[col][self._indices]
def setUserSelection(self, ids): def setUserSelection(self, ids):
@ -148,11 +154,13 @@ class TrackingData(QObject):
and M is number of keypoints and M is number of keypoints
""" """
if selection: if selection:
return np.stack(self._data["keypoints"][self._start:self._stop, :, :]).astype(np.float32) if len(self._indices) < 1:
else: logging.info("TrackingData.coordinates returns empty array, not detections in range!")
return np.ndarray([])
return np.stack(self._data["keypoints"][self._indices]).astype(np.float32)
return np.stack(self._data["keypoints"]).astype(np.float32) return np.stack(self._data["keypoints"]).astype(np.float32)
def keypointScores(self): def keypointScores(self, selection=False):
""" """
Returns the keypoint scores as a NumPy array of type float32. Returns the keypoint scores as a NumPy array of type float32.
@ -162,9 +170,14 @@ class TrackingData(QObject):
A NumPy array of type float32 containing the keypoint scores of the shape (N, M) A NumPy array of type float32 containing the keypoint scores of the shape (N, M)
with N the number of detections and M the number of keypoints. with N the number of detections and M the number of keypoints.
""" """
if selection:
if len(self._indices) < 1:
logging.info("TrackingData.scores returns empty array, not detections in range!")
return np.ndarray([])
return np.stack(self._data["keypoint_score"][self._indices]).astype(np.float32)
return np.stack(self._data["keypoint_score"]).astype(np.float32) return np.stack(self._data["keypoint_score"]).astype(np.float32)
def centerOfGravity(self, threshold=0.8): def centerOfGravity(self, selection=False, threshold=0.8, nodes=[0,1,2]):
""" """
Calculate the center of gravity of keypoints weighted by their scores. Ignores keypoints that have a score Calculate the center of gravity of keypoints weighted by their scores. Ignores keypoints that have a score
less than threshold. less than threshold.
@ -172,16 +185,19 @@ class TrackingData(QObject):
Parameters: Parameters:
----------- -----------
threshold: float threshold: float
keypoints with a score less than threshold are ignored nodes with a score less than threshold are ignored
nodes: list
nodes/keypoints to consider for estimation. Defaults to [0,1,2]
Returns: Returns:
-------- --------
np.ndarray: np.ndarray:
A NumPy array of shape (N, 2) containing the center of gravity for each detection. A NumPy array of shape (N, 2) containing the center of gravity for each detection.
""" """
scores = self.keypointScores() scores = self.keypointScores(selection)
scores[scores < threshold] = 0.0 scores[scores < threshold] = 0.0
weighted_coords = self.coordinates() * scores[:, :, np.newaxis] scores[:, np.setdiff1d(np.arange(scores.shape[1]), nodes)] = 0.0
weighted_coords = self.coordinates(selection=selection) * scores[:, :, np.newaxis]
sum_scores = np.sum(scores, axis=1, keepdims=True) sum_scores = np.sum(scores, axis=1, keepdims=True)
center_of_gravity = np.sum(weighted_coords, axis=1) / sum_scores center_of_gravity = np.sum(weighted_coords, axis=1) / sum_scores
return center_of_gravity return center_of_gravity
@ -285,7 +301,7 @@ def main():
frames = data["frame"] frames = data["frame"]
tracks = data["track"] tracks = data["track"]
bendedness = data.bendedness() bendedness = data.bendedness()
positions = data.coordinates()[[160388, 160389]] # positions = data.coordinates()[[160388, 160389]]
embed() embed()
tracks = data["track"] tracks = data["track"]
@ -309,7 +325,6 @@ def main():
# return distances # return distances
# print("estimating neighorhood distances") # print("estimating neighorhood distances")
# neighbor_distances = compute_neighbor_distances(cogs) # neighbor_distances = compute_neighbor_distances(cogs)
embed()
if __name__ == "__main__": if __name__ == "__main__":
main() main()