[tracks] react on move request and move window

something is still off with the deletion...

FixTracks.py

@@ -0,0 +1,80 @@
+"""
+pyside6-rcc resources.qrc -o resources.py
+
+"""
+import sys
+import logging
+import argparse
+import platform
+
+from PySide6.QtWidgets import QApplication
+from PySide6.QtCore import QSettings
+from PySide6.QtGui import QIcon, QPalette
+
+from fixtracks import info, mainwindow
+
+logging.basicConfig(level=logging.INFO, force=True)
+
+
+def is_dark_mode(app: QApplication) -> bool:
+    palette = app.palette()
+    # Check the brightness of the window text and base colors
+    text_color = palette.color(QPalette.ColorRole.WindowText)
+    base_color = palette.color(QPalette.ColorRole.Base)
+
+    # Calculate brightness (0 for dark, 255 for bright)
+    def brightness(color):
+        return (color.red() * 299 + color.green() * 587 + color.blue() * 114) // 1000
+
+    return brightness(base_color) < brightness(text_color)
+
+def set_logging(loglevel):
+    logging.basicConfig(level=loglevel, force=True)
+
+def main(args):
+    set_logging(args.loglevel)
+    if platform.system() == "Windows":
+        # from PySide6.QtWinExtras import QtWin
+        myappid = f"{info.organization_name}.{info.application_version}"
+        # QtWin.setCurrentProcessExplicitAppUserModelID(myappid)
+    settings = QSettings()
+    width = int(settings.value("app/width", 1024))
+    height = int(settings.value("app/height", 768))
+    x = int(settings.value("app/pos_x", 100))
+    y = int(settings.value("app/pos_y", 100))
+
+    app = QApplication(sys.argv)
+    app.setApplicationName(info.application_name)
+    app.setApplicationVersion(str(info.application_version))
+    app.setOrganizationDomain(info.organization_name)
+
+    # if platform.system() == 'Linux':
+    # icn = QIcon(":/icons/app_icon")
+    # app.setWindowIcon(icn)
+    # Create a Qt widget, which will be our window.
+    window = mainwindow.MainWindow(is_dark_mode(app))
+    window.setGeometry(100, 100, 1024, 768)
+    window.setWindowTitle("FixTracks")
+    window.setMinimumWidth(1024)
+    window.setMinimumHeight(768)
+    window.resize(width, height)
+    window.move(x, y)
+    window.show()
+
+    # Start the event loop.
+    app.exec()
+    pos = window.pos()
+    settings.setValue("app/width", window.width())
+    settings.setValue("app/height", window.height())
+    settings.setValue("app/pos_x", pos.x())
+    settings.setValue("app/pos_y", pos.y())
+
+
+if __name__ == "__main__":
+    levels = {"critical": logging.CRITICAL, "error": logging.ERROR, "warning":logging.WARNING, "info":logging.INFO, "debug":logging.DEBUG}
+    parser = argparse.ArgumentParser(description="FixTracks. Tools for fixing animal tracking")
+    parser.add_argument("-ll", "--loglevel", type=str, default="INFO", help=f"The log level that should be used. Valid levels are {[str(k) for k in levels.keys()]}")
+    args = parser.parse_args()
+    args.loglevel = levels[args.loglevel.lower() if args.loglevel.lower() in levels else "info"]
+
+    main(args)

fixtracks/mainwindow.py

@@ -1,11 +1,9 @@
 from PySide6.QtCore import QSize, Qt
-from PySide6.QtWidgets import QMainWindow, QWidget, QToolBar, QMenu, QMenuBar, QSizePolicy, QFileDialog
-from PySide6.QtWidgets import QDialog, QVBoxLayout
-from PySide6.QtGui import QKeySequence, QAction, QIcon, QPalette
+from PySide6.QtWidgets import QMainWindow, QWidget, QToolBar, QSizePolicy, QFileDialog
+from PySide6.QtGui import QKeySequence, QAction, QIcon
 
 from fixtracks.widgets.centralwidget import CentralWidget
 from fixtracks.dialogs.previewdialog import PreviewDialog
-from fixtracks.utils.reader import ImageReader, DataFrameReader
 from fixtracks.dialogs.about import AboutDialog
 from fixtracks.dialogs.help import HelpDialog
 import fixtracks.resources
@@ -146,8 +144,8 @@ class MainWindow(QMainWindow):
         about.show()
 
     def on_help(self, s):
-        help = HelpDialog(self)
-        help.show()
+        help_dlg = HelpDialog(self)
+        help_dlg.show()
 
     # @Slot(None)
     def exit_request(self):

fixtracks/utils/enums.py

@@ -0,0 +1,31 @@
+from enum import Enum
+
+from PySide6.QtGui import QColor
+
+class DetectionData(Enum):
+    ID = 0
+    FRAME = 1
+    COORDINATES = 2
+    TRACK_ID = 3
+    USERLABELED = 4
+    SCORE = 5
+
+class Tracks(Enum):
+    TRACKONE = 1
+    TRACKTWO = 2
+    UNASSIGNED = -1
+
+    def toColor(self):
+        track_colors = {
+            Tracks.TRACKONE: QColor.fromString("orange"),
+            Tracks.TRACKTWO: QColor.fromString("green"),
+            Tracks.UNASSIGNED: QColor.fromString("red")
+        }
+        return track_colors.get(self, QColor(128, 128, 128))  # Default to black if not found
+
+    @staticmethod
+    def fromValue(value):
+        for track in Tracks:
+            if track.value == value:
+                return track
+        return Tracks.UNASSIGNED

fixtracks/utils/reader.py

@@ -20,9 +20,6 @@ class ImageReader(QRunnable):
 
     @Slot()
     def run(self):
-        '''
-        Your code goes in this function
-        '''
         logging.debug("ImageReader: trying to open file %s", self._filename)
         cap = cv.VideoCapture(self._filename)
         framecount = int(cap.get(cv.CAP_PROP_FRAME_COUNT))

fixtracks/utils/signals.py

@@ -22,6 +22,8 @@ class DetectionSceneSignals(QObject):
 
 class DetectionTimelineSignals(QObject):
     windowMoved = Signal()
+    manualMove = Signal()
+    moveRequest = Signal(float)
 
 class DetectionSignals(QObject):
     hover = Signal((int, QPointF))

fixtracks/utils/styles.py

@@ -0,0 +1,18 @@
+def pushBtnStyle(color):
+    style = f"""
+    QPushButton {{
+        background-color: {color};
+        border-style: outset;
+        border-width: 1px;
+        border-radius: 5px;
+        border-color: white;
+        font: bold 10px;
+        min-width: 10em;
+        padding: 3px;
+    }}
+    QPushButton:pressed {{
+        background-color: rgb(220, 220, 220);
+        border-style: inset;
+    }}
+    """
+    return style

fixtracks/utils/trackingdata.py

@@ -5,21 +5,18 @@ import pandas as pd
 
 from PySide6.QtCore import QObject
 
-class TrackingData(QObject):
-    def __init__(self, parent=None):
-        super().__init__(parent)
-        self._data = None
-        self._columns = []
-        self._start = 0
-        self._stop = 0
-        self._indices = None
-        self._selection_column = None
-        self._user_selections = None
 
-    def setData(self, datadict):
+class TrackingData(QObject):
+    def __init__(self, datadict, parent=None):
+        super().__init__(parent)
         assert isinstance(datadict, dict)
         self._data = datadict
+        if "userlabeled" not in self._data.keys():
+            self._data["userlabeled"] = np.zeros_like(self["frame"], dtype=bool)
         self._columns = [k for k in self._data.keys()]
+        self._indices = np.arange(len(self["index"]), dtype=int)
+        self._selection_indices = np.asarray([])
+        self._selected_ids = None
 
     @property
     def data(self):
@@ -40,63 +37,125 @@ class TrackingData(QObject):
     def numDetections(self):
         return self._data["track"].shape[0]
 
-    @property
-    def selectionRange(self):
-        return self._start, self._stop
-
-    @property
-    def selectionRangeColumn(self):
-        return self._selection_column
+    def _find(self, ids):
+        if len(ids) < 1:
+            return np.array([])
+        ids = np.sort(ids)
+        indexes = np.ones_like(ids, dtype=int) * -1
+        j = 0
+        for idx in self._indices:
+            if self["index"][idx] == ids[j]:
+                indexes[j] = idx
+                j += 1
+            if j == len(indexes):
+                break
+        indexes = indexes[indexes >= 0]
+        return indexes
 
     @property
     def selectionIndices(self):
-        return self._indices
+        return self._selection_indices
+
+    @property
+    def selectionIDs(self):
+        return self._selected_ids
 
     def setSelectionRange(self, col, start, stop):
-        self._start = start
-        self._stop = stop
-        self._selection_column = col
-        self._indices = np.where((self._data[col] >= self._start) & (self._data[col] < self._stop))[0]
+        logging.info("Trackingdata: set selection range based on column %s to %.2f - %.2f", col, start, stop)
+        col_indices = np.where((self[col] >= start) & (self[col] < stop))[0]
+        self._selection_indices = self._indices[col_indices]
+        if len(col_indices) < 1:
+            logging.warning("TrackingData: Selection range is empty!")
 
-    def selectedData(self, col):
-        return self._data[col][self._indices]
+    def selectedData(self, col:str):
+        if col not in self.columns:
+            logging.error("TrackingData:selectedData: Invalid column name! %s", col)
+        return self[col][self._selection_indices]
 
-    def setUserSelection(self, ids):
+    def setSelection(self, ids):
         """
-        Set the user selections. That is, e.g. when the user selected a number of ids.
+        Set the selection based on the detection IDs.
+
         Parameters
         ----------
         ids : array-like
             An array-like object containing the IDs to be set as user selections. 
-            The IDs will be converted to integers.
         """
-        self._user_selections = ids.astype(int)
+        logging.debug("TrackingData.setSelection: %i number of ids", len(ids))
+        self._selection_indices = self._find(ids)
+        self._selected_ids = ids
+        # print(self._selected_ids, self._selection_indices)
 
-    def assignUserSelection(self, track_id:int)-> None:
+    def setTrack(self, track_id:int, setUserLabeled:bool=True)-> None:
         """Assign a new track_id to the user-selected detections
 
         Parameters
         ----------
         track_id : int
             The new track id for the user-selected detections
+        setUserLabeled : bool
+            Should the "userlabeled" state of the detections be set to True? Otherwise they will be left untouched.
         """
-        self._data["track"][self._user_selections] = track_id
+        logging.info("TrackingData: set track id %i for selection, set user-labeled status %s", track_id, str(setUserLabeled))
+        # print(self._selected_ids, self._selection_indices)
+        # print("before: ", self["track"][self._selection_indices], self["frame"][self._selection_indices])
+        self["track"][self._selection_indices] = track_id
+        if setUserLabeled:
+            self.setUserLabeledStatus(True, True)
+        # print("after: ", self["track"][self._selection_indices], self["frame"][self._selection_indices])
 
-    def assignTracks(self, tracks):
-        """assignTracks _summary_
+    def setUserLabeledStatus(self, new_status: bool, selection=True):
+        """Sets the status of the "userlabeled" column to a given value (True|False). This can done for ALL data in one go, or only for the UserSelection.
 
         Parameters
         ----------
-        tracks : _type_
-            _description_
+        new_status : bool
+            The new status, TRUE, if the detections are confirmed by the user (human observer) and can be treated as correct
+        selection : bool, optional
+            Whether the new status should be set for the selection only (True, default) ore not (False)
+        """
+        logging.debug("TrackingData: (Re-)setting assignment status of %s to %s", 
+                      "user selected data" if selection else " ALL", str(new_status))
+        if selection:
+            self["userlabeled"][self._selection_indices] = new_status
+        else:
+            self["userlabeled"][:] = new_status
+
+    def revertUserLabeledStatus(self):
+        logging.debug("TrackingData:Un-setting assignment status of all data!")
+        self["userlabeled"][:] = False
+
+    def revertTrackAssignments(self):
+        logging.debug("TrackingData: Reverting all track assignments!")
+        self["track"][:] = -1
+
+    def deleteDetections(self, ids=None):
+        if ids is not None:
+            logging.debug("TrackingData.deleteDetections of %i detections", len(ids))
+            del_indices = self._find(ids)
+        else:
+            logging.debug("TrackingData.deleteDetections of all selected detections (%i)", len(self._selected_ids))
+            del_indices = self._selected_ids
+        for c in self._columns:
+            self._data[c] = np.delete(self._data[c], del_indices, axis=0)
+        self._indices = self._indices[:-len(del_indices)]
+        self._selected_ids = np.setdiff1d(self._selected_ids, del_indices)
+
+    def assignTracks(self, tracks:np.ndarray):
+        """assigns the given tracks to the user-selected detections. If the sizes of 
+        provided tracks and the user selection do not match and error is logged and the tracks are not set.
+
+        Parameters
+        ----------
+        tracks : np.ndarray
+            The track information.
 
         Returns
         -------
-        _type_
-            _description_
+        None
         """
         if len(tracks) != self.numDetections:
-            logging.error("DataController: Size of passed tracks does not match data!")
+            logging.error("Trackingdata: Size of passed tracks does not match data!")
             return
         self._data["track"] = tracks
 
@@ -113,7 +172,7 @@ class TrackingData(QObject):
             return 0
         return self._data["keypoints"][0].shape[0]
 
-    def coordinates(self):
+    def coordinates(self, selection=False):
         """
         Returns the coordinates of all keypoints as a NumPy array.
 
@@ -121,9 +180,14 @@ class TrackingData(QObject):
             np.ndarray: A NumPy array of shape (N, M, 2) where N is the number of detections, 
                         and M is number of keypoints
         """
-        return np.stack(self._data["keypoints"]).astype(np.float32)
+        if selection:
+            if len(self._selection_indices) < 1:
+                logging.info("TrackingData.coordinates returns empty array, not detections in range!")
+                return np.ndarray([])
+            return np.stack(self["keypoints"][self._selection_indices]).astype(np.float32)
+        return np.stack(self["keypoints"]).astype(np.float32)
 
-    def keypointScores(self):
+    def keypointScores(self, selection=False):
         """
         Returns the keypoint scores as a NumPy array of type float32.
 
