diff --git a/src/orisys/tracker.py b/src/orisys/tracker.py index a799161..d31e117 100644 --- a/src/orisys/tracker.py +++ b/src/orisys/tracker.py @@ -25,7 +25,7 @@ class Tracker: self.__frame_count = 0 dis_preset = config_tracker["dis_preset"] loop_closure_interval = config_tracker["loop_closure_interval"] - if not cuda: + if not self.cuda: # DIS method - optimize for speed if dis_preset == 'ultrafast': preset = cv2.DISOPTICAL_FLOW_PRESET_ULTRAFAST @@ -53,9 +53,10 @@ class Tracker: self.__lower_threshold = config_tracker["lower_threshold"] # error lower than threshold will close the loop self.__upper_threshold = config_tracker["upper_threshold"] # error greater than threshold will add new reference frame ### RAFT - self.__raft_model = self.setup_simple_raft() self.__raft_device = "cuda" if torch.cuda.is_available() else "cpu" + self.__raft_model = self.setup_simple_raft() self.__raft_ref_frame = None + self.__raft_ref_tensor = None def track(self, frame): @@ -232,61 +233,78 @@ class Tracker: self.__fine_flow_queue = [] self.__prev_fine_flow = None self.__frame_count = 0 + self.__raft_ref_frame = None + self.__raft_ref_tensor = None def getFrameQueueLength(self): return len(self.__fine_base_frame_queue) ### RAFT ### - def preprocess(self,batch, resize=None): - """ - Preprocess image batch for RAFT model. - - Args: - batch: Tensor of shape (N, C, H, W) - resize: Optional tuple (height, width) to resize. If None, keeps original size. - - Returns: - Preprocessed tensor - """ - transform_list = [ - T.ConvertImageDtype(torch.float32), - T.Normalize(mean=0.5, std=0.5), # map [0, 1] into [-1, 1] - ] - - if resize is not None: - transform_list.append(T.Resize(size=resize)) - - transforms = T.Compose(transform_list) - batch = transforms(batch) - return batch - def setup_simple_raft(self): """ Setup a simpler version using available models """ - model = raft_large(pretrained=True, progress=False).to('cuda') + # Load model and move to GPU + model = raft_large(pretrained=True, progress=False).to(self.__raft_device) model = model.eval() - # Alternative: Use LiteFlowNet (lighter alternative) - print("RAFT loaded successfully") + + print(f"RAFT model loaded on {self.__raft_device}") return model + def _preprocess_frame_to_tensor(self, frame, resize=None): + """Efficiently convert numpy frame to RAFT-ready tensor""" + # Upload to device as uint8 first to save bandwidth + # permute returns a view, but from_numpy usually shares memory. + # Moving to GPU as uint8 is fast. + tensor = torch.from_numpy(frame).permute(2, 0, 1).to(self.__raft_device) + + # Convert to float and normalize on device + tensor = tensor.float() / 255.0 + + # Resize if needed + if resize is not None: + tensor = T.Resize(resize)(tensor) + + # Normalize to [-1, 1] + tensor = (tensor - 0.5) * 2.0 + + # Add batch dimension + tensor = tensor.unsqueeze(0) + + # Keep in float32 - autocast will handle mixed precision during inference + return tensor + def track_raft(self, frame): + # No need to copy if we just read + # frame is expected to be (H, W, 3) or (H, W) + + if frame.ndim == 2: + frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB) + elif frame.shape[2] == 4: + frame = cv2.cvtColor(frame, cv2.COLOR_BGRA2RGB) + elif frame.shape[2] == 3: + # Assume BGR input from OpenCV + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - gray = frame.copy() if self.__raft_ref_frame is None: - self.__raft_ref_frame = gray.copy() - prev_frame = self.__raft_ref_frame + self.__raft_ref_frame = frame # Keep reference for potential debug/CPU use + # Cache the processed tensor to avoid re-processing reference frame every time + self.__raft_ref_tensor = self._preprocess_frame_to_tensor(frame) - prev_tensor = torch.from_numpy(prev_frame.transpose(2, 0, 1).astype(np.float32) / 255.0) - curr_tensor = torch.from_numpy(gray.transpose(2, 0, 1).astype(np.float32) / 255.0) - prev_tensor = self.preprocess(prev_tensor.unsqueeze(0), resize=None).to(self.__raft_device) - curr_tensor = self.preprocess(curr_tensor.unsqueeze(0), resize=None).to(self.__raft_device) + prev_tensor = self.__raft_ref_tensor + curr_tensor = self._preprocess_frame_to_tensor(frame) + + # Use autocast for automatic mixed precision (FP16 where supported, FP32 where needed) with torch.no_grad(): - flow = self.__raft_model(prev_tensor, curr_tensor) + if self.__raft_device == 'cuda': + with torch.amp.autocast('cuda'): + flow = self.__raft_model(prev_tensor, curr_tensor) + else: + flow = self.__raft_model(prev_tensor, curr_tensor) predicted_flows = flow[-1] - flow_np = predicted_flows[0].cpu().numpy() # or .squeeze(0) + flow_np = predicted_flows[0].float().cpu().numpy() # Convert back to float32 for numpy # Transpose to (H, W, 2) format expected by draw_arrows flow_np = flow_np.transpose(1, 2, 0) # (2, H, W) -> (H, W, 2)