[classifier] improvements

fixtracks/widgets/classifier.py

@@ -10,6 +10,8 @@ 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 WorkerSignals(QObject):
     error = Signal(str)
     running = Signal(bool)
@@ -64,6 +66,35 @@ class ConsistencyWorker(QRunnable):
 
     @Slot()
     def run(self):
+        def needs_checking(original, new):
+            res = False
+            for n, o in zip(new, original):
+                res = (o == 1 or o == 2) and n != o
+            if not res:
+                res = len(new) > 1 and (np.all(new == 1) or np.all(new == 2))
+            return res
+
+        def assign_by_distance(f, p):
+            t1_step = f - last_frame[0]
+            t2_step = f - last_frame[1]
+            if t1_step == 0 or t2_step == 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])/t1_step
+            distance_to_tracktwo = np.linalg.norm(p - last_pos[1])/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(f, o):
+            t1_step = f - last_frame[0]
+            t2_step = f - last_frame[1]
+            orientationchange = np.unwrap((last_angle - o)/np.array([t1_step, t2_step]))
+            most_likely_track = np.argmin(orientationchange) + 1
+            return most_likely_track, orientationchange
+
         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],
@@ -79,53 +110,37 @@ class ConsistencyWorker(QRunnable):
         startframe = np.max(last_frame)
         steps = int((maxframes - startframe) // 200)
 
-        for f in range(startframe + 1, maxframes, 1):
+        for f in np.unique(self.frames[self.frames > startframe]):
             if self._stoprequest:
                 break
             indices = np.where(self.frames == f)[0]
             pp  = self.positions[indices]
             originaltracks = self.tracks[indices]
-            assignments = np.zeros_like(originaltracks)
+            dist_assignments = np.zeros_like(originaltracks)
-            distances = np.zeros((len(originaltracks), 2))
+            angle_assignments = np.zeros_like(originaltracks)
+            # userlabeld = np.zeros_like(originaltracks)
             distances = np.zeros((len(originaltracks), 2))
+            orientations = np.zeros((len(originaltracks), 2))
+
             for i, (idx, p) in enumerate(zip(indices, pp)):
                 if self.userlabeled[idx]:
-                    print("userlabeled")
+                    print("user")
                     processed += 1
                     last_pos[originaltracks[i]-1] = pp[i]
                     last_frame[originaltracks[i]-1] = f
                     last_angle[originaltracks[i]-1] = self.orientations[idx]
                     continue
-                if f < last_frame[0]:
+                dist_assignments[i], distances[i, :] = assign_by_distance(f, p)
-                    print("ping")
+                angle_assignments[i], orientations[i,:] = assign_by_orientation(f, self.orientations[idx])
-                    self.tracks[idx] = 2
+
-                    last_frame[1] = f
-                    last_pos[1] = p
-                    # last_angle[1] = self.orientations[idx]
-                    continue
-                if f < last_frame[1]:
-                    print("pang")
-                    last_frame[0] = f
-                    last_pos[0] = p
-                    # 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
-                distances[i, 0] = distance_to_trackone
-                distances[i, 1] = distance_to_tracktwo
-                assignments[i] = most_likely_track
             # check (re) assignment update and proceed
-            if len(assignments) > 1 and (np.all(assignments == 1) or np.all(assignments == 2)):
+            print("dist", distances)
-                logging.warning("frame %i: Issues assigning based on distances %s", f, str(distances))
+            print("angle", orientations)
+            if needs_checking(originaltracks, dist_assignments):
+                logging.info("frame %i: Issues assigning based on distances %s", f, str(distances))
                 assignment_error = True
                 errors += 1
                 if self._stoponerror:
-                    from IPython import embed
                     embed()
                     break
             else:
@@ -134,10 +149,10 @@ class ConsistencyWorker(QRunnable):
                 if assignment_error:
                     self.tracks[idx] = -1
                 else:
-                    self.tracks[idx] = assignments[i]
+                    self.tracks[idx] = dist_assignments[i]
-                    last_pos[assignments[i]-1] = pp[i]
+                    last_pos[dist_assignments[i]-1] = pp[i]
-                    last_frame[assignments[i]-1] = f
+                    last_frame[dist_assignments[i]-1] = f
-                    last_angle[assignments[i]-1] = self.orientations[idx]
+                    last_angle[dist_assignments[i]-1] = self.orientations[idx]
             assignment_error = False
             if steps > 0 and f % steps == 0:
                 progress += 1