integrated finger process pipeline on remapped 2d image

src/orisys/tracker.py

@@ -82,16 +82,39 @@ class Tracker:
 
 
 
-    def track(self, frame):
+    def _prepare_valid_mask(self, valid_mask, shape, erode_iters=1):
+        if valid_mask is None:
+            return None
+        mask = np.asarray(valid_mask, dtype=np.uint8)
+        if mask.shape != shape:
+            mask = cv2.resize(mask, (shape[1], shape[0]), interpolation=cv2.INTER_NEAREST)
+        if erode_iters > 0:
+            kernel = np.ones((3, 3), dtype=np.uint8)
+            mask = cv2.erode(mask, kernel, iterations=erode_iters)
+        return mask.astype(bool)
+
+    def apply_valid_mask(self, flow, valid_mask, erode_iters=1):
+        """Zero flow outside the valid UV domain."""
+        if flow is None or valid_mask is None:
+            return flow
+        mask = self._prepare_valid_mask(valid_mask, flow.shape[:2], erode_iters=erode_iters)
+        flow_out = flow.copy()
+        flow_out[~mask] = 0
+        return flow_out
+
+    def track(self, frame, valid_mask=None):
         """Main tracking function that dispatches to CUDA or CPU implementation"""
         if self.__algorithm == "raft":
-            return self.track_raft(frame.copy())
+            flow = self.track_raft(frame.copy())
         elif self.__algorithm == "dis":
             if frame.ndim == 3:
                 frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
-            return self.track_cpu(frame.copy())
+            flow = self.track_cpu(frame.copy(), valid_mask=valid_mask)
+        else:
+            flow = None
+        return self.apply_valid_mask(flow, valid_mask)
     
-    def track_cpu(self, frame):
+    def track_cpu(self, frame, valid_mask=None):
         """CPU-based optical flow tracking"""
         frame_height, frame_width = frame.shape[:2]
         
@@ -114,7 +137,7 @@ class Tracker:
                 
                 # Only check loop closure periodically (expensive operation)
                 if self.__frame_count % self.__loop_closure_interval == 0:
-                    self.__checkLoopClosure(frame)
+                    self.__checkLoopClosure(frame, valid_mask=valid_mask)
                 
                 # CPU optical flow computation
                 self.__prev_fine_flow = self.__tracking_inst.calc(
@@ -126,7 +149,7 @@ class Tracker:
                 # Use fine flow for error estimation
                 error = self.__getForwardBackwardError(
                     self.__fine_base_frame_queue[-1], frame,
-                    self.__prev_fine_flow)
+                    self.__prev_fine_flow, valid_mask=valid_mask)
                 
                 if error > self.__upper_threshold:
                     self.__fine_base_frame_queue.append(frame)
@@ -148,7 +171,7 @@ class Tracker:
     #     flag,output = decompose_obj.get_force(flow, verbose=0)
     #     return output
         
-    def __getForwardBackwardError(self, I0, I1, flow):
+    def __getForwardBackwardError(self, I0, I1, flow, valid_mask=None):
         # if self.cuda:
         #     flow_map = cv2.cuda.add(self.map, flow)
         #     map_x, map_y = cv2.cuda.split(flow_map)
@@ -175,15 +198,20 @@ class Tracker:
             diff = z.astype('float32') - I0.astype('float32')
         else:
             diff = z - I0
-        rev = np.mean(np.abs(diff))
+        diff_abs = np.abs(diff)
+        mask = self._prepare_valid_mask(valid_mask, diff_abs.shape, erode_iters=0)
+        if mask is not None and np.any(mask):
+            rev = float(diff_abs[mask].mean())
+        else:
+            rev = float(np.mean(diff_abs))
         return rev
 
-    def __checkLoopClosure(self, frame):  # update (prune) the frame queue with loop closure detection
+    def __checkLoopClosure(self, frame, valid_mask=None):  # update (prune) the frame queue with loop closure detection
         for i in range(len(self.__fine_base_frame_queue) - 2, -1, -1):
             f = self.__tracking_inst.calc(
                 self.__fine_base_frame_queue[i], frame, None)
             error = self.__getForwardBackwardError(
-                self.__fine_base_frame_queue[i], frame, f)
+                self.__fine_base_frame_queue[i], frame, f, valid_mask=valid_mask)
             if error < self.__lower_threshold:
                 self.__fine_base_frame_queue = self.__fine_base_frame_queue[:i+1]
                 self.__fine_flow_queue = self.__fine_flow_queue[:i + 1]