import cv2 import numpy as np import time class Tracker: def __init__(self, adaptive, cuda, config_tracker=None, lower_threshold=12.0, upper_threshold=18.0, loop_closure_interval=1, dis_preset='medium', verbose=False): # Check if OpenCV is built with CUDA support # if cuda and not cv2.cuda.getCudaEnabledDeviceCount(): # # print("Warning: CUDA requested but OpenCV not built with CUDA support or no CUDA device found. Falling back to CPU.") # cuda = False self.cuda = cuda self.verbose = verbose self.__loop_closure_interval = loop_closure_interval self.__frame_count = 0 dis_preset = config_tracker["dis_preset"] loop_closure_interval = config_tracker["loop_closure_interval"] self.__algorithm = config_tracker["algorithm"] # Conditionally import torch only if RAFT algorithm is used self.torch = None if self.__algorithm == "raft": import torch self.torch = torch if self.__algorithm == "raft" and (not self.cuda or not self.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 elif dis_preset == 'fast': preset = cv2.DISOPTICAL_FLOW_PRESET_FAST else: preset = cv2.DISOPTICAL_FLOW_PRESET_MEDIUM self.__tracking_inst = cv2.DISOpticalFlow_create(preset) # Aggressive optimization for CPU self.__tracking_inst.setFinestScale(config_tracker["dis_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() 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.__algorithm == "raft": return self.track_raft(frame) elif self.__algorithm == "dis": if frame.ndim == 3: frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) return self.track_cpu(frame) def track_cpu(self, frame): """CPU-based optical flow tracking""" frame_height, frame_width = frame.shape[:2] 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)) # Pre-allocate with correct dtype self.__fine_flow_queue.append( np.zeros((frame_height, frame_width, 2), dtype=np.float32)) self.__prev_fine_flow = np.zeros( (frame_height, frame_width, 2), dtype=np.float32) 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) # CPU optical flow computation self.__prev_fine_flow = self.__tracking_inst.calc( self.__fine_base_frame_queue[-1], frame, self.__prev_fine_flow) # Direct numpy addition is faster than cv2.add final_flow = self.__fine_flow_queue[-1] + self.__prev_fine_flow # 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 = np.zeros( (frame_height, frame_width, 2), dtype=np.float32) 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) return self.__prev_fine_flow # def track_2dto3d(self, flow): # decompose_obj = nHHD(grid=(flow.shape[0], flow.shape[1]), spacings=(18, 18)) # flag,output = decompose_obj.get_force(flow, verbose=0) # 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] # 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)) return rev 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_flow_queue = self.__fine_flow_queue[:i + 1] return True return False def reset(self): 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) ### RAFT ### def setup_simple_raft(self): """ Setup a simpler version using available models with performance optimizations """ import torchvision.transforms as T from torchvision.models.optical_flow import raft_large # Load model and move to GPU model = raft_large(pretrained=True, progress=False).to(self.__raft_device) model = model.eval() if self.verbose: print(f"RAFT model loaded on {self.__raft_device}") return model 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']: frame = np.ascontiguousarray(frame) # Upload to device as uint8 first to save bandwidth (4x less data than float32) tensor = self.torch.from_numpy(frame).permute(2, 0, 1) # 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) if not tensor.is_pinned(): tensor = tensor.pin_memory() tensor = tensor.to(self.__raft_device, non_blocking=True) else: tensor = tensor.to(self.__raft_device) # Convert to float and normalize on device tensor = tensor.float() / 255.0 # Normalize to [-1, 1] tensor = (tensor - 0.5) * 2.0 # Add batch dimension tensor = tensor.unsqueeze(0) return tensor 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] == 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 self.__raft_ref_tensor = self._preprocess_frame_to_tensor(frame) prev_tensor = self.__raft_ref_tensor curr_tensor = self._preprocess_frame_to_tensor(frame) if self.verbose: t0 = time.perf_counter() with self.torch.inference_mode(): if self.__raft_device == 'cuda': with self.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] 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 # 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