@@ -133,9 +197,14 @@ class TrackingData(QObject):
             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.
         """
-        return np.stack(self._data["keypoint_score"]).astype(np.float32)
+        if selection:
+            if len(self._selection_indices) < 1:
+                logging.info("TrackingData.scores returns empty array, not detections in range!")
+                return None
+            return np.stack(self["keypoint_score"][self._selection_indices]).astype(np.float32)
+        return np.stack(self["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
         less than threshold.
@@ -143,19 +212,27 @@ class TrackingData(QObject):
         Parameters:
         -----------
         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:
         --------
         np.ndarray: 
         A NumPy array of shape (N, 2) containing the center of gravity for each detection.
         """
-        scores = self.keypointScores()
+        scores = self.keypointScores(selection)
+        if scores is None:
+            return None
         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) * scores[:, :, np.newaxis]
         sum_scores = np.sum(scores, axis=1, keepdims=True)
-        center_of_gravity = np.sum(weighted_coords, axis=1) / sum_scores
-        return center_of_gravity
+
+        cogs = np.zeros((weighted_coords.shape[0], 2))
+        val_ids = np.where(sum_scores > 0.0)[0]
+        cogs[val_ids] = np.sum(weighted_coords[val_ids], axis=1) / sum_scores[val_ids]
+        return cogs
 
     def animalLength(self, bodyaxis=None):
         if bodyaxis is None:
@@ -164,35 +241,42 @@ class TrackingData(QObject):
         lengths = np.sum(np.sqrt(np.sum(np.diff(bodycoords, axis=1)**2, axis=2)), axis=1)
         return lengths
 
-    def orientation(self, head_node=1, tail_node=5):
+    def orientation(self, head_node=0, tail_node=5):
         bodycoords = self.coordinates()[:, [head_node, tail_node], :]
         vectors = bodycoords[:, 1, :] - bodycoords[:, 0, :]
-        orientations = np.arctan2(vectors[:, 1], vectors[:, 0])
+        orientations = np.arctan2(vectors[:, 0], vectors[:, 1]) * 180 / np.pi
+        orientations[orientations < 0] += 360
         return orientations
 
     def bendedness(self, bodyaxis=None):
+        """
+        Calculate the bendedness of the body axis.
+        Parameters
+        ----------
+        bodyaxis : list of int, optional
+            Indices of the body axis coordinates to consider. If None, defaults to [0, 1, 2, 5].
+        Returns
+        -------
+        numpy.ndarray
+            Array of mean absolute deviations of the body axis points from the head-tail vector.
+        """
+
         if bodyaxis is None:
             bodyaxis = [0, 1, 2, 5]
         bodycoords = self.coordinates()[:, bodyaxis, :]
+        bodycoords = np.concat((bodycoords, np.zeros((bodycoords.shape[0], len(bodyaxis), 1))), axis=2)
         head_tail_vector = bodycoords[:, -1, :] - bodycoords[:, 0, :]
-        head_tail_vector = head_tail_vector[:, np.newaxis ,:]  # cross-product only defined in space, not in 2D
-        head_tail_length = np.linalg.norm(head_tail_vector, axis=1, keepdims=True)  # pythagoras, length of head- tail connection
-        point_axis_vectors = bodycoords - head_tail_vector
-        deviations = np.cross(head_tail_vector, point_axis_vectors)/head_tail_length
-        deviation = np.mean(deviations, axis=1)
-        return deviation
+        point_axis_vector = bodycoords[:,:,:] - bodycoords[:, 0, :][:,np.newaxis,:]
+        htv = head_tail_vector[:,np.newaxis, :]
+        # Pythagoras, length of head- tail connection
+        head_tail_length = np.linalg.norm(head_tail_vector, axis=1, keepdims=True)
+        deviations = np.cross(htv, point_axis_vector)[:,:,-1] / head_tail_length
+        deviations = np.mean(np.abs(deviations), axis=1)
+        return deviations
 
     def __getitem__(self, key):
         return self._data[key]
 
-    # def __setitem__(self, key, value):
-    #     self._data[key] = value
-        """ 
-         self._data.setSelectionRange("index", 0, self._data.numDetections)
-        self._data.assignTracks(tracks)
-        self._timeline.setDetectionData(self._data.data)
-        self.update()
-        """
 
 def main():
     import pandas as pd
@@ -200,6 +284,8 @@ def main():
     import matplotlib.pyplot as plt
     from fixtracks.info import PACKAGE_ROOT
 
+    logging.basicConfig(level=logging.DEBUG, force=True)
+
     def as_dict(df:pd.DataFrame):
         d = {c: df[c].values for c in df.columns}
         d["index"] = df.index.values
@@ -225,31 +311,45 @@ def main():
             count += 1
         return dists
 
-    datafile = PACKAGE_ROOT / "data/merged_small.pkl"
+    def plot_skeleton(positions):
+        skeleton_grid = [(0, 1), (1, 2), (1, 3), (1, 4), (2, 5)]
+        colors = ["tab:red"]
+        colors.extend(["tab:blue"]*5)
+        plt.scatter(positions[:, 0], positions[:, 1], c=colors)
+        for si, ei in skeleton_grid:
+            plt.plot([positions[si, 0], positions[ei, 0]],
+                     [positions[si, 1], positions[ei, 1]], color="tab:green")
+
+    datafile = PACKAGE_ROOT / "data/merged_small_tracked.pkl"
     with open(datafile, "rb") as f:
         df = pickle.load(f)
 
-    data = TrackingData()
-    data.setData(as_dict(df))
+    data = TrackingData(as_dict(df))
+    test_indices = [32, 88, 99, 2593]
+    data.deleteDetections(test_indices)
+    
+    embed()
+    data.deleteDetections(test_indices)
+    data.setSelection(test_indices)
     all_cogs = data.centerOfGravity()
     orientations = data.orientation()
     lengths = data.animalLength()
     frames = data["frame"]
     tracks = data["track"]
-    
-    # bendedness = data.bendedness()
+    bendedness = data.bendedness()
+    indices = data._indices
+    # positions = data.coordinates()[[160388, 160389]]
 
-    embed()
     tracks = data["track"]
     cogs = all_cogs[tracks==1]
     all_dists = neighborDistances(cogs, 2, False)
-    plt.hist(all_dists[1:, 0], bins=1000)
-    print(np.percentile(all_dists[1:, 0], 99))
-    print(np.percentile(all_dists[1:, 0], 1))
-    plt.gca().set_xscale("log")
-    plt.gca().set_yscale("log")
+    # plt.hist(all_dists[1:, 0], bins=1000)
+    # print(np.percentile(all_dists[1:, 0], 99))
+    # print(np.percentile(all_dists[1:, 0], 1))
+    # plt.gca().set_xscale("log")
+    # plt.gca().set_yscale("log")
     # plt.hist(all_dists[1:, 1], bins=100)
-    plt.show()
+    # plt.show()
     # def compute_neighbor_distances(cogs, window=10):
     #     distances = []
     #     for i in range(len(cogs)):
@@ -261,9 +361,6 @@ def main():
     #     return distances
     # print("estimating neighorhood distances")
     # neighbor_distances = compute_neighbor_distances(cogs)
-    embed()
-
-
 
 if __name__ == "__main__":
     main()

fixtracks/widgets/classifier.py

@@ -1,17 +1,287 @@
 import logging
 import numpy as np
 
-from PySide6.QtWidgets import QWidget, QVBoxLayout, QTabWidget,QPushButton, QGraphicsView
-from PySide6.QtCore import Signal
+from PySide6.QtWidgets import QWidget, QVBoxLayout, QTabWidget, QPushButton, QGraphicsView, QTextEdit
+from PySide6.QtWidgets import QSpinBox, QProgressBar, QGridLayout, QLabel, QCheckBox, QDoubleSpinBox
+from PySide6.QtCore import Qt, Signal, Slot, QRunnable, QObject, QThreadPool
 from PySide6.QtGui import QBrush, QColor
+
 import pyqtgraph as pg  # needs to be imported after pyside to not import pyqt
 
 from fixtracks.utils.trackingdata import TrackingData
 
