ykcheungab
389 lines
17 KiB
Python
389 lines
17 KiB
Python
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)
|
|
# self.__tracking_inst.setUseSpatialPropagation(False)
|
|
# self.__tracking_inst.setUseMeanNormalization(False)
|
|
# self.__tracking_inst.setGradientDescentIterations(16) #16
|
|
# self.__tracking_inst.setPatchSize(12) # 8
|
|
# self.__tracking_inst.setPatchStride(5) # 3
|
|
# self.__tracking_inst.setVariationalRefinementIterations(5)#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 = 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 _prepare_valid_mask(self, valid_mask, shape, erode_iters=1):
|
|
if valid_mask is None:
|
|
return None
|
|
mask = np.asarray(valid_mask, dtype=np.uint8)
|
|
if mask.shape != shape:
|
|
mask = cv2.resize(mask, (shape[1], shape[0]), interpolation=cv2.INTER_NEAREST)
|
|
if erode_iters > 0:
|
|
kernel = np.ones((3, 3), dtype=np.uint8)
|
|
mask = cv2.erode(mask, kernel, iterations=erode_iters)
|
|
return mask.astype(bool)
|
|
|
|
def apply_valid_mask(self, flow, valid_mask, erode_iters=1):
|
|
"""Zero flow outside the valid UV domain (in place when safe).
|
|
|
|
Multiplying by a 0/1 mask is a single streaming pass, whereas
|
|
``flow.copy()`` followed by a boolean fancy-index assignment costs a
|
|
full copy plus a gather/scatter over the masked-out pixels.
|
|
|
|
Note: a masked-out *negative* component becomes ``-0.0`` rather than
|
|
``+0.0``. The two compare, sum and magnitude identically; only the raw
|
|
bit pattern of those zeros differs from the previous implementation.
|
|
"""
|
|
if flow is None or valid_mask is None:
|
|
return flow
|
|
mask = self._prepare_valid_mask(valid_mask, flow.shape[:2], erode_iters=erode_iters)
|
|
if not getattr(self, "_Tracker__adaptive", False):
|
|
# The non-adaptive path hands back the DIS state buffer, which feeds
|
|
# the next frame as its initial guess; never mutate it.
|
|
flow_out = flow.copy()
|
|
flow_out[~mask] = 0
|
|
return flow_out
|
|
if not flow.flags.writeable:
|
|
flow = flow.copy()
|
|
mask_f = mask.astype(np.float32, copy=False)[:, :, None]
|
|
np.multiply(flow, mask_f, out=flow)
|
|
return flow
|
|
|
|
def track(self, frame, valid_mask=None):
|
|
"""Main tracking function that dispatches to CUDA or CPU implementation"""
|
|
if self.__algorithm == "raft":
|
|
flow = self.track_raft(frame.copy())
|
|
elif self.__algorithm == "dis":
|
|
if frame.ndim == 3:
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
|
flow = self.track_cpu(frame.copy(), valid_mask=valid_mask)
|
|
else:
|
|
flow = None
|
|
return self.apply_valid_mask(flow, valid_mask)
|
|
|
|
def track_cpu(self, frame, valid_mask=None):
|
|
"""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, valid_mask=valid_mask)
|
|
|
|
# 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, valid_mask=valid_mask)
|
|
|
|
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, valid_mask=None):
|
|
# 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
|
|
diff_abs = np.abs(diff)
|
|
mask = self._prepare_valid_mask(valid_mask, diff_abs.shape, erode_iters=0)
|
|
if mask is not None and np.any(mask):
|
|
rev = float(diff_abs[mask].mean())
|
|
else:
|
|
rev = float(np.mean(diff_abs))
|
|
return rev
|
|
|
|
def __checkLoopClosure(self, frame, valid_mask=None): # 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, valid_mask=valid_mask)
|
|
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 |