diff --git a/RAFT_test.py b/RAFT_test.py index 7b454fb..416a4cf 100644 --- a/RAFT_test.py +++ b/RAFT_test.py @@ -36,7 +36,7 @@ def setup_simple_raft(): """ Setup a simpler version using available models """ - model = raft_large(pretrained=True, progress=False).to('cuda') + model = raft_small(pretrained=True, progress=False).to('cuda') model = model.eval() # Alternative: Use LiteFlowNet (lighter alternative) print("RAFT loaded successfully") @@ -61,7 +61,6 @@ def track_raft(prev_frame, frame, model, device): # Transpose to (H, W, 2) format expected by draw_arrows flow_np = flow_np.transpose(1, 2, 0) # (2, H, W) -> (H, W, 2) - # torch.onnx.export( # model, # model to export # (prev_tensor, curr_tensor), # inputs of the model, @@ -96,10 +95,11 @@ def main_simple(): dy = 0 frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame1 = frame[420:660, 840:1080, :] - frame2 = frame[420+dy:660+dy, 840+dx:1080+dx, :] - frame3 = frame[420-dy:660-dy, 840-dx:1080-dx, :] + # frame2 = frame[420+dy:660+dy, 840+dx:1080+dx, :] + # frame3 = frame[420-dy:660-dy, 840-dx:1080-dx, :] - frame = np.concatenate([frame3, frame1, frame2], axis=1) + # frame = np.concatenate([frame3, frame1, frame2], axis=1) + frame = frame1 # cv2.imshow("frame", frame3) # cv2.waitKey(0) if prev_gray is None: @@ -137,7 +137,7 @@ def main_simple(): arrow_scale=1.0 ) - arrows = cv2.resize(arrows, (1200, 400)) + arrows = cv2.resize(arrows, (400, 400)) cv2.imshow('Optical Flow', arrows) # prev_gray = gray.copy() diff --git a/config/prototype_of3.json b/config/prototype_of3.json index 9b2dd1d..57b6e23 100644 --- a/config/prototype_of3.json +++ b/config/prototype_of3.json @@ -2,11 +2,11 @@ "name": "prototype", "version": "1.0", "camera": { - "resolution_raw": [960, 540], + "resolution_raw": [1920, 1080], "fps": 30, "units": 15, "units_used": 9, - "resolution_stitch": [200, 200] + "resolution_stitch": [400, 400] }, "stitcher": { "min_area_frac": 0.01, @@ -24,8 +24,10 @@ "lower_threshold": 12.0, "upper_threshold": 18.0, "loop_closure_interval": 1, + "algorithm": "raft", "dis_preset": "medium", - "of_mode": 3 + "of_mode": 3, + "algorithm_list": ["raft", "dis"] }, "contact": { "min_contour_area": 3000, diff --git a/src/orisys/tracker.py b/src/orisys/tracker.py index f984a97..b298376 100644 --- a/src/orisys/tracker.py +++ b/src/orisys/tracker.py @@ -26,7 +26,23 @@ class Tracker: self.__frame_count = 0 dis_preset = config_tracker["dis_preset"] loop_closure_interval = config_tracker["loop_closure_interval"] - if not self.cuda: + self.__algorithm = config_tracker["algorithm"] + + if self.__algorithm == "raft" and (not self.cuda or not torch.cuda.is_available()): + self.__algorithm = "dis" + if self.verbose: + print("No CUDA available, falling back to OPENCV DIS optical flow") + + + if self.__algorithm == "raft": + ### RAFT + self.__raft_device = "cuda" + print(f"RAFT device: {self.__raft_device}") + self.__raft_model = self.setup_simple_raft() + self.__raft_ref_frame = None + self.__raft_ref_tensor = None + + elif self.__algorithm == "dis": # DIS method - optimize for speed if dis_preset == 'ultrafast': preset = cv2.DISOPTICAL_FLOW_PRESET_ULTRAFAST @@ -35,98 +51,35 @@ class Tracker: else: preset = cv2.DISOPTICAL_FLOW_PRESET_MEDIUM - self.__tracking_inst = cv2.DISOpticalFlow_create(preset) - + self.__tracking_inst = cv2.DISOpticalFlow_create(preset) # Aggressive optimization for CPU self.__tracking_inst.setFinestScale(config_tracker["of_mode"]) # Skip finest pyramid level (2x faster) - # else: - # # CUDA flow method with stream support - # self.__tracking_inst = cv2.cuda.FarnebackOpticalFlow_create() - # # Create CUDA stream for async operations - # self.__cuda_stream = cv2.cuda.Stream() + # else: + # # CUDA flow method with stream support + # self.__tracking_inst = cv2.cuda.FarnebackOpticalFlow_create() + # # Create CUDA stream for async operations + # self.__cuda_stream = cv2.cuda.Stream() - self.__fine_base_frame_queue = [] - self.__fine_flow_queue = [] - self.__prev_fine_flow = None - self.__adaptive = adaptive - 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_device = "cuda" if torch.cuda.is_available() else "cpu" - print(f"RAFT device: {self.__raft_device}") - self.__raft_model = self.setup_simple_raft() - self.__raft_ref_frame = None - self.__raft_ref_tensor = None + self.__fine_base_frame_queue = [] + self.__fine_flow_queue = [] + self.__prev_fine_flow = None + self.__adaptive = adaptive + 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 + def track(self, frame): """Main tracking function that dispatches to CUDA or CPU implementation""" - if self.cuda: - # return self.track_cuda(frame) + if self.__algorithm == "raft": return self.track_raft(frame) - else: + elif self.__algorithm == "dis": if frame.ndim == 3: frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) return self.track_cpu(frame) - def track_cuda(self, frame): - """CUDA-accelerated optical flow tracking""" - frame_height, frame_width = frame.shape[:2] - frame = cv2.cuda_GpuMat(frame) - - if self.__adaptive: - if not self.__fine_base_frame_queue: - self.__fine_base_frame_queue.append(frame) - # Pre-compute meshgrid for remap operations - x = np.linspace(0, frame_width - 1, frame_width, dtype='float32') - y = np.linspace(0, frame_height - 1, frame_height, dtype='float32') - self.map = np.array(np.meshgrid(x, y)).transpose((1, 2, 0)) - - cuMat = cv2.cuda_GpuMat() - cuMat.upload(self.map) - self.map = cuMat - self.__fine_flow_queue.append( - cv2.cuda_GpuMat(frame.size(), cv2.CV_32FC2, (0.0, 0.0))) - self.__prev_fine_flow = cv2.cuda_GpuMat( - frame.size(), cv2.CV_32FC2, (0.0, 0.0)) - return self.__fine_flow_queue[-1] - else: - self.__frame_count += 1 - - # Only check loop closure periodically (expensive operation) - if self.__frame_count % self.__loop_closure_interval == 0: - self.__checkLoopClosure(frame) - - # Use CUDA stream for async optical flow computation - self.__prev_fine_flow = self.__tracking_inst.calc( - self.__fine_base_frame_queue[-1], frame, - self.__prev_fine_flow, self.__cuda_stream) - - final_flow = cv2.cuda.add(self.__fine_flow_queue[-1], - self.__prev_fine_flow, - stream=self.__cuda_stream) - - # Use fine flow for error estimation - error = self.__getForwardBackwardError( - self.__fine_base_frame_queue[-1], frame, - self.__prev_fine_flow) - - if error > self.__upper_threshold: - self.__fine_base_frame_queue.append(frame) - self.__fine_flow_queue.append(final_flow) - self.__prev_fine_flow = None - return final_flow - else: - if not self.__fine_base_frame_queue: - self.__fine_base_frame_queue.append(frame) - - self.__prev_fine_flow = self.__tracking_inst.calc( - self.__fine_base_frame_queue[-1], frame, - self.__prev_fine_flow, self.__cuda_stream) - return self.__prev_fine_flow - def track_cpu(self, frame): """CPU-based optical flow tracking""" frame_height, frame_width = frame.shape[:2] @@ -185,58 +138,57 @@ class Tracker: # return output def __getForwardBackwardError(self, I0, I1, flow): - if self.cuda: - flow_map = cv2.cuda.add(self.map, flow) - map_x, map_y = cv2.cuda.split(flow_map) - z = cv2.cuda.remap(I1, - map_x, - map_y, - interpolation=cv2.INTER_CUBIC, - borderMode=cv2.BORDER_REPLICATE) - t = I0.clone() - t = t.convertTo(rtype=cv2.CV_32FC1, dst=t) - z = z.convertTo(rtype=cv2.CV_32FC1, dst=z) - rev = cv2.cuda.absSum(cv2.cuda.subtract( - t, z))[0] / t.size()[0] / t.size()[1] + # if self.cuda: + # flow_map = cv2.cuda.add(self.map, flow) + # map_x, map_y = cv2.cuda.split(flow_map) + # z = cv2.cuda.remap(I1, + # map_x, + # map_y, + # interpolation=cv2.INTER_CUBIC, + # borderMode=cv2.BORDER_REPLICATE) + # t = I0.clone() + # t = t.convertTo(rtype=cv2.CV_32FC1, dst=t) + # z = z.convertTo(rtype=cv2.CV_32FC1, dst=z) + # rev = cv2.cuda.absSum(cv2.cuda.subtract( + # t, z))[0] / t.size()[0] / t.size()[1] + # else: + # Optimized CPU path - avoid redundant type conversions + flow_map = self.map + flow + z = cv2.remap(I1, + flow_map[:, :, 0], + flow_map[:, :, 1], + interpolation=cv2.INTER_CUBIC, # INTER_LINEAR faster than INTER_CUBIC + borderMode=cv2.BORDER_REPLICATE) + # Use in-place operations when possible + if I0.dtype != np.float32: + diff = z.astype('float32') - I0.astype('float32') else: - # Optimized CPU path - avoid redundant type conversions - flow_map = self.map + flow - z = cv2.remap(I1, - flow_map[:, :, 0], - flow_map[:, :, 1], - interpolation=cv2.INTER_CUBIC, # INTER_LINEAR faster than INTER_CUBIC - borderMode=cv2.BORDER_REPLICATE) - # Use in-place operations when possible - if I0.dtype != np.float32: - diff = z.astype('float32') - I0.astype('float32') - else: - diff = z - I0 - rev = np.mean(np.abs(diff)) + diff = z - I0 + rev = np.mean(np.abs(diff)) return rev - def __checkLoopClosure( - self, frame - ): # update (prune) the frame queue with loop closure detection + def __checkLoopClosure(self, frame): # 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) if error < self.__lower_threshold: - self.__fine_base_frame_queue = self.__fine_base_frame_queue[: - i + - 1] + self.__fine_base_frame_queue = self.__fine_base_frame_queue[:i+1] self.__fine_flow_queue = self.__fine_flow_queue[:i + 1] return True return False def reset(self): - self.__fine_base_frame_queue = [] - self.__fine_flow_queue = [] - self.__prev_fine_flow = None - self.__frame_count = 0 - self.__raft_ref_frame = None - self.__raft_ref_tensor = None + if self.__algorithm == "raft": + self.__raft_ref_frame = None + self.__raft_ref_tensor = None + + elif self.__algorithm == "dis": + self.__fine_base_frame_queue = [] + self.__fine_flow_queue = [] + self.__prev_fine_flow = None + self.__frame_count = 0 def getFrameQueueLength(self): return len(self.__fine_base_frame_queue) @@ -251,22 +203,11 @@ class Tracker: model = raft_large(pretrained=True, progress=False).to(self.__raft_device) model = model.eval() - # Optimize matmul precision for faster inference (TF32 on Ampere+ GPUs) - # torch.set_float32_matmul_precision('medium') - - # Enable CUDA-specific optimizations - if self.__raft_device == 'cuda': - # cuDNN optimizations - torch.backends.cudnn.benchmark = True # Auto-select fastest algorithms - torch.backends.cudnn.deterministic = False # Allow non-deterministic for speed - # TensorFloat-32 for Ampere+ GPUs (faster with minimal accuracy loss) - torch.backends.cuda.matmul.allow_tf32 = True - torch.backends.cudnn.allow_tf32 = True - - print(f"RAFT model loaded on {self.__raft_device}") + if self.verbose: + print(f"RAFT model loaded on {self.__raft_device}") return model - def _preprocess_frame_to_tensor(self, frame, resize=None): + def _preprocess_frame_to_tensor(self, frame): """Efficiently convert numpy frame to RAFT-ready tensor""" # Ensure contiguous memory for faster tensor conversion if not frame.flags['C_CONTIGUOUS']: @@ -275,7 +216,6 @@ class Tracker: # Upload to device as uint8 first to save bandwidth (4x less data than float32) tensor = torch.from_numpy(frame).permute(2, 0, 1) - # Move to device with non_blocking for async transfer (if CUDA) # Pin memory for faster CPU-GPU transfer (similar to DataLoader pin_memory=True) if self.__raft_device == 'cuda': # Pin memory before transferring to GPU (must be on CPU first) @@ -287,10 +227,6 @@ class Tracker: # 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 @@ -300,30 +236,25 @@ class Tracker: return tensor - def track_raft(self, frame): + def track_raft(self, frame, gpu = False): # 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) if self.__raft_ref_frame is None: - 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_frame = frame self.__raft_ref_tensor = self._preprocess_frame_to_tensor(frame) prev_tensor = self.__raft_ref_tensor curr_tensor = self._preprocess_frame_to_tensor(frame) - - # Synchronize before timing to ensure preprocessing is complete - if self.__raft_device == 'cuda': - torch.cuda.synchronize() - + if self.verbose: + t0 = time.perf_counter() + with torch.inference_mode(): if self.__raft_device == 'cuda': with torch.amp.autocast('cuda'): @@ -331,14 +262,70 @@ class Tracker: else: flow = self.__raft_model(prev_tensor, curr_tensor) predicted_flows = flow[-1] - - # Synchronize after inference to get accurate timing - if self.__raft_device == 'cuda': - torch.cuda.synchronize() - + + if self.verbose: + t1 = time.perf_counter() + print(f"RAFT time: {t1 - t0}") # Optimize tensor->numpy conversion: contiguous() ensures optimal memory layout flow_np = predicted_flows[0].contiguous().cpu().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) - return flow_np \ No newline at end of file + return flow_np + + + # def track_cuda(self, frame): + # """CUDA-accelerated optical flow tracking""" + # frame_height, frame_width = frame.shape[:2] + # frame = cv2.cuda_GpuMat(frame) + + # if self.__adaptive: + # if not self.__fine_base_frame_queue: + # self.__fine_base_frame_queue.append(frame) + # # Pre-compute meshgrid for remap operations + # x = np.linspace(0, frame_width - 1, frame_width, dtype='float32') + # y = np.linspace(0, frame_height - 1, frame_height, dtype='float32') + # self.map = np.array(np.meshgrid(x, y)).transpose((1, 2, 0)) + + # cuMat = cv2.cuda_GpuMat() + # cuMat.upload(self.map) + # self.map = cuMat + # self.__fine_flow_queue.append( + # cv2.cuda_GpuMat(frame.size(), cv2.CV_32FC2, (0.0, 0.0))) + # self.__prev_fine_flow = cv2.cuda_GpuMat( + # frame.size(), cv2.CV_32FC2, (0.0, 0.0)) + # return self.__fine_flow_queue[-1] + # else: + # self.__frame_count += 1 + + # # Only check loop closure periodically (expensive operation) + # if self.__frame_count % self.__loop_closure_interval == 0: + # self.__checkLoopClosure(frame) + + # # Use CUDA stream for async optical flow computation + # self.__prev_fine_flow = self.__tracking_inst.calc( + # self.__fine_base_frame_queue[-1], frame, + # self.__prev_fine_flow, self.__cuda_stream) + + # final_flow = cv2.cuda.add(self.__fine_flow_queue[-1], + # self.__prev_fine_flow, + # stream=self.__cuda_stream) + + # # Use fine flow for error estimation + # error = self.__getForwardBackwardError( + # self.__fine_base_frame_queue[-1], frame, + # self.__prev_fine_flow) + + # if error > self.__upper_threshold: + # self.__fine_base_frame_queue.append(frame) + # self.__fine_flow_queue.append(final_flow) + # self.__prev_fine_flow = None + # return final_flow + # else: + # if not self.__fine_base_frame_queue: + # self.__fine_base_frame_queue.append(frame) + + # self.__prev_fine_flow = self.__tracking_inst.calc( + # self.__fine_base_frame_queue[-1], frame, + # self.__prev_fine_flow, self.__cuda_stream) + # return self.__prev_fine_flow \ No newline at end of file