+from IPython import embed
+
+class Detection():
+    def __init__(self, id, frame, track, position, orientation, length, userlabeled, confidence):
+        self.id = id
+        self.frame = frame
+        self.track = track
+        self.position = position
+        self.confidence = confidence
+        self.angle = orientation
+        self.length = length
+        self.userlabeled = userlabeled
+
+
+class WorkerSignals(QObject):
+    message = Signal(str)
+    running = Signal(bool)
+    progress = Signal(int, int, int)
+    currentframe = Signal(int)
+    stopped = Signal(int)
+
+
+class ConsistencyDataLoader(QRunnable):
+    def __init__(self, data):
+        super().__init__()
+        self.signals = WorkerSignals()
+        self.data = data
+        self.bendedness = None
+        self.positions = None
+        self.lengths = None
+        self.orientations = None
+        self.userlabeled = None
+        self.confidence = None
+        self.frames = None
+        self.tracks = None
+
+    @Slot()
+    def run(self):
+        if self.data is None:
+            logging.error("ConsistencyTracker.DataLoader failed. No Data!")
+            return
+        self.positions = self.data.centerOfGravity()
+        self.orientations = self.data.orientation()
+        self.lengths = self.data.animalLength()
+        # self.bendedness = self.data.bendedness()
+        self.userlabeled = self.data["userlabeled"]
+        self.confidence = self.data["confidence"]  # ignore for now, let's see how far this carries.
+        self.frames = self.data["frame"]
+        self.tracks = self.data["track"]
+        self.signals.stopped.emit(0)
+
+
+class ConsistencyWorker(QRunnable):
+
+    def __init__(self, positions, orientations, lengths, bendedness, frames, tracks,
+                 userlabeled, confidence, startframe=0, stoponerror=False, min_confidence=0.0) -> None:
+        super().__init__()
+        self.signals = WorkerSignals()
+        self.positions = positions
+        self.orientations = orientations
+        self.lengths = lengths
+        self.bendedness = bendedness
+        self.userlabeled = userlabeled
+        self.confidence = confidence
+        self._min_confidence = min_confidence
+        self.frames = frames
+        self.tracks = tracks
+        self._startframe = startframe
+        self._stoprequest = False
+        self._stoponerror = stoponerror
+
+    @Slot()
+    def stop(self):
+        self._stoprequest = True
+
+    @Slot()
+    def run(self):
+        def get_detections(frame, indices):
+            detections = []
+            for i in indices:
+                if np.any(self.positions[i] < 0.1):
+                    logging.debug("Encountered probably invalid position %s", str(self.positions[i]))
+                    continue
+                if self._min_confidence > 0.0 and self.confidence[i] < self._min_confidence:
+                    self.tracks[i] = -1
+                    continue
+                d = Detection(i, frame, self.tracks[i], self.positions[i],
+                              self.orientations[i], self.lengths[i],
+                              self.userlabeled[i], self.confidence[i])
+                detections.append(d)
+            return detections
+
+        def assign_by_distance(d):
+            t1_step = d.frame - last_detections[1].frame
+            t2_step = d.frame - last_detections[2].frame
+            if t1_step == 0 or t2_step == 0:
+                print(f"framecount is zero! current frame {f}, last frame {last_detections[1].frame} and {last_detections[2].frame}")
+            distance_to_trackone = np.linalg.norm(d.position - last_detections[1].position) /t1_step
+            distance_to_tracktwo = np.linalg.norm(d.position - last_detections[2].position) /t2_step
+            most_likely_track = np.argmin([distance_to_trackone, distance_to_tracktwo]) + 1
+            distances = np.zeros(2)
+            distances[0] = distance_to_trackone
+            distances[1] = distance_to_tracktwo
+            return most_likely_track, distances
+
+        def assign_by_orientation(d):
+            t1_step = d.frame - last_detections[1].frame
+            t2_step = d.frame - last_detections[2].frame
+            orientationchanges = np.zeros(2)
+            for i in [1, 2]:
+                orientationchanges[i-1] = (last_detections[i].angle - d.angle)
+
+            orientationchanges[orientationchanges > 180] = 360 - orientationchanges[orientationchanges > 180]
+            orientationchanges /= np.array([t1_step, t2_step])
+            most_likely_track = np.argmin(np.abs(orientationchanges)) + 1
+            return most_likely_track, orientationchanges
+
+        def assign_by_length(d):
+            length_differences = np.zeros(2)
+            length_differences[0] = np.abs((last_detections[1].length - d.length))
+            length_differences[1] = np.abs((last_detections[2].length - d.length))
+            most_likely_track = np.argmin(length_differences) + 1
+            return most_likely_track, length_differences
+
+        def check_multiple_detections(detections):
+            if self._min_confidence > 0.0:
+                for i, d in enumerate(detections):
+                    if d.confidence < self._min_confidence:
+                        del detections[i]
+            distances = np.zeros((len(detections), len(detections)))
+            for i, d1 in enumerate(detections):
+                for j, d2 in enumerate(detections):
+                    distances[i, j] = np.abs(np.linalg.norm(d2.position - d1.position))
+            lowest_dist = np.argmin(np.sum(distances, axis=1))
+            del detections[lowest_dist]
+            return detections
+
+        def find_last_userlabeled(startframe):
+            t1index = np.where((self.frames < startframe) & (self.userlabeled) & (self.tracks == 1))[0][-1]
+            t2index = np.where((self.frames < startframe) & (self.userlabeled) & (self.tracks == 2))[0][-1]
+            d1 = Detection(t1index, self.frames[t1index], self.tracks[t1index], self.positions[t1index],
+                           self.orientations[t1index], self.lengths[t1index], self.userlabeled[t1index], 
+                           self.confidence[t1index])
+            d2 = Detection(t1index, self.frames[t2index], self.tracks[t2index], self.positions[t2index],
+                           self.orientations[t2index], self.lengths[t2index], self.userlabeled[t2index],
+                           self.confidence[t1index])
+            last_detections[1] = d1
+            last_detections[2] = d2
+
+        unique_frames = np.unique(self.frames)
+        steps = int((len(unique_frames) - self._startframe) // 100)
+        errors = 0
+        processed = 1
+        progress = 0
+        self._stoprequest = False
+        last_detections = {1: None, 2: None, -1: None}
+        find_last_userlabeled(self._startframe)
+
+        for f in unique_frames[unique_frames >= self._startframe]:
+            if self._stoprequest:
+                break
+            error = False
+            message = ""
+            self.signals.currentframe.emit(f)
+            indices = np.where(self.frames == f)[0]
+            detections = get_detections(f, indices)
+            done = [False, False]
+            if len(detections) == 0:
+                continue
+            if len(detections) > 2:
+                message = f"Frame {f}: More than 2 detections ({len(detections)}) in the same frame!"
+                logging.info("ConsistencyTracker: %s", message)
+                self.signals.message.emit(message)
+                while len(detections) > 2:
+                    detections = check_multiple_detections(detections)
+
+            if len(detections) > 1 and np.any([detections[0].userlabeled, detections[1].userlabeled]):
+                # more than one detection
+                if detections[0].userlabeled and detections[1].userlabeled:
+                    if detections[0].track == detections[1].track:
+                        error = True
+                        message = f"Frame {f}: Classification error both detections in the same frame are assigned to the same track!"
+                        logging.info("ConsistencyTracker: %s", message)
+                        self.signals.message.emit(message)
+                elif detections[0].userlabeled and not detections[1].userlabeled:
+                    detections[1].track = 1 if detections[0].track == 2 else 2
+                elif not detections[0].userlabeled and detections[1].userlabeled:
+                    detections[0].track = 1 if detections[1].track == 2 else 2
+
+                if not error:
+                    last_detections[detections[0].track] = detections[0]
+                    last_detections[detections[1].track] = detections[1]
+                    self.tracks[detections[0].id] = detections[0].track
+                    self.tracks[detections[1].id] = detections[1].track
+                    done[0] = True
+                    done[1] = True
+            elif len(detections) == 1 and detections[0].userlabeled:  # ony one detection and labeled
+                last_detections[detections[0].track] = detections[0]
+                done[0] = True
+            if np.sum(done) == len(detections):
+                continue
+
+            if error and self._stoponerror:
+                self.signals.message.emit(f"Tracking stopped at frame {f}.")
+                break
+            elif error:
+                continue
+            dist_assignments = np.zeros(2, dtype=int)
+            orientation_assignments = np.zeros_like(dist_assignments)
+            length_assignments = np.zeros_like(dist_assignments)
+            distances = np.zeros((2, 2))
+            orientations = np.zeros_like(distances)
+            lengths = np.zeros_like(distances)
+            assignments = np.zeros(2)
+            for i, d in enumerate(detections):
+                dist_assignments[i], distances[i, :] = assign_by_distance(d)
+                orientation_assignments[i], orientations[i,:] = assign_by_orientation(d)
+                length_assignments[i], lengths[i, :] = assign_by_length(d)
+            assignments = dist_assignments  # (dist_assignments * 10 + orientation_assignments + length_assignments) / 3
+
+            error = False
+            temp = {}
+            message = ""
+            if len(detections) > 1:
+                if assignments[0] == assignments[1]:
+                    d.track = -1
+                    error = True
+                    errors += 1
+                    message = f"Frame {f}: Classification error: both detections in the same frame are assigned to the same track!"
+                    break
+                elif assignments[0] != assignments[1]:
+                    detections[0].track = assignments[0]
+                    detections[1].track = assignments[1]
+                    temp[detections[0].track] = detections[0]
+                    temp[detections[1].track] = detections[1]
+                    self.tracks[detections[0].id] = detections[0].track
+                    self.tracks[detections[1].id] = detections[1].track
+            else:
+                if np.abs(np.diff(distances[0,:])) > 50:  # maybe include the time difference into this?
+                    detections[0].track = assignments[0]
+                    temp[detections[0].track] = detections[0]
+                    self.tracks[detections[0].id] = detections[0].track
+                else:
+                    self.tracks[detections[0].id] = -1
+                    message = f"Frame: {f}: Decision based on distance not safe. Track set to unassigned."
+                    error = True
+                    errors += 1
+
+            if not error:
+                for k in temp:
+                    last_detections[temp[k].track] = temp[k]
+            else:
+                logging.info("frame %i: Cannot decide who is who! %s", f, message)
+                for idx in indices:
+                    self.tracks[idx] = -1
+                errors += 1
+                if error and self._stoponerror:
+                    self.signals.message.emit(message)
+                    break
+            processed += 1
+
+            if steps > 0 and f % steps == 0:
+                progress += 1
+                self.signals.progress.emit(progress, processed, errors)
+
+        self.signals.message.emit(f"Tracking stopped at frame {f}.")
+        self.signals.stopped.emit(f)
+
 
 class SizeClassifier(QWidget):
     apply = Signal()
-    name = "SizeClassifier"
+    name = "Size classifier"
 
     def __init__(self, parent=None):
         super().__init__(parent)
@@ -60,7 +330,7 @@ class SizeClassifier(QWidget):
 
     def setCoordinates(self, coordinates):
         self._coordinates = coordinates
-        self._sizes = self.estimate_length(coordinates)
+        self._sizes = self.estimate_length(coordinates, bodyaxis=[0, 1, 2, 5])
         n, e = self.estimate_histogram(self._sizes)
         plot = self._plot_widget.addPlot()
         bgi = pg.BarGraphItem(x0=e[:-1], x1=e[1:], height=n, pen='w', brush=(0,0,255,150))
@@ -84,6 +354,7 @@ class SizeClassifier(QWidget):
         tracks[(self._sizes >= t2lower) & (self._sizes < t2upper)] = 2
         return tracks
 
+
 class NeighborhoodValidator(QWidget):
     apply = Signal()
     name = "Neighborhood Validator"
@@ -198,67 +469,261 @@ class NeighborhoodValidator(QWidget):
 
 class ConsistencyClassifier(QWidget):
     apply = Signal()
-    name = "Consistency classifier"
+    name = "Consistency tracker"
 
     def __init__(self, parent=None):
         super().__init__(parent)
+        self._data = None
+        self._all_pos = None
+        self._all_orientations = None
+        self._all_lengths = None
+        self._all_bendedness = None
+        self._all_scores = None
+        self._confidence = None
+        self._userlabeled = None
+        self._maxframes = 0
+        self._frames = None
+        self._tracks = None
+        self._worker = None
+        self._dataworker = None
+        self._processed_frames = 0
 
+        self._errorlabel = QLabel()
+        self._framelabel = QLabel()
+        self._errorlabel.setStyleSheet("QLabel { color : red; }")
+        self._assignedlabel = QLabel()
+        self._maxframeslabel = QLabel()
+        self._startframe_spinner = QSpinBox()
 
-    def setData(self, keypoints, tracks, frames):
+        self._startbtn = QPushButton("start")
+        self._startbtn.clicked.connect(self.start)
+        self._startbtn.setEnabled(False)
+
+        self._stopbtn = QPushButton("stop")
+        self._stopbtn.clicked.connect(self.stop)
+        self._stopbtn.setEnabled(False)
+
+        self._proceedbtn = QPushButton("proceed")
+        self._proceedbtn.clicked.connect(self.proceed)
+        self._proceedbtn.setEnabled(False)
+
+        self._refreshbtn = QPushButton("refresh")
+        self._refreshbtn.clicked.connect(self.refresh)
+        self._refreshbtn.setEnabled(True)
+
+        self._apply_btn = QPushButton("apply")
+        self._apply_btn.clicked.connect(lambda: self.apply.emit())
+        self._apply_btn.setEnabled(False)
+
+        self._progressbar = QProgressBar()
+        self._progressbar.setMinimum(0)
+        self._progressbar.setMaximum(100)
+
+        self._stoponerror = QCheckBox("Stop processing whenever an error is encountered")
+        self._stoponerror.setToolTip("Stop process upon errors")
+        self._stoponerror.setCheckable(True)
+        self._stoponerror.setChecked(True)
+        self.threadpool = QThreadPool()
+
+        self._ignore_confidence = QCheckBox("Ignore detections widh confidence less than")
+        self._confidence_spinner = QDoubleSpinBox()
+        self._confidence_spinner.setRange(0.0, 1.0)
+        self._confidence_spinner.setSingleStep(0.01)
+        self._confidence_spinner.setDecimals(2)
+        self._confidence_spinner.setValue(0.5)
+        self._messagebox = QTextEdit()
+        self._messagebox.setFocusPolicy(Qt.NoFocus)
+        self._messagebox.setReadOnly(True)
+
+        lyt = QGridLayout()
+        lyt.addWidget(QLabel("Start frame:"), 0, 0 )
+        lyt.addWidget(self._startframe_spinner, 0, 1, 1, 1)
+        lyt.addWidget(QLabel("of"), 0, 2, 1, 1)
+        lyt.addWidget(self._maxframeslabel, 0, 3, 1, 1)
+        lyt.addWidget(self._stoponerror, 1, 0, 1, 3)
+        lyt.addWidget(self._ignore_confidence, 3, 0, 1, 3)
+        lyt.addWidget(self._confidence_spinner, 3, 3, 1, 1)
+        lyt.addWidget(QLabel("Current frame"), 4, 0)
+        lyt.addWidget(self._framelabel, 4, 1)
+        lyt.addWidget(QLabel("(Re-)Assigned"), 5, 0)
+        lyt.addWidget(self._assignedlabel, 5, 1)
+        lyt.addWidget(QLabel("Errors/issues"), 5, 2)
+        lyt.addWidget(self._errorlabel, 5, 3, 1, 1)
+        lyt.addWidget(self._messagebox, 6, 0, 2, 4)
+
+        lyt.addWidget(self._startbtn, 8, 0, 1, 2)
+        lyt.addWidget(self._stopbtn, 8, 2)
+        # lyt.addWidget(self._proceedbtn, 8, 2)
+        lyt.addWidget(self._refreshbtn, 8, 3, 1, 1)
+        lyt.addWidget(self._apply_btn, 9, 0, 1, 4)
+        lyt.addWidget(self._progressbar, 10, 0, 1, 4)
+        self.setLayout(lyt)
+
+    def setData(self, data:TrackingData):
         """Set the data, the classifier/should be working on.
 
         Parameters
         ----------
-        positions : np.ndarray
-            The position estimates, e.g. the center of gravity for each detection
-        tracks : np.ndarray
-            The current track assignment.
-        frames : np.ndarray
-            respective frame.
+        data : Trackingdata
+            The tracking data.
         """
-        def mouseClicked(event):
-            pos = event.pos()
-            if self._plot.sceneBoundingRect().contains(pos):
-                mousePoint = vb.mapSceneToView(pos)
-                print("mouse clicked at", mousePoint)
-                vLine.setPos(mousePoint.x())
-        track2_brush = QBrush(QColor.fromString("green"))
-        track1_brush = QBrush(QColor.fromString("orange"))
-        self._positions = positions
-        self._tracks = tracks
-        self._frames = frames
-        t1_positions = self._positions[self._tracks == 1]
-        t1_frames = self._frames[self._tracks == 1]
-        t1_distances = self.neighborDistances(t1_positions, t1_frames, 1, False)
-        t2_positions = self._positions[self._tracks == 2]
-        t2_frames = self._frames[self._tracks == 2]
-        t2_distances = self.neighborDistances(t2_positions, t2_frames, 1, False)
+        self.setEnabled(False)
+        self._progressbar.setRange(0,0)
+        self._data = data
 
+    @Slot()
+    def data_processed(self):
+        if self._dataworker is not None:
+            self._progressbar.setRange(0,100)
+            self._progressbar.setValue(0)
+            self._all_pos = self._dataworker.positions
+            self._all_orientations = self._dataworker.orientations
+            self._all_lengths = self._dataworker.lengths
+            self._all_bendedness = self._dataworker.bendedness
+            self._userlabeled = self._dataworker.userlabeled
+            self._confidence = self._dataworker.confidence
+            self._frames = self._dataworker.frames
+            self._tracks = self._dataworker.tracks
+            self._dataworker = None
+        if np.sum(self._userlabeled) < 1:
+            msg = "ConsistencyTracker: I need at least 1 user-labeled frame to start with!"
+            logging.error(msg)
+            self._messagebox.append(msg)
+            self.setEnabled(False)
+        else:
+            t1_userlabeled = self._frames[self._userlabeled & (self._tracks == 1)]
+            t2_userlabeled = self._frames[self._userlabeled & (self._tracks == 2)]
+            if any([len(t1_userlabeled) == 0, len(t2_userlabeled)== 0]):
+                self._messagebox.append("Error preparing data! Make sure that the first user-labeled frames contain both tracks!")
+                self.setEnabled(False)
+                return
+            max_startframe = np.min([t1_userlabeled[-1], t2_userlabeled[-1]]) -1
+            first_guess = np.max([t1_userlabeled[0], t2_userlabeled[0]])
+            while first_guess not in t1_userlabeled or first_guess not in t2_userlabeled:
+                first_guess += 1
+            min_startframe = first_guess + 1
+            self._maxframes = np.max(self._frames)
+            self._maxframeslabel.setText(str(self._maxframes))
+            self._startframe_spinner.setMinimum(min_startframe)
+            self._startframe_spinner.setMaximum(max_startframe)
+            self._startframe_spinner.setValue(min_startframe)
+            self._startframe_spinner.setSingleStep(20)
+            self._startframe_spinner.setToolTip(f"Maximum possible start frame: {max_startframe}")
+            self._startbtn.setEnabled(True)
+            self._assignedlabel.setText("0")
+            self._errorlabel.setText("0")
+            self.setEnabled(True)
+
+    @Slot(float)
+    def on_progress(self, value):
+        if self._progressbar is not None:
+            self._progressDialog.setValue(int(value * 100))
+
+    def stop(self):
+        if self._worker is not None:
+            self._worker.stop()
+            self._messagebox.append("Stopping tracking.")
+
+    def start(self):
+        confidence_level = self._confidence_spinner.value() if self._ignore_confidence.isChecked() else 0.0
+        self._startbtn.setEnabled(False)
+        self._refreshbtn.setEnabled(False)
+        self._stopbtn.setEnabled(True)
+        self._worker = ConsistencyWorker(self._all_pos, self._all_orientations, self._all_lengths,
+                                         self._all_bendedness, self._frames, self._tracks, self._userlabeled,
+                                         self._confidence, self._startframe_spinner.value(), self._stoponerror.isChecked(), 
+                                         min_confidence=confidence_level)
+        self._worker.signals.stopped.connect(self.worker_stopped)
+        self._worker.signals.progress.connect(self.worker_progress)
+        self._worker.signals.message.connect(self.worker_error)
+        self._worker.signals.currentframe.connect(self.worker_frame)
+        self._messagebox.append("Tracking in progress ...")
+        self.threadpool.start(self._worker)
+
+    def worker_frame(self, frame):
+        self._framelabel.setText(str(frame))
+
+    def worker_error(self, msg):
+        self._messagebox.append(msg)
+
+    def proceed(self):
+        self.start()
+
+    def refresh(self):
+        self.setEnabled(False)
+        self._dataworker = ConsistencyDataLoader(self._data)
+        self._dataworker.signals.stopped.connect(self.data_processed)
+        self._messagebox.clear()
+        self._messagebox.append("Refreshing...")
+        self.threadpool.start(self._dataworker)
+
+    def worker_progress(self, progress, processed, errors):
+        self._progressbar.setValue(progress)
+        self._errorlabel.setText(str(errors))
+        self._assignedlabel.setText(str(processed))
+
+    def worker_stopped(self, frame):
+        self._startbtn.setEnabled(True)
+        self._proceedbtn.setEnabled(True)
+        self._stopbtn.setEnabled(False)
+        self._apply_btn.setEnabled(True)
+        self._refreshbtn.setEnabled(True)
+        self._startframe_spinner.setValue(frame-1)
+        self._proceedbtn.setEnabled(bool(frame < self._maxframes-1))
+        self._processed_frames = frame
+        self._messagebox.append("... done.")
+
+    def assignedTracks(self):
+        return self._tracks
 
 
 class ClassifierWidget(QTabWidget):
-    apply_sizeclassifier = Signal(np.ndarray)
+    apply_classifier = Signal(np.ndarray)
 
     def __init__(self, parent=None):
         super().__init__(parent)
+        self._data = None
         self._size_classifier = SizeClassifier()
-        self._neigborhood_validator = NeighborhoodValidator()
+        # self._neigborhood_validator = NeighborhoodValidator()
+        self._consistency_tracker = ConsistencyClassifier()
         self.addTab(self._size_classifier, SizeClassifier.name)
-        self.addTab(self._neigborhood_validator, NeighborhoodValidator.name)
+        self.addTab(self._consistency_tracker, ConsistencyClassifier.name)
+        # self.tabBarClicked.connect(self.update)
+        self.currentChanged.connect(self.tabChanged)
         self._size_classifier.apply.connect(self._on_applySizeClassifier)
+        self._consistency_tracker.apply.connect(self._on_applyConsistencyTracker)
 
     def _on_applySizeClassifier(self):
         tracks = self.size_classifier.assignedTracks()
-        self.apply_sizeclassifier.emit(tracks)
+        self.apply_classifier.emit(tracks)
+
+    def _on_applyConsistencyTracker(self):
+        tracks = self._consistency_tracker.assignedTracks()
+        self.apply_classifier.emit(tracks)
 
     @property
     def size_classifier(self):
         return self._size_classifier
 
     @property
-    def neighborhood_validator(self):
-        return self._neigborhood_validator
+    def consistency_tracker(self):
+        return self._consistency_tracker
 
+    @Slot()
+    def tabChanged(self):
+        if isinstance(self.currentWidget(), ConsistencyClassifier):
+            self.consistency_tracker.refresh()
+
+    @Slot()
+    def update(self):
+        if isinstance(self.currentWidget(), ConsistencyClassifier):
+            self.consistency_tracker.refresh()
+
+    def setData(self, data:TrackingData):
+        self._data = data
+        self.consistency_tracker.setData(data)
+        coordinates = self._data.coordinates()
+        self._size_classifier.setCoordinates(coordinates)
 
 def as_dict(df):
     d = {c: df[c].values for c in df.columns}
@@ -270,19 +735,15 @@ def main():
     test_size = False
     import pickle
     from fixtracks.info import PACKAGE_ROOT
-    
-    datafile = PACKAGE_ROOT / "data/merged_small_tracked.pkl"
+
+    datafile = PACKAGE_ROOT / "data/merged_small_beginning.pkl"
 
     with open(datafile, "rb") as f:
         df = pickle.load(f)
-    data = TrackingData()
-    data.setData(as_dict(df))
-
-    positions = data.centerOfGravity()
-    tracks = data["track"]
-    frames = data["frame"]
+    data = TrackingData(as_dict(df))
     coords = data.coordinates()
-
+    cogs = data.centerOfGravity()
+    userlabeled = data["userlabeled"]
     app = QApplication([])
     window = QWidget()
     window.setMinimumSize(200, 200)
@@ -291,7 +752,8 @@ def main():
     #     win.setCoordinates(coords)
     # else:
     w = ClassifierWidget()
-    w.neighborhood_validator.setData(positions, tracks, frames)
+    w.setData(data)
+    # w.size_classifier.setCoordinates(coords)
 
     layout = QVBoxLayout()
     layout.addWidget(w)

fixtracks/widgets/detectiontimeline.py

@@ -1,13 +1,14 @@
 import logging
 import numpy as np
-import pandas as pd
+
 from PySide6.QtCore import Qt
 from PySide6.QtWidgets import QWidget, QVBoxLayout, QSizePolicy, QLabel
-from PySide6.QtWidgets import QGraphicsView, QGraphicsScene, QGraphicsItem, QGraphicsRectItem, QGraphicsLineItem
+from PySide6.QtWidgets import QGraphicsView, QGraphicsScene, QGraphicsItem, QGraphicsRectItem, QGraphicsLineItem, QGraphicsEllipseItem
 from PySide6.QtCore import Qt, QRectF, QRectF
 from PySide6.QtGui import QBrush, QColor, QPen, QFont
 
 from fixtracks.utils.signals import DetectionTimelineSignals
+from fixtracks.utils.trackingdata import TrackingData
 
 
 class Window(QGraphicsRectItem):
@@ -20,16 +21,16 @@ class Window(QGraphicsRectItem):
         self.setBrush(brush)
         self.setZValue(1.0)
         self.setAcceptHoverEvents(True)  # Enable hover events if needed
-        self.setFlags(
-            QGraphicsItem.ItemIsMovable |  # Enables item dragging
-            QGraphicsItem.ItemIsSelectable  # Enables item selection
-        )
+        # self.setFlags(
+        #     QGraphicsItem.ItemIsMovable |  # Enables item dragging
+        #     QGraphicsItem.ItemIsSelectable  # Enables item selection
+        # )
         self._y = y
 
     def setWindowX(self, newx):
         logging.debug("timeline.window: set position to %.3f", newx)
         self.setX(newx)
-        self.signals.windowMoved.emit()
+        # self.signals.windowMoved.emit()
 
     def setWindowWidth(self, newwidth):
         logging.debug("timeline.window: update window width to %f", newwidth)
@@ -37,34 +38,34 @@ class Window(QGraphicsRectItem):
         r = self.rect()
         r.setWidth(newwidth)
         self.setRect(r)
-        self.signals.windowMoved.emit()
+        # self.signals.windowMoved.emit()
 
     def setWindow(self, newx:float, newwidth:float):
-        def setWindow(self, newx: float, newwidth: float):
-            """
-            Update the window to the specified range.
-            Parameters
-            ----------
-            newx : float
-                The new x-coordinate of the window.
-            newwidth : float
-                The new width of the window.
-            Returns
-            -------
-            None
-            """
-
+        """
+        Update the window to the specified range.
+        Parameters
+        ----------
+        newx : float
+            The new x-coordinate of the window.
+        newwidth : float
+            The new width of the window.
+        Returns
+        -------
+        None
+        """
         logging.debug("timeline.window: update window to range %.5f to %.5f", newx, newwidth)
         self._width = newwidth
         r = self.rect()
         self.setRect(newx, r.y(), self._width, r.height())
-        self.signals.windowMoved.emit()
+        self.update()
+        # self.signals.windowMoved.emit()
 
     def mouseMoveEvent(self, event):
         super().mouseMoveEvent(event)
 
     def mousePressEvent(self, event):
         self.setCursor(Qt.ClosedHandCursor)
+        # print(event.pos())
         super().mousePressEvent(event)
 
     def mouseReleaseEvent(self, event):
@@ -77,7 +78,7 @@ class Window(QGraphicsRectItem):
         if r.y() != self._y:
             self.setY(self._y)
         super().mouseReleaseEvent(event)
-        self.signals.windowMoved.emit()
+        self.signals.manualMove.emit()
 
     def hoverEnterEvent(self, event):
         super().hoverEnterEvent(event)
@@ -86,40 +87,45 @@ class Window(QGraphicsRectItem):
 class DetectionTimeline(QWidget):
     signals = DetectionTimelineSignals()
 
-    def __init__(self, detectiondata=None, trackone_id=1, tracktwo_id=2, parent=None):
+    def __init__(self, trackone_id=1, tracktwo_id=2, parent=None):
         super().__init__(parent)
         self._trackone = trackone_id
         self._tracktwo = tracktwo_id
-        self._data = detectiondata
+        self._data = None
         self._rangeStart = 0.0
         self._rangeStop = 0.005
         self._total_width = 2000
-        self._stepCount = 200
+        self._stepCount = 1000
         self._bg_brush = QBrush(QColor(20, 20, 20, 255))
         transparent_brush = QBrush(QColor(200, 200, 200, 64))
         self._white_pen = QPen(QColor.fromString("white"))
         self._white_pen.setWidth(0.1)
         self._t1_pen = QPen(QColor.fromString("orange"))
-        self._t1_pen.setWidth(2)
+        self._t1_pen.setWidth(1)
         self._t2_pen = QPen(QColor(0, 255, 0, 255))
-        self._t2_pen.setWidth(2)
+        self._t2_pen.setWidth(1)
         self._other_pen = QPen(QColor.fromString("red"))
-        self._other_pen.setWidth(2)
-        axis_pen = QPen(QColor.fromString("white"))
-        axis_pen.setWidth(2)
+        self._other_pen.setWidth(1)
+        window_pen = QPen(QColor.fromString("white"))
+        window_pen.setWidth(2)
+        self._user_brush = QBrush(QColor.fromString("white"))
+        user_pen = QPen(QColor.fromString("white"))
+        user_pen.setWidth(2)
+
         font = QFont()
         font.setPointSize(15)
-        font.setBold(False)
+        font.setBold(True)
 
-        self._window = Window(0, 0, 100, 60, axis_pen, transparent_brush)
-        self._window.signals.windowMoved.connect(self.on_windowMoved)
+        self._window = Window(0, 0, 100, 60, window_pen, transparent_brush)
+        self._window.signals.manualMove.connect(self.on_windowMoved)
 
-        self._scene = QGraphicsScene(QRectF(0, 0, self._total_width, 55.))
+        self._scene = QGraphicsScene(QRectF(0, 0, self._total_width, 85.))
         self._scene.setBackgroundBrush(self._bg_brush)
         self._scene.addItem(self._window)
+        self._scene.mousePressEvent = self.on_sceneMousePress
 
         self._view = QGraphicsView()
-        # self._view.setRenderHints(QPainter.RenderHint.Antialiasing | QPainter.RenderHint.SmoothPixmapTransform);
+        # self._view.setRenderHints(QPainter.RenderHint.Antialiasing | QPainter.RenderHint.SmoothPixmapTransform)
         self._view.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
         self._view.setScene(self._scene)
         self._view.fitInView(self._scene.sceneRect(), aspectRadioMode=Qt.AspectRatioMode.KeepAspectRatio)
@@ -127,15 +133,19 @@ class DetectionTimeline(QWidget):
 
         t1_label = self._scene.addText("track 1", font)
         t1_label.setDefaultTextColor(self._t1_pen.color())
-        t1_label.setPos(0, 0)
+        t1_label.setPos(0, 50)
         t2_label = self._scene.addText("track 2", font)
         t2_label.setFont(font)
         t2_label.setDefaultTextColor(self._t2_pen.color())
-        t2_label.setPos(0, 17)
+        t2_label.setPos(100, 50)
         other_label = self._scene.addText("unassigned", font)
         other_label.setFont(font)
         other_label.setDefaultTextColor(self._other_pen.color())
-        other_label.setPos(0, 30)
+        other_label.setPos(200, 50)
+        user_label = self._scene.addText("user-labeled", font)
+        user_label.setFont(font)
+        user_label.setDefaultTextColor(user_pen.color())
+        user_label.setPos(350, 50)
 
         self._position_label = QLabel("")
         f = self._position_label.font()
@@ -143,54 +153,62 @@ class DetectionTimeline(QWidget):
         self._position_label.setFont(f)
 
         layout = QVBoxLayout()
+        layout.setSpacing(0)
+        layout.setContentsMargins(5, 2, 5, 2)
         layout.addWidget(self._view)
         layout.addWidget(self._position_label, Qt.AlignmentFlag.AlignRight)
         self.setLayout(layout)
-
-        if self._data is not None:
-            self.draw_coverage()
         # self.setMaximumHeight(100)
         # self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed)
 
-    def setDetectionData(self, data):
-        self._data = data
+    def on_sceneMousePress(self, event):
+        scene_pos = event.scenePos()
+        relpos = scene_pos.x() / self._total_width
+        relpos = 0 if relpos < 0.0 else relpos
+        relpos = 2000/self._total_width if scene_pos.x() > self._total_width else relpos
+        self.signals.moveRequest.emit(relpos)
+        logging.debug("Timeline: Scene clicked at position: %.2f, %.2f --> rel x-pos %.3f", scene_pos.x(), scene_pos.y(), relpos)
+    
+    def clear(self):
         for i in self._scene.items():
-            if isinstance(i, QGraphicsLineItem):
+            if isinstance(i, (QGraphicsLineItem, QGraphicsEllipseItem)):
                 self._scene.removeItem(i)
+
+    def setData(self, data:TrackingData):
+        logging.debug("Timeline: setData!")
+        self._data = data
+        self.update()
+        self.resizeEvent(None)
+
+    def update(self):
+        self.clear()
         self.draw_coverage()
 
     def draw_coverage(self):
-        # FIXME this must be disentangled. timeline should not have to deal with two different ways of data storage
-        if isinstance(self._data, pd.DataFrame):
-            maxframe = np.max(self._data.frame.values)
-
-            bins = np.linspace(0, maxframe, self._stepCount)
-            pos = np.linspace(0, self._scene.width(), self._stepCount)
-            track1_frames = self._data.frame.values[self._data.track == self._trackone]
-            track2_frames = self._data.frame.values[self._data.track == self._tracktwo]
-            other_frames = self._data.frame.values[(self._data.track != self._trackone) &
-                                                   (self._data.track != self._tracktwo)]
-        elif isinstance(self._data, dict):
+        logging.debug("Timeline: drawCoverage!")
+        if isinstance(self._data, TrackingData):
             maxframe = np.max(self._data["frame"])
             bins = np.linspace(0, maxframe, self._stepCount)
-            pos = np.linspace(0, self._scene.width(), self._stepCount)
+            pos = np.linspace(0, self._scene.width(), self._stepCount)  # of the vertical dashes is this correct?
             track1_frames = self._data["frame"][self._data["track"] == self._trackone]
             track2_frames = self._data["frame"][self._data["track"] == self._tracktwo]
             other_frames = self._data["frame"][(self._data["track"] != self._trackone) &
                                                (self._data["track"] != self._tracktwo)]
+            userlabeled = self._data["frame"][self._data["userlabeled"]]
         else:
+            print("Data is not trackingdata")
             return
         t1_coverage, _ = np.histogram(track1_frames, bins=bins)
-        t1_coverage = t1_coverage > 0
         t2_coverage, _ = np.histogram(track2_frames, bins=bins)
-        t2_coverage = t2_coverage > 0
         other_coverage, _ = np.histogram(other_frames, bins=bins)
-        other_coverage = other_coverage > 0
+        labeled_coverage, _ = np.histogram(userlabeled, bins=bins)
 
-        for i in range(len(t1_coverage)-1):
+        for i in range(len(bins)-1):
             if t1_coverage[i]: self._scene.addLine(pos[i], 0, pos[i], 15., pen=self._t1_pen)
             if t2_coverage[i]: self._scene.addLine(pos[i], 17, pos[i], 32., pen=self._t2_pen)
             if other_coverage[i]: self._scene.addLine(pos[i], 34, pos[i], 49., pen=self._other_pen)
+            if other_coverage[i]: self._scene.addLine(pos[i], 34, pos[i], 49., pen=self._other_pen)
+            if labeled_coverage[i]: self._scene.addEllipse(pos[i]-2, 52, 4, 4, brush=self._user_brush)
 
     def updatePositionLabel(self):
         start = np.round(self._rangeStart * 100, 4)
@@ -207,6 +225,7 @@ class DetectionTimeline(QWidget):
 
     def fit_scene_to_view(self):
         """Scale the image to fit the QGraphicsView."""
+        logging.debug("Timeline: fit scene to view")
         self._view.fitInView(self._scene.sceneRect(), Qt.KeepAspectRatio)
 
     def resizeEvent(self, event):
@@ -289,6 +308,11 @@ def main():
         view.setWindowPos(0.0)
         print(view.windowBounds())
 
+    def as_dict(df):
+        d = {c: df[c].values for c in df.columns}
+        d["index"] = df.index.values
+        return d
+
     import pickle
     import numpy as np
     from PySide6.QtWidgets import QApplication, QPushButton, QHBoxLayout
@@ -298,13 +322,17 @@ def main():
     datafile = PACKAGE_ROOT / "data/merged_small.pkl"
     with open(datafile, "rb") as f:
         df = pickle.load(f)
-
+    data = TrackingData(as_dict(df))
+    data.setSelection(np.arange(0,100, 1))
+    data.setUserLabeledStatus(True)
     start_x = 0.1
     app = QApplication([])
     window = QWidget()
     window.setMinimumSize(200, 75)
 
-    view = DetectionTimeline(df)
+    view = DetectionTimeline()
+    view.setData(data)
+
     fwdBtn = QPushButton(">>")
     fwdBtn.clicked.connect(lambda: fwd(0.5))
     zeroBtn = QPushButton("0->|")
@@ -313,12 +341,14 @@ def main():
     backBtn.clicked.connect(lambda: back(0.2))
 
     btnLyt = QHBoxLayout()
+    btnLyt.setSpacing(1)
     btnLyt.addWidget(backBtn)
     btnLyt.addWidget(zeroBtn)
     btnLyt.addWidget(fwdBtn)
 
     view.setWindowPos(start_x)
     layout = QVBoxLayout()
+    layout.setSpacing(1)
     layout.addWidget(view)
     layout.addLayout(btnLyt)
     window.setLayout(layout)

fixtracks/widgets/detectionview.py

@@ -1,20 +1,14 @@
-import enum
 import logging
 import numpy as np
 
 from PySide6.QtWidgets import QWidget, QVBoxLayout, QSizePolicy, QGraphicsView, QGraphicsScene, QGraphicsEllipseItem, QGraphicsRectItem
 from PySide6.QtCore import Qt, QPointF, QRectF, QPointF
-from PySide6.QtGui import QPixmap, QBrush, QColor, QImage
+from PySide6.QtGui import QPixmap, QBrush, QColor, QImage, QPen
 
 from fixtracks.info import PACKAGE_ROOT
 from fixtracks.utils.signals import DetectionSignals, DetectionViewSignals, DetectionSceneSignals
-
-
-class DetectionData(enum.Enum):
-    ID = 0
-    FRAME = 1
-    COORDINATES = 2
-    TRACK_ID = 3
+from fixtracks.utils.enums import DetectionData, Tracks
+from fixtracks.utils.trackingdata import TrackingData
 
 
 class Detection(QGraphicsEllipseItem):
@@ -86,6 +80,7 @@ class DetectionView(QWidget):
     def __init__(self, parent=None):
         super().__init__(parent)
         self._img = None
+        self._data = None
         self._pixmapitem = None
         self._scene = DetectionScene()
         # self.setRenderHint(QGraphicsView.RenderFlag.Ren Antialiasing)
@@ -97,7 +92,6 @@ class DetectionView(QWidget):
         self._minZoom = 0.1
         self._maxZoom = 10
         self._currentZoom = 1.0
-
         lyt = QVBoxLayout()
         lyt.addWidget(self._view)
         self.setLayout(lyt)
@@ -123,6 +117,9 @@ class DetectionView(QWidget):
         self._view.setScene(self._scene)
         self._view.fitInView(self._scene.sceneRect(), aspectRadioMode=Qt.AspectRatioMode.KeepAspectRatio)
 
+    def setData(self, data:TrackingData):
+        self._data = data
+
     def clearDetections(self):
         items = self._scene.items()
         if items is not None:
@@ -131,23 +128,39 @@ class DetectionView(QWidget):
                     self._scene.removeItem(it)
                     del it
 
-    def addDetections(self, coordinates:np.array, track_ids:np.array, detection_ids:np.array, frames: np.array,
-                      keypoint:int, brush:QBrush):
+    def updateDetections(self, keypoint=-1):
+        logging.info("DetectionView.updateDetections!")
+        self.clearDetections()
+        if self._data is None:
+            return
+        frames = self._data.selectedData("frame")
+        tracks = self._data.selectedData("track")
+        ids = self._data.selectedData("index")
+        coordinates = self._data.coordinates(selection=True)
+        centercoordinates = self._data.centerOfGravity(selection=True)
+        userlabeled = self._data.selectedData("userlabeled")
+        scores = self._data.selectedData("confidence")
+
         image_rect = self._pixmapitem.boundingRect() if self._pixmapitem is not None else QRectF(0,0,0,0)
-        num_detections = coordinates.shape[0]
-        for i in range(num_detections):
-            x = coordinates[i, keypoint, 0]
-            y = coordinates[i, keypoint, 1]
-            c = brush.color()
+        num_detections = len(frames)
+        for i, (id, f, t, l, s) in enumerate(zip(ids, frames, tracks, userlabeled, scores)):
+            c = Tracks.fromValue(t).toColor()
             c.setAlpha(int(i * 255 / num_detections))
-            brush.setColor(c)
-            item = Detection(image_rect.left() + x, image_rect.top() + y, 20, 20, brush=brush)
-            item.setData(DetectionData.TRACK_ID.value, track_ids[i])
-            item.setData(DetectionData.ID.value, detection_ids[i])
+            if keypoint >= 0:
+                x = coordinates[i, keypoint, 0]
+                y = coordinates[i, keypoint, 1]
+            else:
+                x = centercoordinates[i, 0]
+                y = centercoordinates[i, 1]
+
+            item = Detection(image_rect.left() + x, image_rect.top() + y, 20, 20, brush=QBrush(c))
+            item.setData(DetectionData.TRACK_ID.value, t)
+            item.setData(DetectionData.ID.value, id)
             item.setData(DetectionData.COORDINATES.value, coordinates[i, :, :])
-            item.setData(DetectionData.FRAME.value, frames[i])
+            item.setData(DetectionData.FRAME.value, f)
+            item.setData(DetectionData.USERLABELED.value, l)
+            item.setData(DetectionData.SCORE.value, s)
             item = self._scene.addItem(item)
-        # logging.debug("DetectionView: Number of items in scene: %i", len(self._scene.items()))
 
     def fit_image_to_view(self):
         """Scale the image to fit the QGraphicsView."""
@@ -166,7 +179,7 @@ class DetectionView(QWidget):
 def main():
     def items_selected(items):
         print("items selected")
-
+    # FIXME The following code will no longer work...
     import pickle
     import numpy as np
     from IPython import embed
@@ -205,7 +218,7 @@ def main():
     view.setImage(img)
     view.addDetections(bg_coords, bg_tracks, bg_ids, background_brush)
     view.addDetections(focus_coords, focus_tracks, focus_ids, focus_brush)
-    view.addDetections(scnd_coords, scnd_tracks,  scnd_ids, second_brush)
+    view.addDetections(scnd_coords, scnd_tracks, scnd_ids, second_brush)
     window.setLayout(layout)
     window.show()
     app.exec()

fixtracks/widgets/selection_control.py

@@ -0,0 +1,263 @@
+import logging
+import numpy as np
+
+from PySide6.QtCore import Qt, Signal, QSize
+from PySide6.QtGui import QFont
+from PySide6.QtWidgets import QWidget, QLabel, QPushButton, QSizePolicy
+from PySide6.QtWidgets import QGridLayout, QVBoxLayout, QApplication
+
+from fixtracks.utils.styles import pushBtnStyle
+
+class SelectionControls(QWidget):
+    fwd = Signal(float)
+    back = Signal(float)
+    assignOne = Signal()
+    assignTwo = Signal()
+    assignOther = Signal()
+    accept = Signal()
+    accept_until = Signal()
+    unaccept = Signal()
+    delete = Signal()
+    revertall = Signal()
+
+    def __init__(self, parent = None,):
+        super().__init__(parent)
+        font = QFont()
+        font.setBold(True)
+        font.setPointSize(10)
+
+        fullstep = 1.0
+        halfstep = 0.5
+        quarterstep = 0.25
+
+        backBtn = QPushButton("|<<")
+        backBtn.setFont(font)
+        backBtn.setShortcut(Qt.Key.Key_Left)
+        backBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
+        backBtn.setToolTip(f"Go back to previous window ({backBtn.shortcut().toString()})")
+        backBtn.setStyleSheet(pushBtnStyle("darkgray"))
+        backBtn.clicked.connect(lambda: self.on_Back(fullstep))
+        halfstepBackBtn = QPushButton("<<")
+        halfstepBackBtn.setFont(font)
+        halfstepBackBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
+        halfstepBackBtn.setShortcut(Qt.KeyboardModifier.AltModifier | Qt.Key.Key_Left)
+        halfstepBackBtn.setToolTip(f"Go back by half a window ({halfstepBackBtn.shortcut().toString()})")
+        halfstepBackBtn.setStyleSheet(pushBtnStyle("darkgray"))
+        halfstepBackBtn.clicked.connect(lambda: self.on_Back(halfstep))
+        quarterstepBackBtn = QPushButton("<")
+        quarterstepBackBtn.setFont(font)
+        quarterstepBackBtn.setShortcut(Qt.KeyboardModifier.ShiftModifier | Qt.Key.Key_Left)
+        quarterstepBackBtn.setToolTip(f"Go back by a quarter window ({quarterstepBackBtn.shortcut().toString()})")
+        quarterstepBackBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
+        quarterstepBackBtn.setStyleSheet(pushBtnStyle("darkgray"))
+        quarterstepBackBtn.clicked.connect(lambda: self.on_Back(quarterstep))
+
+        fwdBtn = QPushButton(">>|")
+        fwdBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
+        fwdBtn.setShortcut(Qt.Key.Key_Right)
+        fwdBtn.setFont(font)
+        fwdBtn.setToolTip(f"Proceed to next window ({fwdBtn.shortcut().toString()})")
+        fwdBtn.setStyleSheet(pushBtnStyle("darkgray"))
+        fwdBtn.clicked.connect(lambda: self.on_Fwd(fullstep))
+        halfstepFwdBtn = QPushButton(">>")
+        halfstepFwdBtn.setFont(font)
+        halfstepFwdBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
+        halfstepFwdBtn.setShortcut(Qt.KeyboardModifier.AltModifier | Qt.Key.Key_Right)
+        halfstepFwdBtn.setToolTip(f"Proceed by half a window ({halfstepFwdBtn.shortcut().toString()})")
+        halfstepFwdBtn.setStyleSheet(pushBtnStyle("darkgray"))
+        halfstepFwdBtn.clicked.connect(lambda: self.on_Fwd(halfstep))
+        quarterstepFwdBtn = QPushButton(">")
+        quarterstepFwdBtn.setFont(font)
+        quarterstepFwdBtn.setShortcut(Qt.KeyboardModifier.ShiftModifier | Qt.Key.Key_Right)
+        quarterstepFwdBtn.setToolTip(f"Proceed by a quarter window ({quarterstepFwdBtn.shortcut().toString()})")
+        quarterstepFwdBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
+        quarterstepFwdBtn.setStyleSheet(pushBtnStyle("darkgray"))
+        quarterstepFwdBtn.clicked.connect(lambda: self.on_Fwd(quarterstep))
+
+        assignOneBtn = QPushButton("Track One")
+        assignOneBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
+        assignOneBtn.setStyleSheet(pushBtnStyle("orange"))
+        assignOneBtn.setShortcut("Ctrl+1")
+        assignOneBtn.setToolTip(f"Assign current selection to Track One ({assignOneBtn.shortcut().toString()})")
+        assignOneBtn.setFont(font)
+        assignOneBtn.clicked.connect(self.on_TrackOne)
+
+        assignTwoBtn = QPushButton("Track Two")
+        assignTwoBtn.setShortcut("Ctrl+2")
+        assignTwoBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
+        assignTwoBtn.setStyleSheet(pushBtnStyle("green"))
+        assignTwoBtn.setToolTip(f"Assign current selection to Track Two ({assignTwoBtn.shortcut().toString()})")
+        assignTwoBtn.setFont(font)
+        assignTwoBtn.clicked.connect(self.on_TrackTwo)
+
+        assignOtherBtn = QPushButton("Other")
+        assignOtherBtn.setShortcut("Ctrl+0")
+        assignOtherBtn.setFont(font)
+        assignOtherBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
+        assignOtherBtn.setStyleSheet(pushBtnStyle("gray"))
+        assignOtherBtn.setToolTip(f"Assign current selection to Unassigned ({assignOtherBtn.shortcut().toString()})")
+        assignOtherBtn.clicked.connect(self.on_TrackOther)
+
+        acceptBtn = QPushButton("accept")
+        acceptBtn.setFont(font)
+        acceptBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
+        acceptBtn.setStyleSheet(pushBtnStyle("darkgray"))
+        acceptBtn.setToolTip(f"Accept assignments of current selection as TRUE, Hold shift while clicking to accept all until here.")
+        acceptBtn.clicked.connect(self.on_Accept)
+
+        unacceptBtn = QPushButton("un-accept")
+        unacceptBtn.setFont(font)
+        unacceptBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
+        unacceptBtn.setStyleSheet(pushBtnStyle("darkgray"))
+        unacceptBtn.setToolTip(f"Revoke current selection TRUE status")
+        unacceptBtn.clicked.connect(self.on_Unaccept)
+
+        deleteBtn = QPushButton("delete")
+        deleteBtn.setFont(font)
+        deleteBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
+        deleteBtn.setStyleSheet(pushBtnStyle("red"))
+        deleteBtn.setToolTip(f"DANGERZONE! Delete current selection of detections!")
+        deleteBtn.setShortcut("Ctrl+D")
+        deleteBtn.clicked.connect(self.on_Delete)
+        deleteBtn.setEnabled(False)
+
+        revertBtn = QPushButton("revert assignments")
+        revertBtn.setFont(font)
+        revertBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
+        revertBtn.setStyleSheet(pushBtnStyle("red"))
+        revertBtn.setToolTip(f"DANGERZONE! Remove ALL assignments!")
+        revertBtn.clicked.connect(self.on_RevertAll)
+
+        self.tone_selection = QLabel("0")
+        self.ttwo_selection = QLabel("0")
+        self.tother_selection = QLabel("0")
+        self.startframe = QLabel("0")
+        self.endframe = QLabel("0")
+        self._total = 0
+
+        grid = QGridLayout()
+        grid.addWidget(backBtn, 0, 0, 3, 2)
+        grid.addWidget(fwdBtn, 0, 6, 3, 2)
+        grid.addWidget(halfstepBackBtn, 3, 0, 2, 2)
+        grid.addWidget(halfstepFwdBtn, 3, 6, 2, 2)
+        grid.addWidget(quarterstepBackBtn, 5, 0, 2, 2)
+        grid.addWidget(quarterstepFwdBtn, 5, 6, 2, 2)
+
+        cwLabel = QLabel("Current window:")
+        cwLabel.setFont(font)
+        grid.addWidget(cwLabel, 0, 2, 1, 4)
+        grid.addWidget(QLabel("start:"), 1, 2, 1, 2)
+        grid.addWidget(self.startframe, 1, 4, 1, 2, Qt.AlignmentFlag.AlignRight)
+        grid.addWidget(QLabel("end:"), 2, 2, 1, 2)
+        grid.addWidget(self.endframe, 2, 4, 1, 2, Qt.AlignmentFlag.AlignRight)
+
+        csLabel = QLabel("Current selection:")
+        csLabel.setFont(font)
+        grid.addWidget(csLabel, 3, 2, 1, 4)
+        grid.addWidget(QLabel("track One:"), 4, 2, 1, 2)
+        grid.addWidget(self.tone_selection, 4, 4, 1, 2, Qt.AlignmentFlag.AlignRight)
+        grid.addWidget(QLabel("track Two:"), 5, 2, 1, 2)
+        grid.addWidget(self.ttwo_selection, 5, 4, 1, 2, Qt.AlignmentFlag.AlignRight)
+        grid.addWidget(QLabel("Unassigned:"), 6, 2, 1, 2)
+        grid.addWidget(self.tother_selection, 6, 4, 1, 2, Qt.AlignmentFlag.AlignRight)
+
+        grid.addWidget(assignOneBtn, 7, 0, 1, 3)
+        grid.addWidget(assignOtherBtn, 7, 3, 1, 2)
+        grid.addWidget(assignTwoBtn, 7, 5, 1, 3)
+        grid.addWidget(acceptBtn, 8, 0, 1, 4)
+        grid.addWidget(unacceptBtn, 8, 4, 1, 4)
+        grid.addWidget(deleteBtn, 9, 0, 1, 4)
+        grid.addWidget(revertBtn, 9, 4, 1, 4)
+
+        grid.setColumnStretch(0, 1)
+        grid.setColumnStretch(7, 1)
+        self.setLayout(grid)
+        # self.setMaximumSize(QSize(500, 500))
+
+    def setWindow(self, start:int=0, end:int=0):
+        self.startframe.setText(f"{start:.0f}")
+        self.endframe.setText(f"{end:g}")
+
+    def _updateNumbers(self, track):
+        labels = {1: self.tone_selection, 2: self.ttwo_selection, 3: self.tother_selection}
+        for k in labels:
+            if k == track:
+                labels[k].setText(str(self._total))
+            else:
+                labels[k].setText("0")
+
+    def on_Fwd(self, stepsize):
+        logging.debug("SelectionControls: forward step by %.2f", stepsize)
+        self.fwd.emit(stepsize)
+
+    def on_Back(self, stepsize):
+        logging.debug("SelectionControls: backward step by %.2f", stepsize)
+        self.back.emit(stepsize)
+
+    def on_TrackOne(self):
+        logging.debug("SelectionControl: TrackONEBtn")
+        self.assignOne.emit()
+        self._updateNumbers(1)
+
+    def on_TrackTwo(self):
+        logging.debug("SelectionControl: TrackTWOBtn")
+        self.assignTwo.emit()
+        self._updateNumbers(2)
+
+    def on_TrackOther(self):
+        logging.debug("SelectionControl: TrackOtherBtn")
+        self.assignOther.emit()
+        self._updateNumbers(3)
+
+    def on_Accept(self):
+        logging.debug("SelectionControl: accept AssignmentBtn")
+        modifiers = QApplication.keyboardModifiers()
+        if modifiers == Qt.KeyboardModifier.ShiftModifier:
+            logging.debug("Shift key was pressed during accept")
+            self.accept_until.emit()
+        else:
+            self.accept.emit()
+
+    def on_Unaccept(self):
+        logging.debug("SelectionControl: revoke user assignmentBtn")
+        self.unaccept.emit()
+
+    def on_RevertAll(self):
+        logging.debug("SelectionControl: revert Btn")
+        self.revertall.emit()
+
+    def on_Delete(self):
+        logging.debug("SelectionControl: delete Btn")
+        self.delete.emit()
+
+    def setSelectedTracks(self, tracks):
+        logging.debug("SelectionControl: setSelectedTracks")
+        if tracks is not None:
+            tone = np.sum(tracks == 1)
+            ttwo = np.sum(tracks == 2)
+        else:
+            tone = 0
+            ttwo = 0
+
+        self.tone_selection.setText(str(tone))
+        self.ttwo_selection.setText(str(ttwo))
+        self.tother_selection.setText(str(len(tracks) - tone - ttwo if tracks is not None else 0))
+        self._total = len(tracks) if tracks is not None else 0
+
+
+def main():
+    from PySide6.QtWidgets import QApplication
+    logging.basicConfig(level=logging.DEBUG, force=True)
+
+    app = QApplication([])
+    window = QWidget()
+    window.setMinimumSize(200, 200)
+    layout = QVBoxLayout()
+    controls = SelectionControls()
+    layout.addWidget(controls)
+    window.setLayout(layout)
+    window.show()
+    app.exec()
+
+if __name__ == "__main__":
+    main()

fixtracks/widgets/skeleton.py

@@ -6,7 +6,7 @@ from PySide6.QtWidgets import QGraphicsScene, QGraphicsEllipseItem, QGraphicsRec
 from PySide6.QtCore import Qt
 from PySide6.QtGui import QBrush, QColor, QPen, QPainter, QFont
 
-from fixtracks.widgets.detectionview import DetectionData
+from fixtracks.utils.enums import DetectionData
 
 class Skeleton(QGraphicsRectItem):
     skeleton_grid = [(0, 1), (1, 2), (1, 3), (1, 4), (2, 5)]
@@ -94,7 +94,8 @@ class SkeletonWidget(QWidget):
             i = s.data(DetectionData.ID.value)
             t = s.data(DetectionData.TRACK_ID.value)
             f = s.data(DetectionData.FRAME.value)
-            self._info_label.setText(f"Id {i}, track {t} on frame {f}, length {l:.1f} px")
+            sc = s.data(DetectionData.SCORE.value)
+            self._info_label.setText(f"Id {i}, track {t} on frame {f}, length {l:.1f} px, confidence {sc:.2f}")
         else:
             self._info_label.setText("")
 
@@ -129,7 +130,7 @@ class SkeletonWidget(QWidget):
         self._scene.setSceneRect(self._minx, self._miny, self._maxx - self._minx, self._maxy - self._miny)
         self._view.fitInView(self._scene.sceneRect(), Qt.AspectRatioMode.KeepAspectRatio)
 
-    def addSkeleton(self, coords, detection_id, frame, track, brush, update=True):
+    def addSkeleton(self, coords, detection_id, frame, track, score, brush, update=True):
         def check_extent(x, y, w, h):
             if x == 0 and y == 0:
                 return
@@ -157,22 +158,27 @@ class SkeletonWidget(QWidget):
         item.setData(DetectionData.ID.value, detection_id)
         item.setData(DetectionData.TRACK_ID.value, track)
         item.setData(DetectionData.FRAME.value, frame)
+        item.setData(DetectionData.SCORE.value, score)
         self._skeletons.append(item)
         if update:
             self.update()
 
     def addSkeletons(self, coordinates:np.ndarray, detection_ids:np.ndarray,
-                     frames:np.ndarray, tracks:np.ndarray, brush:QBrush):
+                     frames:np.ndarray, tracks:np.ndarray, scores:np.ndarray,
+                     brush:QBrush):
         num_detections = 0 if coordinates is None else coordinates.shape[0]
         logging.debug("SkeletonWidget: add %i Skeletons", num_detections)
