[skeleton] correctly catch zero len detection

fixtracks/widgets/classifier.py

@@ -1,7 +1,8 @@
 import logging
 import numpy as np
 
-from PySide6.QtWidgets import QWidget, QVBoxLayout, QTabWidget, QPushButton, QGraphicsView, QSpinBox, QProgressBar, QGridLayout, QLabel
+from PySide6.QtWidgets import QWidget, QVBoxLayout, QTabWidget, QPushButton, QGraphicsView
+from PySide6.QtWidgets import QSpinBox, QProgressBar, QGridLayout, QLabel, QCheckBox
 from PySide6.QtCore import 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
@@ -13,12 +14,13 @@ class WorkerSignals(QObject):
     error = Signal(str)
     running = Signal(bool)
     progress = Signal(int, int, int)
-    finished = Signal(bool)
+    stopped = Signal(int)
 
 class ConsistencyWorker(QRunnable):
     signals = WorkerSignals()
 
-    def __init__(self, positions, orientations, lengths, bendedness, frames, tracks, startframe=0) -> None:
+    def __init__(self, positions, orientations, lengths, bendedness, frames, tracks,
+                 startframe=0, stoponerror=False) -> None:
         super().__init__()
         self.positions = positions
         self.orientations = orientations
@@ -28,25 +30,29 @@ class ConsistencyWorker(QRunnable):
         self.tracks = tracks
         self._startframe = startframe
         self._stoprequest = False
+        self._stoponerror = stoponerror
 
     @Slot()
-    def cancel(self):
+    def stop(self):
         self._stoprequest = True
 
     @Slot()
     def run(self):
-        last_pos = [self.positions[self.tracks == 1][0], self.positions[self.tracks == 2][0]]
-        last_frame = [self.frames[self.tracks == 1][0], self.frames[self.tracks == 2][0]]
-        last_angle = [self.orientations[self.tracks == 1][0], self.orientations[self.tracks == 2][0]]
+        last_pos = [self.positions[(self.tracks == 1) & (self.frames <= self._startframe)][-1],
+                    self.positions[(self.tracks == 2) & (self.frames <= self._startframe)][-1]]
+        last_frame = [self.frames[(self.tracks == 1) & (self.frames <= self._startframe)][-1],
+                      self.frames[(self.tracks == 2) & (self.frames <= self._startframe)][-1]]
+        # last_angle = [self.orientations[self.tracks == 1][0], self.orientations[self.tracks == 2][0]]
         errors = 0
