import cv2 import numpy as np # from .HHD import * # import src.utils as utils class Tracker: def __init__(self, adaptive, cuda, 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 if not cuda: # 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(1) # Skip finest pyramid level (2x faster) # self.__tracking_inst.setPatchSize(8) # Smaller patches (default 8) # self.__tracking_inst.setPatchStride(4) # Larger stride = fewer patches # self.__tracking_inst.setGradientDescentIterations(12) # Reduce iterations (default 25) # self.__tracking_inst.setVariationalRefinementIterations(3) # Reduce refinement (default 5) 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 = lower_threshold # error lower than threshold will close the loop self.__upper_threshold = 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) else: 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] 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: print("loop closed at frame ", str(i)) 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.__vpi_prev_frame = None self.__vpi_prev_pyramid = None def getFrameQueueLength(self): return len(self.__fine_base_frame_queue)