+        if num_detections < 1:
+            return
         sorting = np.argsort(frames)
         coordinates = coordinates[sorting,:, :]
         detection_ids = detection_ids[sorting]
         frames = frames[sorting]
         tracks = tracks[sorting]
+        scores = scores[sorting]
         for i in range(num_detections):
             self.addSkeleton(coordinates[i,:,:], detection_ids[i], frames[i],
-                             tracks[i], brush=brush, update=False)
+                             tracks[i], scores[i], brush=brush, update=False)
         self.update()
 
     # def addSkeleton(self, coords, detection_id, brush):

fixtracks/widgets/tracks.py

@@ -1,11 +1,11 @@
 import logging
 import numpy as np
-import pandas as pd
 
-from PySide6.QtCore import Qt, QThreadPool, Signal, QSize, QObject
-from PySide6.QtGui import QImage, QBrush, QColor, QFont
+from PySide6.QtCore import Qt, QThreadPool, Signal
+from PySide6.QtGui import QImage, QBrush, QColor
 from PySide6.QtWidgets import QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, QSizePolicy, QComboBox
-from PySide6.QtWidgets import QSpinBox, QSpacerItem, QProgressBar, QSplitter, QGridLayout, QFileDialog, QGridLayout
+from PySide6.QtWidgets import QSpinBox, QSpacerItem, QProgressBar, QSplitter, QFileDialog, QMessageBox
+
 
 from fixtracks.utils.reader import PickleLoader
 from fixtracks.utils.writer import PickleWriter