-        processed = 0
-        self._stoprequest = False
-        maxframes = np.max(self.frames)
-        steps = int((maxframes - self._startframe) // 100)
+        processed = 1
         progress = 0
         assignment_error = False
-        for f in range(self._startframe, np.max(self.frames), 1):
-            processed += 1
+        self._stoprequest = False
+        maxframes = np.max(self.frames)
+        startframe = np.max(last_frame)
+        steps = int((maxframes - startframe) // 200)
+
+        for f in range(startframe + 1, maxframes, 1):
             if self._stoprequest:
                 break
             indices = np.where(self.frames == f)[0]
@@ -59,15 +65,17 @@ class ConsistencyWorker(QRunnable):
                     self.tracks[idx] = 2
                     last_frame[1] = f
                     last_pos[1] = p
-                    last_angle[1] = self.orientations[idx]
+                    # last_angle[1] = self.orientations[idx]
                     continue
                 if f < last_frame[1]:
                     last_frame[0] = f
                     last_pos[0] = p
-                    last_angle[0] = self.orientations[idx]
+                    # last_angle[0] = self.orientations[idx]
                     self.tracks[idx] = 1
                     continue
                 # else, we have already seen track one and track two entries
+                if f - last_frame[0] == 0 or f - last_frame[1] == 0:
+                    print(f"framecount is zero! current frame {f}, last frame {last_frame[0]} and {last_frame[1]}")
                 distance_to_trackone = np.linalg.norm(p - last_pos[0])/(f - last_frame[0])
                 distance_to_tracktwo = np.linalg.norm(p - last_pos[1])/(f - last_frame[1])
                 most_likely_track = np.argmin([distance_to_trackone, distance_to_tracktwo]) + 1
@@ -79,6 +87,8 @@ class ConsistencyWorker(QRunnable):
                 logging.warning("frame %i: Issues assigning based on distances %s", f, str(distances))
                 assignment_error = True
                 errors += 1
+                if self._stoponerror:
+                    break
             else:
                 processed += 1
             for i, idx in enumerate(indices):
@@ -89,17 +99,16 @@ class ConsistencyWorker(QRunnable):
                     last_pos[assignments[i]-1] = pp[i]
                     last_frame[assignments[i]-1] = f
             assignment_error = False
-            if f % steps == 0:
+            if steps > 0 and f % steps == 0:
                 progress += 1
                 self.signals.progress.emit(progress, processed, errors)
 
-        self.signals.finished.emit(True)
-
+        self.signals.stopped.emit(f)
 
 
 class SizeClassifier(QWidget):
     apply = Signal()
-    name = "SizeClassifier"
+    name = "Size classifier"
 
     def __init__(self, parent=None):
         super().__init__(parent)
@@ -286,12 +295,12 @@ 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_cogs = None
+        self._all_pos = None
         self._all_orientations = None
         self._all_lengths = None
         self._all_bendedness = None
@@ -299,41 +308,61 @@ class ConsistencyClassifier(QWidget):
         self._frames = None
         self._tracks = None
         self._worker = None
+        self._processed_frames = 0
 
         self._errorlabel = QLabel()
         self._errorlabel.setStyleSheet("QLabel { color : red; }")
         self._assignedlabel = QLabel()
+        self._maxframeslabel = QLabel()
         self._startframe_spinner = QSpinBox()
 
-        self._startbtn = QPushButton("run")
-        self._startbtn.clicked.connect(self.run)
+        self._startbtn = QPushButton("start")
+        self._startbtn.clicked.connect(self.start)
         self._startbtn.setEnabled(False)
 
-        self._cancelbtn = QPushButton("cancel")
-        self._cancelbtn.clicked.connect(self.cancel)
-        self._cancelbtn.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._apply_btn.clicked.connect(lambda: self.apply.emit())
-        self._apply_btn.setEnabled(False)
+        self._stoponerror = QCheckBox("Stop processing whenever an error is encountered")
+        self._stoponerror.setToolTip("Stop process whenever ")
+        self._stoponerror.setCheckable(True)
+        self._stoponerror.setChecked(True)
         self.threadpool = QThreadPool()
 
         lyt = QGridLayout()
         lyt.addWidget(QLabel("Start frame:"), 0, 0 )
-        lyt.addWidget(self._startframe_spinner, 0, 1 )
-        lyt.addWidget(QLabel("assigned"), 1, 0)
-        lyt.addWidget(self._assignedlabel, 1, 1)
-        lyt.addWidget(QLabel("errors/issues"), 2, 0)
-        lyt.addWidget(self._errorlabel, 2, 1)
-        
-        lyt.addWidget(self._startbtn, 3, 0)
-        lyt.addWidget(self._cancelbtn, 3, 1)
-        lyt.addWidget(self._progressbar, 4, 0, 1, 2)
-        lyt.addWidget(self._apply_btn, 5, 0, 1, 2)
+        lyt.addWidget(self._startframe_spinner, 0, 1, 1, 2)
+        lyt.addWidget(QLabel("of"), 1, 1, 1, 1)
+        lyt.addWidget(self._maxframeslabel, 1, 2, 1, 1)
+        lyt.addWidget(self._stoponerror, 2, 0, 1, 3)
+        lyt.addWidget(QLabel("assigned"), 3, 0)
+        lyt.addWidget(self._assignedlabel, 3, 1)
+        lyt.addWidget(QLabel("errors/issues"), 4, 0)
+        lyt.addWidget(self._errorlabel, 4, 1)
+
+        lyt.addWidget(self._startbtn, 5, 0)
+        lyt.addWidget(self._stopbtn, 5, 1)
+        lyt.addWidget(self._proceedbtn, 5, 2)
+        lyt.addWidget(self._apply_btn, 6, 0, 1, 2)
+        lyt.addWidget(self._refreshbtn, 6, 2, 1, 1)
+        lyt.addWidget(self._progressbar, 7, 0, 1, 3)
         self.setLayout(lyt)
 
     def setData(self, data:TrackingData):
@@ -344,19 +373,23 @@ class ConsistencyClassifier(QWidget):
         data : Trackingdata
             The tracking data.
         """
-
-        self._all_cogs = data.centerOfGravity()
+        self._data = data
+        self._all_pos = data.centerOfGravity()
         self._all_orientations = data.orientation()
         self._all_lengths = data.animalLength()
         self._all_bendedness = data.bendedness()
         self._all_scores = data["confidence"]  # ignore for now, let's see how far this carries.
         self._frames = data["frame"]
         self._tracks = data["track"]
+        self._maxframes = np.max(self._frames)
         min_frame = max([self._frames[self._tracks == 1][0], self._frames[self._tracks == 2][0]]) + 1
+        self._maxframeslabel.setText(str(self._maxframes))
         self._startframe_spinner.setMinimum(min_frame)
         self._startframe_spinner.setMaximum(self._frames[-1])
         self._startframe_spinner.setValue(self._frames[0] + 1)
         self._startbtn.setEnabled(True)
+        self._assignedlabel.setText("0")
+        self._errorlabel.setText("0")
         self._worker = None
 
     @Slot(float)
@@ -364,30 +397,44 @@ class ConsistencyClassifier(QWidget):
         if self._progressbar is not None:
             self._progressDialog.setValue(int(value * 100))
 
-    def cancel(self):
+    def stop(self):
         if self._worker is not None:
-            self._worker.cancel()
+            self._worker.stop()
         self._startbtn.setEnabled(True)
-        self._cancelbtn.setEnabled(False)
+        self._proceedbtn.setEnabled(True)
+        self._stopbtn.setEnabled(False)
+        self._refreshbtn.setEnabled(True)
 
-    def run(self):
+    def start(self):
         self._startbtn.setEnabled(False)
-        self._cancelbtn.setEnabled(True)
-        self._worker = ConsistencyWorker(self._all_cogs, self._all_orientations, self._all_lengths,
-                                         self._all_bendedness, self._frames, self._tracks, self._startframe_spinner.value())
-        self._worker.signals.finished.connect(self.worker_done)
+        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._startframe_spinner.value(), self._stoponerror.isChecked())
+        self._worker.signals.stopped.connect(self.worker_stopped)
         self._worker.signals.progress.connect(self.worker_progress)
         self.threadpool.start(self._worker)
 
+    def proceed(self):
+        self.start()
+
+    def refresh(self):
+        self.setData(self._data)
+
     def worker_progress(self, progress, processed, errors):
         self._progressbar.setValue(progress)
         self._errorlabel.setText(str(errors))
         self._assignedlabel.setText(str(processed))
 
-    def worker_done(self):
+    def worker_stopped(self, frame):
         self._apply_btn.setEnabled(True)
         self._startbtn.setEnabled(True)
-        self._cancelbtn.setEnabled(False)
+        self._stopbtn.setEnabled(False)
+        self._startframe_spinner.setValue(frame-1)
+        self._proceedbtn.setEnabled(bool(frame < self._maxframes-1))
+        self._refreshbtn.setEnabled(True)
+        self._processed_frames = frame
 
     def assignedTracks(self):
         return self._tracks
@@ -441,7 +488,7 @@ def main():
     from IPython import embed
     from fixtracks.info import PACKAGE_ROOT
 
-    datafile = PACKAGE_ROOT / "data/merged_small_tracked.pkl"
+    datafile = PACKAGE_ROOT / "data/merged_small.pkl"
 
     with open(datafile, "rb") as f:
         df = pickle.load(f)

fixtracks/widgets/detectiontimeline.py

@@ -94,7 +94,7 @@ class DetectionTimeline(QWidget):
         self._rangeStart = 0.0
         self._rangeStop = 0.005
         self._total_width = 2000
-        self._stepCount = 200
+        self._stepCount = 300
         self._bg_brush = QBrush(QColor(20, 20, 20, 255))
         transparent_brush = QBrush(QColor(200, 200, 200, 64))
         self._white_pen = QPen(QColor.fromString("white"))
@@ -114,7 +114,7 @@ class DetectionTimeline(QWidget):
         self._window = Window(0, 0, 100, 60, axis_pen, transparent_brush)
         self._window.signals.windowMoved.connect(self.on_windowMoved)
 
-        self._scene = QGraphicsScene(QRectF(0, 0, self._total_width, 55.))
+        self._scene = QGraphicsScene(QRectF(0, 0, self._total_width, 65.))
         self._scene.setBackgroundBrush(self._bg_brush)
         self._scene.addItem(self._window)
 
@@ -127,15 +127,15 @@ 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)
 
         self._position_label = QLabel("")
         f = self._position_label.font()

fixtracks/widgets/skeleton.py

@@ -165,6 +165,8 @@ class SkeletonWidget(QWidget):
                      frames:np.ndarray, tracks: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]

fixtracks/widgets/tracks.py

@@ -95,31 +95,43 @@ class SelectionControls(QWidget):
         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, 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.addWidget(backBtn, 0, 0, 3, 2)
+        grid.addWidget(halfstepBackBtn, 3, 0, 2, 2)
+        grid.addWidget(quarterstepBackBtn, 5, 0, 2, 2)
+        grid.addWidget(fwdBtn, 0, 6, 3, 2)
+        grid.addWidget(halfstepFwdBtn, 3, 6, 2, 2)
+        grid.addWidget(quarterstepFwdBtn, 5, 6, 2, 2)
+
+        grid.addWidget(QLabel("Current window:"), 0, 2, 1, 4)
+        grid.addWidget(QLabel("start:"), 1, 3, 1, 1)
+        grid.addWidget(self.startframe, 1, 4, 1, 2, Qt.AlignmentFlag.AlignRight)
+        grid.addWidget(QLabel("end:"), 2, 3, 1, 1)
+        grid.addWidget(self.endframe, 2, 4, 1, 2, Qt.AlignmentFlag.AlignRight)
+        grid.addWidget(QLabel("Current selection:"), 3, 2, 1, 4)
+        grid.addWidget(QLabel("Track One:"), 4, 3, 1, 2)
+        grid.addWidget(self.tone_selection, 4, 5, 1, 1, Qt.AlignmentFlag.AlignRight)
+        grid.addWidget(QLabel("Track Two:"), 5, 3, 1, 2)
+        grid.addWidget(self.ttwo_selection, 5, 5, 1, 1, Qt.AlignmentFlag.AlignRight)
+        grid.addWidget(QLabel("Unassigned:"), 6, 3, 1, 2)
+        grid.addWidget(self.tother_selection, 6, 5, 1, 1, Qt.AlignmentFlag.AlignRight)
+
+        grid.addWidget(assignOneBtn, 7, 0, 4, 3)
+        grid.addWidget(assignOtherBtn, 7, 3, 4, 2)
+        grid.addWidget(assignTwoBtn, 7, 5, 4, 3)
         grid.setColumnStretch(0, 1)
         grid.setColumnStretch(7, 1)
         self.setLayout(grid)
         self.setMaximumSize(QSize(400, 200))
 
+    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:
@@ -150,12 +162,17 @@ class SelectionControls(QWidget):
 
     def setSelectedTracks(self, tracks):
         logging.debug("SelectionControl: setSelectedTracks")
-        tone = np.sum(tracks == 1)
-        ttwo = np.sum(tracks == 2)
+        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))
-        self._total = len(tracks)
+        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
 
 
 class FixTracks(QWidget):
@@ -313,6 +330,7 @@ class FixTracks(QWidget):
 
         start_frame = self._currentWindowPos
         stop_frame = start_frame + self._currentWindowWidth
+        self._controls_widget.setWindow(start_frame, stop_frame)
         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")
@@ -448,6 +466,7 @@ class FixTracks(QWidget):
         self._currentWindowWidth = value
         logging.debug("Tracks:OnWindowSizeChanged %i franes", value)
         self._timeline.setWindowWidth(self._currentWindowWidth / self._maxframes)
+        self._controls_widget.setSelectedTracks(None)
 
     def on_detectionsSelected(self, detections):
         logging.debug("Tracks: Detections selected")
@@ -476,6 +495,7 @@ class FixTracks(QWidget):
         new_start_frame = self._currentWindowPos + step
         self._timeline.setWindowPos(new_start_frame / self._maxframes)
         self._currentWindowPos = new_start_frame
+        self._controls_widget.setSelectedTracks(None)
         self.update()
 
     def on_forward(self, stepsize):