initial commit
This commit is contained in:
unknown
@@ -0,0 +1,235 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
from .nHHD 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'):
|
||||
# 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.__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)
|
||||
|
||||
Reference in New Issue
Block a user