@@ -14,149 +14,7 @@ from fixtracks.widgets.detectionview import DetectionView, DetectionData
 from fixtracks.widgets.detectiontimeline import DetectionTimeline
 from fixtracks.widgets.skeleton import SkeletonWidget
 from fixtracks.widgets.classifier import ClassifierWidget
-
-
-class SelectionControls(QWidget):
-    fwd = Signal(float)
-    back = Signal(float)
-    assignOne = Signal()
-    assignTwo = Signal()
-    assignOther = Signal()
-
-    def __init__(self, parent = None,):
-        super().__init__(parent)
-        font = QFont()
-        font.setBold(True)
-        font.setPointSize(10)
-        fullstep = 1.0
-        halfstep = 0.5
-        quarterstep = 0.25
-
-        backBtn = QPushButton("|<<")
-        backBtn.setFont(font)
-        backBtn.setShortcut(Qt.Key.Key_Left)
-        backBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
-        backBtn.setToolTip(f"Go back to previous window ({backBtn.shortcut().toString()})")
-        backBtn.clicked.connect(lambda: self.on_Back(fullstep))
-        halfstepBackBtn = QPushButton("<<")
-        halfstepBackBtn.setFont(font)
-        halfstepBackBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
-        halfstepBackBtn.setShortcut(Qt.KeyboardModifier.AltModifier | Qt.Key.Key_Left)
-        halfstepBackBtn.setToolTip(f"Go back by half a window ({halfstepBackBtn.shortcut().toString()})")
-        halfstepBackBtn.clicked.connect(lambda: self.on_Back(halfstep))
-        quarterstepBackBtn = QPushButton("<")
-        quarterstepBackBtn.setFont(font)
-        quarterstepBackBtn.setShortcut(Qt.KeyboardModifier.ShiftModifier | Qt.Key.Key_Left)
-        quarterstepBackBtn.setToolTip(f"Go back by a quarter window ({quarterstepBackBtn.shortcut().toString()})")
-        quarterstepBackBtn.clicked.connect(lambda: self.on_Back(quarterstep))
-
-        fwdBtn = QPushButton(">>|")
-        fwdBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
-        fwdBtn.setShortcut(Qt.Key.Key_Right)
-        fwdBtn.setFont(font)
-        fwdBtn.setToolTip(f"Proceed to next window ({fwdBtn.shortcut().toString()})")
-        fwdBtn.clicked.connect(lambda: self.on_Fwd(fullstep))
-        halfstepFwdBtn = QPushButton(">>")
-        halfstepFwdBtn.setFont(font)
-        halfstepFwdBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
-        halfstepFwdBtn.setShortcut(Qt.KeyboardModifier.AltModifier | Qt.Key.Key_Right)
-        halfstepFwdBtn.setToolTip(f"Proceed by half a window ({halfstepFwdBtn.shortcut().toString()})")
-        halfstepFwdBtn.clicked.connect(lambda: self.on_Fwd(halfstep))
-        quarterstepFwdBtn = QPushButton(">")
-        quarterstepFwdBtn.setFont(font)
-        quarterstepFwdBtn.setShortcut(Qt.KeyboardModifier.ShiftModifier | Qt.Key.Key_Right)
-        quarterstepFwdBtn.setToolTip(f"Proceed by a quarter window ({quarterstepFwdBtn.shortcut().toString()})")
-        quarterstepFwdBtn.clicked.connect(lambda: self.on_Fwd(quarterstep))
-
-        assignOneBtn = QPushButton("Track One")
-        assignOneBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
-        assignOneBtn.setStyleSheet("QPushButton { background-color: orange; }")
-        assignOneBtn.setShortcut("Ctrl+1")
-        assignOneBtn.setToolTip(f"Assign current selection to Track One ({assignOneBtn.shortcut().toString()})")
-        assignOneBtn.setFont(font)
-        assignOneBtn.clicked.connect(self.on_TrackOne)
-
-        assignTwoBtn = QPushButton("Track Two")
-        assignTwoBtn.setShortcut("Ctrl+2")
-        assignTwoBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
-        assignTwoBtn.setStyleSheet("QPushButton { background-color: green; }")
-        assignTwoBtn.setToolTip(f"Assign current selection to Track Two ({assignTwoBtn.shortcut().toString()})")
-        assignTwoBtn.setFont(font)
-        assignTwoBtn.clicked.connect(self.on_TrackTwo)
-
-        assignOtherBtn = QPushButton("Other")
-        assignOtherBtn.setShortcut("Ctrl+0")
-        assignOtherBtn.setFont(font)
-        assignOtherBtn.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
-        assignOtherBtn.setStyleSheet("QPushButton { background-color: red; }")
-        assignOtherBtn.setToolTip(f"Assign current selection to Unassigned ({assignOtherBtn.shortcut().toString()})")
-        assignOtherBtn.clicked.connect(self.on_TrackOther)
-
-        self.tone_selection = QLabel("0")
-        self.ttwo_selection = QLabel("0")
-        self.tother_selection = QLabel("0")
-        self._total = 0
-
-        grid = QGridLayout()
-        grid.addWidget(backBtn, 0, 0, 2, 2)
-        grid.addWidget(halfstepBackBtn, 2, 0, 1, 2)
-        grid.addWidget(quarterstepBackBtn, 3, 0, 1, 2)
-        grid.addWidget(fwdBtn, 0, 6, 2, 2)
-        grid.addWidget(halfstepFwdBtn, 2, 6, 1, 2)
-        grid.addWidget(quarterstepFwdBtn, 3, 6, 1, 2)
-        grid.addWidget(QLabel("Current selection:"), 0, 2, 1, 4)
-        grid.addWidget(QLabel("Track One:"), 1, 2, 1, 3)
-        grid.addWidget(self.tone_selection, 1, 5, 1, 1)
-        grid.addWidget(QLabel("Track Two:"), 2, 2, 1, 3)
-        grid.addWidget(self.ttwo_selection, 2, 5, 1, 1)
-        grid.addWidget(QLabel("Unassigned:"), 3, 2, 1, 3)
-        grid.addWidget(self.tother_selection, 3, 5, 1, 1)
-
-        grid.addWidget(assignOneBtn, 4, 0, 4, 3)
-        grid.addWidget(assignOtherBtn, 4, 3, 4, 2)
-        grid.addWidget(assignTwoBtn, 4, 5, 4, 3)
-        grid.setColumnStretch(0, 1)
-        grid.setColumnStretch(7, 1)
-        self.setLayout(grid)
-        self.setMaximumSize(QSize(400, 200))
-
-    def _updateNumbers(self, track):
-        labels = {1: self.tone_selection, 2: self.ttwo_selection, 3: self.tother_selection}
-        for k in labels:
-            if k == track:
-                labels[k].setText(str(self._total))
-            else:
-                labels[k].setText("0")
-
-    def on_Fwd(self, stepsize):
-        logging.debug("SelectionControls: forward step by %.2f", stepsize)
-        self.fwd.emit(stepsize)
-
-    def on_Back(self, stepsize):
-        logging.debug("SelectionControls: backward step by %.2f", stepsize)
-        self.back.emit(stepsize)
-
-    def on_TrackOne(self):
-        self.assignOne.emit()
-        self._updateNumbers(1)
-
-    def on_TrackTwo(self):
-        self.assignTwo.emit()
-        self._updateNumbers(2)
-
-    def on_TrackOther(self):
-        self.assignOther.emit()
-        self._updateNumbers(3)
-
-    def setSelectedTracks(self, tracks):
-        logging.debug("SelectionControl: setSelectedTracks")
-        tone = np.sum(tracks == 1)
-        ttwo = np.sum(tracks == 2)
-        self.tone_selection.setText(str(tone))
-        self.ttwo_selection.setText(str(ttwo))
-        self.tother_selection.setText(str(len(tracks) - tone - ttwo))
-        self._total = len(tracks)
-
+from fixtracks.widgets.selection_control import SelectionControls
 
 class FixTracks(QWidget):
     back = Signal()
@@ -170,24 +28,15 @@ class FixTracks(QWidget):
         self._threadpool = QThreadPool()
         self._reader = None
         self._image = None
-        self._clear_detections = True
         self._currentWindowPos = 0  # in frames
         self._currentWindowWidth = 0  # in frames
         self._maxframes = 0
-        self._data = TrackingData()
-        self._brushes = {"assigned_left": QBrush(QColor.fromString("orange")),
-                         "assigned_right": QBrush(QColor.fromString("green")),
-                         "unassigned": QBrush(QColor.fromString("red"))
-        }
+        self._manualmove = False
+        self._data = None
+
         self._detectionView = DetectionView()
         self._detectionView.signals.itemsSelected.connect(self.on_detectionsSelected)
         self._skeleton = SkeletonWidget()
-        # self._skeleton.setMaximumSize(QSize(400, 400))
-        top_splitter = QSplitter(Qt.Orientation.Horizontal)
-        top_splitter.addWidget(self._detectionView)
-        top_splitter.addWidget(self._skeleton)
-        top_splitter.setStretchFactor(0, 2)
-        top_splitter.setStretchFactor(1, 1)
 
         self._progress_bar = QProgressBar(self)
         self._progress_bar.setMaximumHeight(20)
@@ -196,25 +45,42 @@ class FixTracks(QWidget):
 
         self._timeline = DetectionTimeline()
         self._timeline.signals.windowMoved.connect(self.on_windowChanged)
+        self._timeline.signals.moveRequest.connect(self.on_moveRequest)
 
         self._windowspinner = QSpinBox()
         self._windowspinner.setRange(10, 10000)
         self._windowspinner.setSingleStep(50)
         self._windowspinner.setValue(500)
         self._windowspinner.valueChanged.connect(self.on_windowSizeChanged)
-        # self._timeline.setWindowWidth(0.01)
+
         self._keypointcombo = QComboBox()
         self._keypointcombo.currentIndexChanged.connect(self.on_keypointSelected)
 
-        combo_layout = QGridLayout()
-        combo_layout.addWidget(QLabel("Window:"), 0, 0)
-        combo_layout.addWidget(self._windowspinner, 0, 1)
-        combo_layout.addWidget(QLabel("Keypoint:"), 1, 0)
-        combo_layout.addWidget(self._keypointcombo, 1, 1)
+        self._goto_spinner = QSpinBox()
+        self._goto_spinner.setSingleStep(1)
 
-        timelinebox = QHBoxLayout()
-        timelinebox.addWidget(self._timeline)
+        self._gotobtn = QPushButton("go!")
+        self._gotobtn.setToolTip("Jump to a given frame")
+        self._gotobtn.clicked.connect(self.on_goto)
+
+        combo_layout = QHBoxLayout()
+        combo_layout.addWidget(QLabel("Window width:"))
+        combo_layout.addWidget(self._windowspinner)
+        combo_layout.addWidget(QLabel("frames"))
+        combo_layout.addItem(QSpacerItem(10, 10, QSizePolicy.Policy.Fixed, QSizePolicy.Policy.Fixed))
+        combo_layout.addWidget(QLabel("Keypoint:"))
+        combo_layout.addWidget(self._keypointcombo)
+        combo_layout.addItem(QSpacerItem(10, 10, QSizePolicy.Policy.Fixed, QSizePolicy.Policy.Fixed))
+        combo_layout.addWidget(QLabel("Jump to frame:"))
+        combo_layout.addWidget(self._goto_spinner)
+        combo_layout.addWidget(self._gotobtn)
+        combo_layout.addItem(QSpacerItem(100, 10, QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed))
+        combo_layout.setSpacing(1)
+
+        timelinebox = QVBoxLayout()
+        timelinebox.setSpacing(2)
         timelinebox.addLayout(combo_layout)
+        timelinebox.addWidget(self._timeline)
 
         self._controls_widget = SelectionControls()
         self._controls_widget.assignOne.connect(self.on_assignOne)
@@ -222,6 +88,11 @@ class FixTracks(QWidget):
         self._controls_widget.assignOther.connect(self.on_assignOther)
         self._controls_widget.fwd.connect(self.on_forward)
         self._controls_widget.back.connect(self.on_backward)
+        self._controls_widget.accept.connect(self.on_setUserFlag)
+        self._controls_widget.accept_until.connect(self.on_setUserFlagsUntil)
+        self._controls_widget.unaccept.connect(self.on_unsetUserFlag)
+        self._controls_widget.delete.connect(self.on_deleteDetection)
+        self._controls_widget.revertall.connect(self.on_revertUserFlags)
 
         self._saveBtn = QPushButton("Save")
         self._saveBtn.setShortcut("Ctrl+S")
@@ -244,6 +115,7 @@ class FixTracks(QWidget):
         data_selection_box.addWidget(QLabel("Select data file"))
         data_selection_box.addWidget(self._data_combo)
         data_selection_box.addItem(QSpacerItem(100, 10, QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed))
+        data_selection_box.setSpacing(0)
 
         btnBox = QHBoxLayout()
         btnBox.setAlignment(Qt.AlignmentFlag.AlignLeft)
@@ -254,14 +126,19 @@ class FixTracks(QWidget):
         btnBox.addWidget(self._saveBtn)
 
         self._classifier = ClassifierWidget()
-        self._classifier.apply_sizeclassifier.connect(self.on_classifyBySize)
+        self._classifier.apply_classifier.connect(self.on_autoClassify)
         self._classifier.setMaximumWidth(500)
         cntrlBox = QHBoxLayout()
         cntrlBox.addWidget(self._classifier)
         cntrlBox.addWidget(self._controls_widget, alignment=Qt.AlignmentFlag.AlignCenter)
-        cntrlBox.addItem(QSpacerItem(300, 100, QSizePolicy.Policy.Fixed, QSizePolicy.Policy.Expanding))
+        cntrlBox.addWidget(self._skeleton)
+        cntrlBox.addItem(QSpacerItem(50, 100, QSizePolicy.Policy.Fixed, QSizePolicy.Policy.Expanding))
+        cntrlBox.setSpacing(0)
+        cntrlBox.setContentsMargins(0,0,0,0)
 
         vbox = QVBoxLayout()
+        vbox.setSpacing(0)
+        vbox.setContentsMargins(0,0,0,0)
         vbox.addLayout(timelinebox)
         vbox.addLayout(cntrlBox)
         vbox.addLayout(btnBox)
@@ -269,19 +146,22 @@ class FixTracks(QWidget):
         container.setLayout(vbox)
 
         splitter = QSplitter(Qt.Orientation.Vertical)
-        splitter.addWidget(top_splitter)
+        splitter.addWidget(self._detectionView)
         splitter.addWidget(container)
         splitter.setStretchFactor(0, 3)
         splitter.setStretchFactor(1, 1)
+
         layout = QVBoxLayout()
         layout.addLayout(data_selection_box)
         layout.addWidget(splitter)
+        layout.setSpacing(0)
+        layout.setContentsMargins(5,2,2,5)
         self.setLayout(layout)
 
-    def on_classifyBySize(self, tracks):
+    def on_autoClassify(self, tracks):
         self._data.setSelectionRange("index", 0, self._data.numDetections)
         self._data.assignTracks(tracks)
-        self._timeline.setDetectionData(self._data.data)
+        self._timeline.update()
         self.update()
 
     def on_dataSelection(self):
@@ -301,38 +181,20 @@ class FixTracks(QWidget):
         self._detectionView.setImage(img)
 
     def update(self):
-        def update_detectionView(df, name):
-            if len(df) == 0:
-                return
-            keypoint = self._keypointcombo.currentIndex()
-            coords = np.stack(df["keypoints"].values).astype(np.float32)[:, :,:]
-            tracks = df["track"].values.astype(int)
-            ids = df.index.values.astype(int)
-            frames = df["frame"].values.astype(int)
-            self._detectionView.addDetections(coords, tracks, ids, frames, keypoint, self._brushes[name])
+        kp = self._keypointcombo.currentText().lower()
+        if len(kp) == 0:
+            return
+        kpi = -1 if "center" in kp else int(kp)
 
         start_frame = self._currentWindowPos
         stop_frame = start_frame + self._currentWindowWidth
+
+        self._timeline.setWindow(start_frame / self._maxframes,
+                                 self._currentWindowWidth/self._maxframes)
         logging.debug("Tracks:update: Updating View for detection range %i, %i frames", start_frame, stop_frame)
         self._data.setSelectionRange("frame", start_frame, stop_frame)
-        frames = self._data.selectedData("frame")
-        tracks = self._data.selectedData("track")
-        keypoints = self._data.selectedData("keypoints")
-        index = self._data.selectedData("index")
-
-        df = pd.DataFrame({"frame": frames,
-                           "track": tracks,
-                           "keypoints": keypoints},
-                            index=index)
-        assigned_left = df[(df.track == self.trackone_id)]
-        assigned_right = df[(df.track == self.tracktwo_id)]
-        unassigned = df[(df.track != self.trackone_id) & (df.track != self.tracktwo_id)]
-
-        if self._clear_detections:
-            self._detectionView.clearDetections()
-        update_detectionView(unassigned, "unassigned")
-        update_detectionView(assigned_left, "assigned_left")
-        update_detectionView(assigned_right, "assigned_right")
+        self._controls_widget.setWindow(start_frame, stop_frame)
+        self._detectionView.updateDetections(kpi)
 
     @property
     def fileList(self):
@@ -359,6 +221,7 @@ class FixTracks(QWidget):
 
     def populateKeypointCombo(self, num_keypoints):
         self._keypointcombo.clear()
+        self._keypointcombo.addItem("Center")
         for i in range(num_keypoints):
             self._keypointcombo.addItem(str(i))
         self._keypointcombo.setCurrentIndex(0)
@@ -368,23 +231,20 @@ class FixTracks(QWidget):
         self._progress_bar.setRange(0, 100)
         self._progress_bar.setValue(0)
         if state and self._reader is not None:
-            self._data.setData(self._reader.asdict)
+            self._data = TrackingData(self._reader.asdict)
+            self._saveBtn.setEnabled(True)
             self._currentWindowPos = 0
             self._currentWindowWidth = self._windowspinner.value()
-            self._maxframes = self._data.max("frame")
+            self._maxframes = np.max(self._data["frame"])
+            self._goto_spinner.setMaximum(self._maxframes)
             self.populateKeypointCombo(self._data.numKeypoints())
-            self._timeline.setDetectionData(self._data.data)
-            self._timeline.setWindow(self._currentWindowPos / self._maxframes,
-                                     self._currentWindowWidth / self._maxframes)
-            coordinates = self._data.coordinates()
-            positions = self._data.centerOfGravity()
-            tracks = self._data["track"]
-            frames = self._data["frame"]
-            self._classifier.size_classifier.setCoordinates(coordinates)
-            self._classifier.neighborhood_validator.setData(positions, tracks, frames)
+            self._timeline.setData(self._data)
+            # self._timeline.setWindow(self._currentWindowPos / self._maxframes,
+            #                          self._currentWindowWidth / self._maxframes)
+            self._detectionView.setData(self._data)
+            self._classifier.setData(self._data)
             self.update()
-            self._saveBtn.setEnabled(True)
-            logging.info("Finished loading data: %i frames, %i detections", self._maxframes, len(positions))
+            logging.info("Finished loading data: %i frames", self._maxframes)
 
     def on_keypointSelected(self):
         self.update()
@@ -415,25 +275,78 @@ class FixTracks(QWidget):
 
     def on_assignOne(self):
         logging.debug("Assigning user selection to track One")
-        self._data.assignUserSelection(self.trackone_id)
-        self._timeline.setDetectionData(self._data.data)
+        self._data.setTrack(self.trackone_id)
+        self._timeline.update()
         self.update()
 
     def on_assignTwo(self):
         logging.debug("Assigning user selection to track Two")
-        self._data.assignUserSelection(self.tracktwo_id)
-        self._timeline.setDetectionData(self._data.data)
+        self._data.setTrack(self.tracktwo_id)
+        self._timeline.update()
         self.update()
 
     def on_assignOther(self):
         logging.debug("Assigning user selection to track Other")
-        self._data.assignUserSelection(self.trackother_id)
-        self._timeline.setDetectionData(self._data.data)
+        self._data.setTrack(self.trackother_id, False)
+        self._timeline.update()
+        self.update()
+
+    def on_setUserFlag(self):
+        self._data.setUserLabeledStatus(True)
+        self._timeline.update()
+        self.update()
+
+    def on_setUserFlagsUntil(self):
+        self._data.setSelectionRange("frame", 0, self._currentWindowPos + self._currentWindowWidth)
+        self._data.setUserLabeledStatus(True)
+        self._timeline.update()
+        self.update()
+
+    def on_unsetUserFlag(self):
+        logging.debug("Tracks:unsetUserFlag")
+        self._data.setUserLabeledStatus(False)
+        self._timeline.update()
+        self.update()
+
+    def on_revertUserFlags(self):
+        logging.debug("Tracks:revert ALL UserFlags and track assignments")
+        msg_box = QMessageBox()
+        msg_box.setIcon(QMessageBox.Icon.Warning)
+        msg_box.setText(f"Are you sure you want to revert ALL track assignments?")
+        msg_box.setStandardButtons(QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No)
+        msg_box.setDefaultButton(QMessageBox.StandardButton.No)
+        ret = msg_box.exec()
+
+        if ret == QMessageBox.StandardButton.Yes:
+            self._data.revertUserLabeledStatus()
+            self._data.revertTrackAssignments()
+        self._timeline.update()
+        self.update()
+
+    def on_deleteDetection(self):
+        logging.info("Tracks:deleting detections!")
+        msg_box = QMessageBox()
+        msg_box.setIcon(QMessageBox.Icon.Warning)
+        msg_box.setText(f"Are you sure you want to delete the selected ({len(self._data.selectionIndices)})detections?")
+        msg_box.setStandardButtons(QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No)
+        msg_box.setDefaultButton(QMessageBox.StandardButton.No)
+        ret = msg_box.exec()
+
+        if ret == QMessageBox.StandardButton.Yes:
+            self._data.deleteDetections()
+        self._timeline.update()
         self.update()
 
     def on_windowChanged(self):
-        logging.info("Timeline reports window change ")
-        self._currentWindowPos = np.round(self._timeline.rangeStart * self._maxframes)
+        logging.debug("Tracks:Timeline reports window change ")
+        if not self._manualmove:
+            self._currentWindowPos = np.round(self._timeline.rangeStart * self._maxframes)
+            self.update()
+        self._manualmove = False
+
+    def on_moveRequest(self, pos):
+        new_pos = int(np.round(pos * self._maxframes))
+        self._currentWindowPos = new_pos
         self.update()
 
     def on_windowSizeChanged(self, value):
@@ -446,13 +359,29 @@ class FixTracks(QWidget):
         """
         self._currentWindowWidth = value
         logging.debug("Tracks:OnWindowSizeChanged %i franes", value)
-        self._timeline.setWindowWidth(self._currentWindowWidth / self._maxframes)
+        # if self._maxframes == 0:
+        #     self._timeline.setWindowWidth(self._currentWindowWidth / 2000)
+        # else:
+        #     self._timeline.setWindowWidth(self._currentWindowWidth / self._maxframes)
+        # self._controls_widget.setSelectedTracks(None)
+        self.update()
+
+    def on_goto(self):
+        target = self._goto_spinner.value()
+        if target > self._maxframes - self._currentWindowWidth:
+           target =  self._maxframes - self._currentWindowWidth
+        logging.info("Jump to frame %i", target)
+        self._currentWindowPos = target
+        self._timeline.setWindow(self._currentWindowPos / self._maxframes, 
+                                 self._currentWindowWidth / self._maxframes)
+        self.update()
 
     def on_detectionsSelected(self, detections):
-        logging.debug("Tracks: Detections selected")
+        logging.debug("Tracks: %i Detections selected", len(detections))
         tracks = np.zeros(len(detections), dtype=int)
         ids = np.zeros_like(tracks)
         frames = np.zeros_like(tracks)
+        scores = np.zeros(tracks.shape, dtype=float)
         coordinates = None
         if len(detections) > 0:
             c = detections[0].data(DetectionData.COORDINATES.value)
@@ -463,18 +392,22 @@ class FixTracks(QWidget):
             ids[i] = d.data(DetectionData.ID.value)
             frames[i] = d.data(DetectionData.FRAME.value)
             coordinates[i, :, :] = d.data(DetectionData.COORDINATES.value)
-        self._data.setUserSelection(ids)
+            scores[i] = d.data(DetectionData.SCORE.value)
+        self._data.setSelection(ids)
         self._controls_widget.setSelectedTracks(tracks)
         self._skeleton.clear()
-        self._skeleton.addSkeletons(coordinates, ids, frames, tracks, QBrush(QColor(10, 255, 65, 255)))
-        self.update()
+        self._skeleton.addSkeletons(coordinates, ids, frames, tracks, scores, QBrush(QColor(10, 255, 65, 255)))
 
     def moveWindow(self, stepsize):
-        self._clear_detections = True
-        step = np.round(stepsize * (self._currentWindowWidth))
-        new_start_frame = self._currentWindowPos + step
-        self._timeline.setWindowPos(new_start_frame / self._maxframes)
+        logging.info("Tracks.moveWindow: move window with stepsize %.2f", stepsize)
+        self._manualmove = True
+        new_start_frame = self._currentWindowPos + np.round(stepsize * self._currentWindowWidth)
+        if new_start_frame < 0:
+            new_start_frame = 0
+        elif new_start_frame + self._currentWindowWidth > self._maxframes:
+            new_start_frame = self._maxframes - self._currentWindowWidth
         self._currentWindowPos = new_start_frame
+        self._controls_widget.setSelectedTracks(None)
         self.update()
 
     def on_forward(self, stepsize):

main.py

@@ -1,65 +0,0 @@
-"""
-pyside6-rcc resources.qrc -o resources.py
-
-"""
-import sys
-import platform
-import logging
-from PySide6.QtWidgets import QApplication
-from PySide6.QtCore import QSettings
-from PySide6.QtGui import QIcon, QPalette
-
-from fixtracks import fixtracks, info
-
-logging.basicConfig(level=logging.DEBUG, force=True)
-
-
-def is_dark_mode(app: QApplication) -> bool:
-    palette = app.palette()
-    # Check the brightness of the window text and base colors
-    text_color = palette.color(QPalette.ColorRole.WindowText)
-    base_color = palette.color(QPalette.ColorRole.Base)
-
-    # Calculate brightness (0 for dark, 255 for bright)
-    def brightness(color):
-        return (color.red() * 299 + color.green() * 587 + color.blue() * 114) // 1000
-
-    return brightness(base_color) < brightness(text_color)
-
-
-# import resources  # needs to be imported somewhere in the project to be picked up by qt
-
-if platform.system() == "Windows":
-    # from PySide6.QtWinExtras import QtWin
-    myappid = f"{info.organization_name}.{info.application_version}"
-    # QtWin.setCurrentProcessExplicitAppUserModelID(myappid)
-settings = QSettings()
-width = int(settings.value("app/width", 1024))
-height = int(settings.value("app/height", 768))
-x = int(settings.value("app/pos_x", 100))
-y = int(settings.value("app/pos_y", 100))
-app = QApplication(sys.argv)
-app.setApplicationName(info.application_name)
-app.setApplicationVersion(str(info.application_version))
-app.setOrganizationDomain(info.organization_name)
-
-# if platform.system() == 'Linux':
-# icn = QIcon(":/icons/app_icon")
-# app.setWindowIcon(icn)
-# Create a Qt widget, which will be our window.
-window = fixtracks.MainWindow(is_dark_mode(app))
-window.setGeometry(100, 100, 1024, 768)
-window.setWindowTitle("FixTracks")
-window.setMinimumWidth(1024)
-window.setMinimumHeight(768)
-window.resize(width, height)
-window.move(x, y)
-window.show()
-
-# Start the event loop.
-app.exec()
-pos = window.pos()
-settings.setValue("app/width", window.width())
-settings.setValue("app/height", window.height())
-settings.setValue("app/pos_x", pos.x())
-settings.setValue("app/pos_y", pos.y())

pyproject.toml

@@ -2,7 +2,7 @@
 name = "fixtracks"
 version = "0.1.0"
 description = "A project to fix track metadata"
-authors = ["Your Name <your.email@example.com>"]
+authors = ["Your Name <jan.grewe@uni-tuebingen.de>"]
 license = "MIT"
 
 [tool.poetry.dependencies]