"""Keyframe-based point tracking with proper fresh point handling."""

def __init__(

self, H, W, num_landmarks=2000, dtype=torch.float32, device='cuda', model_execution_lock=None,

is_portrait_mode=False, verbose=True

):

self.H, self.W = H, W

self.verbose = verbose

self.num_landmarks = num_landmarks

self.dtype = dtype

self.device = device

if model_execution_lock is None:

model_execution_lock = nullcontext()

self.model_execution_lock = model_execution_lock

self.EDGE_MARGIN = 40

self.KEYFRAME_INTERVAL = 10

self.is_portrait_mode = is_portrait_mode

if is_portrait_mode:

h, w = W, H

else:

h, w = H, W

self.raft_tracker = RaftTracker(

# We use float32 Because model for some reason, flops on float32

device='dml', dtype=torch.float32, H=h, W=w, num_landmarks=num_landmarks

)

self.current_flows = self.generate_grid_points()

self.current_masks = torch.ones(num_landmarks, dtype=torch.bool, device=device)

self.keyframe_tensor = None

self.keyframe_tensor_cp = torch.empty((H, W, 3), dtype=self.raft_tracker.dtype, device=self.raft_tracker.exec_device)

self.frame_tensor_cp = torch.empty((H, W, 3), dtype=self.raft_tracker.dtype, device=self.raft_tracker.exec_device)

self.keyframe_flows = self.current_flows.clone()

self.fresh_integration_mask_cp = torch.zeros(num_landmarks, dtype=torch.bool, device=device)

self.fresh_integration_mask = torch.zeros(num_landmarks, dtype=torch.bool, device=device)

self.arange_buf = torch.arange(num_landmarks + num_landmarks*5 + 10, device=self.device)

self.pending_integration_command = None

self.candidate_lock = threading.Lock()

self.frame_count = 0

print(f"🎯 Tracker: {num_landmarks} points @ {W}x{H} (Keyframe-based, Fixed)\n")

# Debug tools ===========================================================================

self.alive_mask_ones_debug = torch.ones(num_landmarks, dtype=torch.bool, device=device)

self.alive_mask_zeros_debug = torch.zeros(num_landmarks, dtype=torch.bool, device=device)

self.landmarks_debug = torch.empty((num_landmarks, 2), dtype=dtype, device=device)

self.landmarks_debug[..., 0] = W / 2

self.landmarks_debug[..., 1] = H / 2

def generate_grid_points(self):

"""Generate grid in [X, Y] format."""

aspect = self.W / self.H

if self.num_landmarks == 2:

return torch.tensor([[self.W / 4, self.H / 2],

[3 * self.W / 4, self.H / 2]], dtype=self.dtype, device=self.device)

elif self.num_landmarks == 1:

return torch.tensor([[self.W / 2, self.H / 2]], dtype=self.dtype, device=self.device)

elif self.num_landmarks == 0:

return torch.empty((0, 2), dtype=self.dtype, device=self.device)

ny = int(np.sqrt(self.num_landmarks / aspect))

nx = int(np.ceil(self.num_landmarks / ny))

y = np.linspace(30, self.H - 30, ny)

x = np.linspace(30, self.W - 30, nx)

xv, yv = np.meshgrid(x, y)

points = np.stack([xv.ravel(), yv.ravel()], axis=1).astype(np.float32)[:self.num_landmarks]

if len(points) < self.num_landmarks:

points = np.vstack([points, np.tile(points[-1:], (self.num_landmarks - len(points), 1))])

return torch.from_numpy(points).to(self.device)

def get_and_clear_fresh_mask(self):

"""Get and clear fresh mask atomically."""

fresh = self.fresh_integration_mask_cp.copy_(self.fresh_integration_mask)

self.fresh_integration_mask.zero_()

return fresh

def set_candidate_points(self, points, z_values, amp_values_norm, frame_tensor):

"""Store candidates WITH their pre-sampled z-values."""

with self.candidate_lock:

self.pending_integration_command = {

"candidates_S": points,

"z_values_S": z_values,

"amp_values_norm_S": amp_values_norm,

"keyframe_S": self.frame_tensor_cp.copy_(frame_tensor),

}

def _execute_integration_command(self, current_frame_tensor):

"""Execute integration and return z-values for fresh points."""

command = None

with self.candidate_lock:

if self.pending_integration_command is not None:

command = self.pending_integration_command

self.pending_integration_command = None

if command is None:

return None

candidates_S = command["candidates_S"]

z_values_S = command["z_values_S"]

amp_values_norm_S = command["amp_values_norm_S"]

keyframe_S = command["keyframe_S"]

is_dead_mask = ~self.current_masks

# num_dead = is_dead_mask.sum()

num_candidates = len(candidates_S)

# Non-blocking replacement

# rand_indices = torch.randperm(num_candidates, device=self.device)

# Create a boolean mask of size 'num_candidates' where the first 'num_dead' are True

# We use arange to compare against num_dead purely on GPU

# mask_indices = self.arange_buf[:num_candidates] < num_dead

# Now filter. Note: This creates a dynamic shape tensor, see point 2.

# indices = rand_indices[mask_indices]

indices = self.arange_buf[:num_candidates]

selected_S = candidates_S[indices]

selected_zs_S = z_values_S[indices]

selected_amps_norm_S = amp_values_norm_S[indices]

self.verb_print(f" ... running S→T mini-track for points")

# For some reason, portrait fails so we simply permute to landscape temporarily

if self.is_portrait_mode:

selected_S = selected_S[:, [1, 0]]

keyframe_S = keyframe_S.permute(0, 2, 1)

current_frame_tensor = current_frame_tensor.permute(0, 2, 1)

selected_T, valid_T, _, _ = self.raft_tracker.track_points(

keyframe_S,

current_frame_tensor,

selected_S.to(dtype=torch.float32)

)

if self.is_portrait_mode:

selected_T = selected_T[:, [1, 0]]

dead_indices = torch.where(is_dead_mask)[0]

valid_indices_into_dead = torch.where(valid_T)[0]

valid_new_indices = dead_indices[valid_indices_into_dead]

valid_new_points_T = selected_T[valid_T]

valid_zs_S = selected_zs_S[valid_T]

valid_amps_norm_S = selected_amps_norm_S[valid_T]

final_integration_mask = valid_new_indices

# final_integration_mask = torch.zeros(self.num_landmarks, dtype=torch.bool, device=self.device)

# final_integration_mask[valid_new_indices] = True

self.current_flows[final_integration_mask] = valid_new_points_T

self.current_masks[final_integration_mask] = True

self.fresh_integration_mask[final_integration_mask] = True

self.verb_print(f" ✓ Synced new points to present (T)")

return {

'integrated': True,

'fresh_indices': valid_new_indices,

'fresh_z_values': valid_zs_S,

'normalized_amp_values': valid_amps_norm_S

}

def update(self, frame_tensor):

"""Main tracking update with z-value passthrough."""

# 1. HANDLING INITIALIZATION OR REBASE

# We rebase if it's the first frame, OR if we hit the interval

# You might also want to OR this with: `or self.fresh_integration_occurred`

if self.keyframe_tensor is None or (self.frame_count % self.KEYFRAME_INTERVAL == 0):

if self.keyframe_tensor is None:

self.keyframe_tensor = self.keyframe_tensor_cp.copy_(frame_tensor)

else:

self.keyframe_tensor.copy_(frame_tensor)

self.keyframe_flows[:] = self.current_flows

# 2. TRACKING (Keyframe -> Current)

if self.is_portrait_mode:

keyframe_tensor_ = self.keyframe_tensor.permute(1, 0, 2)

frame_tensor_ = frame_tensor.permute(1, 0, 2)

keyframe_flows_ = self.keyframe_flows[:, [1, 0]]

else:

keyframe_tensor_ = self.keyframe_tensor

frame_tensor_ = frame_tensor

keyframe_flows_ = self.keyframe_flows

new_points, valid, flow_field, sampled_flow = self.raft_tracker.track_points(

# keyframe_flows must be float32 for accurate accumulative precision

keyframe_tensor_[None].permute(0, 3, 1, 2).clone(), frame_tensor_[None].permute(0, 3, 1, 2).clone(), keyframe_flows_

)

if self.is_portrait_mode:

new_points = new_points[:, [1, 0]]

self.current_masks[~valid] = False

self.current_flows[self.current_masks] = new_points[self.current_masks]

# 4. INTEGRATION

fresh_mask = self.get_and_clear_fresh_mask()

integration_result = self._execute_integration_command(frame_tensor)

# OPTIONAL PRO TIP:

# If integration happened, the new points exist in 'frame_tensor' but NOT in 'self.keyframe_tensor'.

# You must force a rebase on the NEXT frame, or immediately update the keyframe now.

# Otherwise, the next loop will try to track these new points from the OLD keyframe.

if integration_result and integration_result['integrated']:

# Force keyframe update immediately so new points stick

self.keyframe_tensor.copy_(frame_tensor)

self.keyframe_flows[:] = self.current_flows

self.verb_print(" 🔄 Forced Rebase due to Integration")

self.frame_count += 1

return self.current_flows, self.current_masks, fresh_mask, integration_result

def verb_print(self, *args, **kwargs):

if self.verbose:

"""

def __init__(

self, H, W, num_landmarks=256, dtype=torch.float32, device='cpu', model_execution_lock=None,

is_portrait_mode=False, verbose=True):

self.H, self.W = H, W

self.verbose = verbose

self.num_landmarks = num_landmarks

self.dtype = dtype

self.device = device

self.is_portrait_mode = is_portrait_mode

if model_execution_lock is None:

model_execution_lock = nullcontext()

self.model_execution_lock = model_execution_lock

print(f" 🔧 Initializing RAFT on {device}...")

self.raft_tracker = RaftTracker(

device='dml', dtype=torch.float32, H=H, W=W, num_landmarks=num_landmarks

)

# State

self.current_flows = self.generate_grid_points()

self.current_masks = torch.ones(num_landmarks, dtype=torch.bool, device=device)

self.keyframe_tensor = None

# Pre-allocate keyframe flows

self.keyframe_flows = self.current_flows.clone()

self.frame_count = 0

# --- INTEGRATION BUFFERS (From original tracker.py) ---

self.fresh_integration_mask = torch.zeros(num_landmarks, dtype=torch.bool, device=device)

self.fresh_integration_mask_cp = torch.zeros(num_landmarks, dtype=torch.bool, device=device)

# Allocate enough buffer space for operations

self.arange_buf = torch.arange(num_landmarks + num_landmarks * 5 + 10, device=self.device)

self.pending_integration_command = None

self.candidate_lock = threading.Lock()

print(f" ✅ Tracker Initialized: {num_landmarks} points")

def generate_grid_points(self):

"""Generate grid in [X, Y] format."""

aspect = self.W / self.H

if self.num_landmarks == 2:

return torch.tensor([[self.W / 4, self.H / 2],

[3 * self.W / 4, self.H / 2]], dtype=self.dtype, device=self.device)

elif self.num_landmarks == 1:

return torch.tensor([[self.W / 2, self.H / 2]], dtype=self.dtype, device=self.device)

elif self.num_landmarks == 0:

return torch.empty((0, 2), dtype=self.dtype, device=self.device)

ny = int(np.sqrt(self.num_landmarks / aspect))

nx = int(np.ceil(self.num_landmarks / ny))

y = np.linspace(30, self.H - 30, ny)

x = np.linspace(30, self.W - 30, nx)

xv, yv = np.meshgrid(x, y)

points = np.stack([xv.ravel(), yv.ravel()], axis=1).astype(np.float32)[:self.num_landmarks]

if len(points) < self.num_landmarks:

points = np.vstack([points, np.tile(points[-1:], (self.num_landmarks - len(points), 1))])

return torch.from_numpy(points).to(self.device)

def get_and_clear_fresh_mask(self):

"""Get and clear fresh mask atomically."""

fresh = self.fresh_integration_mask_cp.copy_(self.fresh_integration_mask)

self.fresh_integration_mask.zero_()

return fresh

def set_candidate_points(self, points, z_values, amp_values_norm, frame_tensor):

"""Store candidates WITH their pre-sampled z-values."""

with self.candidate_lock:

# We clone frame_tensor to ensure it persists safely until integration runs

self.pending_integration_command = {

"candidates_S": points,

"z_values_S": z_values,

"amp_values_norm_S": amp_values_norm,

"keyframe_S": frame_tensor.detach().clone().contiguous(),

}

def _execute_integration_command(self, current_frame_tensor):

"""Execute integration and return z-values for fresh points."""

command = None

with self.candidate_lock:

if self.pending_integration_command is not None:

command = self.pending_integration_command

self.pending_integration_command = None

if command is None:

return None

candidates_S = command["candidates_S"]

z_values_S = command["z_values_S"]

amp_values_norm_S = command["amp_values_norm_S"]

keyframe_S = command["keyframe_S"]

is_dead_mask = ~self.current_masks

num_dead = is_dead_mask.sum()

num_candidates = len(candidates_S)

# Non-blocking replacement logic

rand_indices = torch.randperm(num_candidates, device=self.device)

mask_indices = self.arange_buf[:num_candidates] < num_dead

indices = rand_indices[mask_indices]

selected_S = candidates_S[indices]

selected_zs_S = z_values_S[indices]

selected_amps_norm_S = amp_values_norm_S[indices]

self.verb_print(f" ... running S→T mini-track for {len(selected_S)} points")

# --- PREPARE INPUTS FOR RAFT (APPLYING THE FIX) ---

# Ensure dimensions match what RAFT expects: (1, C, H, W)

# keyframe_S is (H, W, C) -> unsqueeze -> (1, H, W, C) -> permute -> (1, C, H, W)

if keyframe_S.ndim == 3:

keyframe_S_in = keyframe_S.unsqueeze(0).permute(0, 3, 1, 2)

current_frame_in = current_frame_tensor.unsqueeze(0).permute(0, 3, 1, 2)

else:

keyframe_S_in = keyframe_S.permute(0, 3, 1, 2)

current_frame_in = current_frame_tensor.permute(0, 3, 1, 2)

if self.is_portrait_mode:

selected_S = selected_S[:, [1, 0]]

# Note: Permuting the image itself for portrait mode if needed is handled by logic

# outside RAFT usually, but if your model is landscape-only, the above permute

# handles the channel ordering. If you need 90deg rotation, that's separate.

# Assuming standard behavior here as per main update loop.

# Track Candidates

selected_T, valid_T, _, _ = self.raft_tracker.track_points(

keyframe_S_in,

current_frame_in,

selected_S.to(dtype=torch.float32)

)

if self.is_portrait_mode:

selected_T = selected_T[:, [1, 0]]

# Flatten outputs if batch dim exists

if selected_T.ndim == 3: selected_T = selected_T[0]

if valid_T.ndim == 2: valid_T = valid_T[0]

dead_indices = torch.where(is_dead_mask)[0]

valid_indices_into_dead = torch.where(valid_T)[0]

valid_new_indices = dead_indices[valid_indices_into_dead]

valid_new_points_T = selected_T[valid_T]

valid_zs_S = selected_zs_S[valid_T]

valid_amps_norm_S = selected_amps_norm_S[valid_T]

final_integration_mask = valid_new_indices

self.current_flows[final_integration_mask] = valid_new_points_T

self.current_masks[final_integration_mask] = True

self.fresh_integration_mask[final_integration_mask] = True

# self.verb_print(f" ✓ Synced {len(valid_new_indices)} new points")

return {

'integrated': True,

'fresh_indices': valid_new_indices,

'fresh_z_values': valid_zs_S,

'normalized_amp_values': valid_amps_norm_S

}

def update(self, frame_tensor):

"""

# 1. SANITIZE INPUT

current_frame_clean = frame_tensor.detach().clone().contiguous()

# 2. INITIALIZATION

if self.keyframe_tensor is None:

self.keyframe_tensor = torch.empty_like(current_frame_clean)

self.keyframe_tensor.copy_(current_frame_clean)

self.keyframe_flows[:] = self.current_flows

self.frame_count += 1

return self.current_flows, self.current_masks, None, None

# 3. PREPARE INPUTS (Permute Here)

if self.keyframe_tensor.ndim == 3:

keyframe_input = self.keyframe_tensor.unsqueeze(0).permute(0, 3, 1, 2)

current_input = current_frame_clean.unsqueeze(0).permute(0, 3, 1, 2)

else:

keyframe_input = self.keyframe_tensor.permute(0, 3, 1, 2)

current_input = current_frame_clean.permute(0, 3, 1, 2)

points_input = self.keyframe_flows.clone()

# 4. TRACK (T-1 -> T)

new_points, valid, _, _ = self.raft_tracker.track_points(

keyframe_input,

current_input,

points_input

)

if new_points.ndim == 3: new_points = new_points[0]

if valid.ndim == 2: valid = valid[0]

# Filter

self.current_masks[~valid] = False

self.current_flows[self.current_masks] = new_points[self.current_masks]

# 5. INTEGRATION (Added Back)

fresh_mask = self.get_and_clear_fresh_mask()

integration_result = self._execute_integration_command(current_frame_clean)

# 6. UPDATE KEYFRAME (T becomes T-1)

self.keyframe_tensor.copy_(current_frame_clean)

self.keyframe_flows[:] = self.current_flows

self.frame_count += 1

return self.current_flows, self.current_masks, fresh_mask, integration_result

def verb_print(self, *args, **kwargs):

if self.verbose:

"""Main 3D tracking pipeline."""

def __init__(

self, W, H, K, num_landmarks=2000, dtype=torch.float32, device='cuda', dsp=None, is_portrait_mode=False,

verbose=True

):

self.W, self.H = W, H

self.verbose = verbose

self.K = K

self.device = device

self.num_landmarks = num_landmarks

self.dtype = dtype

self.tracker_lock = threading.Lock()

self.model_execution_lock = threading.Lock()

# Grids sample does not support float16 on CPU, neither is it better

self.tracker = TrackingPipeline(

H, W, num_landmarks=num_landmarks, dtype=dtype, device=device,

model_execution_lock=self.model_execution_lock, is_portrait_mode=is_portrait_mode, verbose=verbose

)

self.accum_z_amp = torch.ones((num_landmarks, 2), dtype=dtype, device=device)

self.accum_z_amp[:, 1] = 0.3

self.accum_valid = torch.zeros(num_landmarks, device=device, dtype=torch.bool)

self.depth_sampler = MultiThreadedZSampler(

H=H, W=W, K=K, num_landmarks=num_landmarks, tracker=self.tracker, parent=self, dsp=dsp,

model_execution_lock=self.model_execution_lock, verbose=verbose

)

self.color_heuristics_echo_db = -1.0

self.alive_mask = torch.ones(num_landmarks, dtype=torch.bool, device=device)

self.frame_count = 0

self._absorbance = torch.tensor(-6, dtype=dtype, device=device)

self.fresh_mask = torch.ones(num_landmarks, dtype=torch.bool, device=device)

self.xyz_amp = torch.empty((self.num_landmarks, 4),

device=self.device,

dtype=self.dtype)

# Buffers

N = self.num_landmarks

# buffer for indices of not_fresh elements (worst case: all landmarks)

self.not_fresh_idx = torch.empty(N, dtype=torch.int64, device=self.device)

self.min_distance = torch.empty(N, dtype=dtype, device=device)

# Debug tools ===========================================================================

self.alive_mask_ones_debug = torch.ones(num_landmarks, dtype=torch.bool, device=device)

self.alive_mask_zeros_debug = torch.zeros(num_landmarks, dtype=torch.bool, device=device)

self.landmarks_debug = torch.empty((num_landmarks, 3), dtype=dtype, device=device)

self.landmarks_debug[..., 0] = W/2

self.landmarks_debug[..., 1] = H/2

self.landmarks_debug[..., 2] = 1.0

self.landmarks4_debug = torch.empty((num_landmarks, 4), dtype=dtype, device=device)

self.landmarks4_debug[..., 0] = W / 2

self.landmarks4_debug[..., 1] = H / 2

self.landmarks4_debug[..., 2] = 1.0

self.landmarks4_debug[..., 3] = 1.0

def verb_print(self, *args, **kwargs):

if self.verbose:

print(*args, **kwargs)

def close(self):

self.depth_sampler.close()

def absorbance(self):

return self._absorbance

def room_head_geometry(self):

return self.depth_sampler.room_head_geometry()

@stable_checker

def track(

self, frame, frame_timestamp_sec, should_queue_regeneration=False, rotation_matrix=None

) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool]:

"""Main tracking with IMMEDIATE z-value initialization for fresh points."""

# Ignore absorbance estimation completely

# if self.frame_count % 4 == 0:

# self.color_heuristics_echo_db = self.depth_sampler.absorption_estimator.estimate_echo_db(frame) + self.color_heuristics_echo_db * 0.8

# self._absorbance = self.color_heuristics_echo_db + self.depth_sampler.depth_heuristics_echo_db

self.frame_count += 1

# Next is from depth map

with self.tracker_lock:

# This is Not a bottleneck

landmarks, alive_mask, _, integration_result = self.tracker.update(frame)

self.alive_mask = alive_mask

fresh_mask = self.fresh_mask

fresh_mask.zero_()

just_initialized = False

if integration_result is not None and integration_result.get('integrated'):

just_initialized = True

fresh_indices = integration_result['fresh_indices']

fresh_zs = integration_result['fresh_z_values']

normalized_amp_values = integration_result['normalized_amp_values']

if fresh_zs.dtype != self.dtype:

fresh_zs = fresh_zs.to(self.dtype)

self.accum_z_amp[fresh_indices, 0] = fresh_zs

self.accum_z_amp[fresh_indices, 1] = normalized_amp_values

self.accum_valid[fresh_indices] = True

self.depth_sampler.integration_period = True

fresh_mask[fresh_indices] = True

self.verb_print(f" 🎉 IMMEDIATE HARD RESET: fresh points initialized!")

self.depth_sampler.queue_work(

frame, self.tracker.current_flows, self.tracker.current_masks, frame_timestamp_sec,

should_queue_regeneration=should_queue_regeneration, rotation_matrix=rotation_matrix

)

z_data = self.depth_sampler.get_depth_data()

# self.xyz_amp[:, :2] = landmarks[:, :2]

# self.xyz_amp[:, 2:] = self.accum_z_amp

# return self.xyz_amp, alive_mask, fresh_mask, just_initialized

if z_data is not None:

# not_fresh = ~fresh_mask

# if (

# integration_result is None or not integration_result.get('integrated')

# ):

# ) and not self.depth_sampler.integration_period: #

# TODO: For now

self.accum_z_amp[:, 0] = z_data['zs']

# self.accum_z_amp[not_fresh, 0] = z_data['zs'][not_fresh]

# new_zs = z_data['zs'][not_fresh]

# old_zs = self.accum_z_amp[not_fresh, 0]

# # Calculate alpha | len is fine here because we use CPU in this pipeline

# # min_distance = torch.min(new_zs, old_zs, out=self.min_distance[:len(not_fresh)])

# # alpha = (min_distance.neg_()).mul_(0.1).exp_()

# alpha = 0.8 # Since computation is really slow and depth does not flicker, We do this. Low pass filter handles softness

# new_zs.mul_(alpha)

# if isinstance(alpha, torch.Tensor):

# alpha.neg_().add_(1.0)

# else:

# alpha = 1.0 - alpha

# old_zs.mul_(alpha)

# self.accum_z_amp[not_fresh, 0] = old_zs.add_(new_zs)

# Amps will forever stay consistent

self.xyz_amp[:, :2] = landmarks[:, :2]

self.xyz_amp[:, 2:] = self.accum_z_amp

def __init__(

self, W, H, K, view_sensor, head_mounted_sensor, num_landmarks,

dtype=torch.float32, device='cuda', cutoff_freq=4.3, dsp=None, is_portrait_mode=False, verbose=True

):

# Make the entire tracking Pipeline with all depth estimation code on CPU and IGPU

# This guarantees DSP runs at maximum CUDA performance

tracker_device = 'cpu'

super().__init__(

W, H, K, num_landmarks=num_landmarks, dtype=dtype, device=tracker_device, dsp=dsp,

is_portrait_mode=is_portrait_mode, verbose=verbose

)

self.tracker_device = tracker_device

self.main_device = device

self.cutoff_freq = cutoff_freq

self.lowpass_filter = LowPassFilter(

cutoff_freq=cutoff_freq, dtype=dtype, device=device, max_nodes=10

)

self.camera_transform_tracker = CameraProjection(

num_landmarks=num_landmarks, H=H, W=W, K=K, dtype=dtype, device=tracker_device

)

self.camera_transform_dsp = CameraProjection(

num_landmarks=num_landmarks, H=H, W=W, K=K, dtype=dtype, device=device

)

self.dtype = dtype

self.view_sensor = view_sensor

self.head_mounted_sensor = head_mounted_sensor

self.sensor_id = makeid(10)

# Cutoff freq and history window must match precisely

self.view_sensor.make_reading(

self.sensor_id,

cutoff_freq=cutoff_freq

)

if self.view_sensor is not self.head_mounted_sensor:

self.head_mounted_sensor.make_reading(

self.sensor_id,

cutoff_freq=cutoff_freq

)

self.velocities = torch.zeros([num_landmarks, 3], dtype=dtype, device=tracker_device)

self.position_offsets = torch.zeros([num_landmarks, 3], dtype=dtype, device=tracker_device)

self.frozen_world_landmarks = torch.zeros([num_landmarks, 4], dtype=dtype, device=tracker_device)

self.positions_buf = torch.zeros([num_landmarks, 4], dtype=dtype, device=self.main_device)

self.gpu_reset_values = torch.zeros([num_landmarks, 3], dtype=dtype, device=self.main_device)

self.gpu_fresh_mask_ = torch.zeros(num_landmarks, dtype=torch.bool, device=self.main_device)

self.frozen_world_landmarks[:, 0] = W / 2

self.frozen_world_landmarks[:, 1] = H / 2

self.frozen_world_landmarks[:, 2] = 1.0

self.frozen_world_landmarks[:, 3] = 0.4

self.last_points_in_world_frame = torch.zeros([num_landmarks, 3], dtype=dtype, device=tracker_device)

self.last_frame_timestamp_sec = None

self.is_updated_mask = torch.zeros(num_landmarks, dtype=torch.bool, device=device)

self.is_new_updates = False

self.is_updated_mask_change_lock = threading.Lock()

self.rotation_gpu = torch.empty((3, 3), dtype=self.dtype, device=device)

self.global_point_buf = torch.empty((num_landmarks, 4), dtype=dtype, device=tracker_device)

self.output_positions = torch.empty((num_landmarks, 4), dtype=dtype, device=tracker_device)

self.point_in_camera_frame_buf = torch.empty((num_landmarks, 4), dtype=dtype, device=device)

self.rotation_out = torch.empty(

(3, 3),

dtype=self.dtype,

device=self.device

)

def track(self, frame, frame_timestamp_sec, should_queue_regeneration=False, rotation_matrix=None) -> tuple:

points_, alive_mask, fresh_mask, just_initialized = super().track(

frame, frame_timestamp_sec, should_queue_regeneration=should_queue_regeneration,

rotation_matrix=rotation_matrix

)

if points_ is None:

return None, None

imu_to_world_rotation = self.view_sensor.rotation_matrix_when(frame_timestamp_sec)

if imu_to_world_rotation is None:

return None, None

point_in_camera_frame = self.camera_transform_tracker.pixels_to_camera_(points_)

point_in_camera_frame[..., :2] *= -1 # Flip XY

point_in_imu_frame = point_in_camera_frame[..., :3]

point_in_world_frame = point_in_imu_frame @ imu_to_world_rotation.T

# This is not thread safe

global_point = self.global_point_buf

global_point[..., :3] = point_in_world_frame

global_point[..., 3:] = point_in_camera_frame[..., 3:]

# # We want, velocity on dead nodes and also "NOT" Lowpass new nodes, reset the lowpass for specific indices

if just_initialized:

fresh_mask_ = fresh_mask

reset_values = point_in_world_frame[fresh_mask, :3]

if reset_values.device != self.main_device and 'cuda' in str(self.main_device):

# This is already on the CPU

length = fresh_mask.sum()

if 'cpu' in str(reset_values.device):

reset_values = reset_values.pin_memory()

reset_values = self.gpu_reset_values[:length].copy_(reset_values, non_blocking=True)

if 'cpu' in str(fresh_mask_.device):

fresh_mask_ = fresh_mask_.pin_memory()

fresh_mask_ = self.gpu_fresh_mask_.copy_(fresh_mask_, non_blocking=True)

self.lowpass_filter.values[:, fresh_mask_, :3] = reset_values

self.lowpass_filter.filtered[:, fresh_mask_, :3] = reset_values

fresh_amps = point_in_camera_frame[fresh_mask, 3]

# Ensure device matching (gpu_reset_values logic handles XYZ, we handle Amp here)

if fresh_amps.device != self.lowpass_filter.values.device:

fresh_amps = fresh_amps.to(self.lowpass_filter.values.device, non_blocking=True)

self.lowpass_filter.values[:, fresh_mask_, 3] = fresh_amps

self.lowpass_filter.filtered[:, fresh_mask_, 3] = fresh_amps

# Freeze world positions for newly dead points

self.frozen_world_landmarks[alive_mask] = global_point[alive_mask]

self.velocities[fresh_mask] = 0.0

# Reset for fresh points

if just_initialized:

with self.is_updated_mask_change_lock:

self.is_new_updates = True

self.is_updated_mask[fresh_mask] = True

# For velocity calculation, use current tracked position

velocity_updatable_mask = ~fresh_mask & alive_mask

position_updatable_mask = ~fresh_mask & ~alive_mask

if self.last_points_in_world_frame is not None and self.last_frame_timestamp_sec is not None:

delta = (frame_timestamp_sec - self.last_frame_timestamp_sec)

if delta > 0.00001:

vel = (point_in_world_frame[velocity_updatable_mask] - self.last_points_in_world_frame[velocity_updatable_mask]) / delta

self.velocities[velocity_updatable_mask] = vel * 0.01 + self.velocities[velocity_updatable_mask] * 0.99

self.frozen_world_landmarks[position_updatable_mask, :3] += self.velocities[position_updatable_mask] * delta

self.last_frame_timestamp_sec = frame_timestamp_sec

self.last_points_in_world_frame = point_in_world_frame # Only update for alive!

# Use frozen positions for dead points

output_positions = self.output_positions.copy_(global_point)

output_positions[position_updatable_mask] = self.frozen_world_landmarks[position_updatable_mask]

self.lowpass_filter.add_node(output_positions, frame_timestamp_sec)

def clear_updated_mask(self):

if self.is_new_updates: # Put this check up here,

# The idea is we must never clear right after thread ands new.

with self.is_updated_mask_change_lock:

self.is_updated_mask[:] = False

self.is_new_updates = False

return True

return False

@stable_checker

def get_position_global(self, t) -> Union[None, torch.Tensor]:

# This will be derived in the DSP thread will combine with high passed Accel

highpass_position = self.head_mounted_sensor.get_highpass_position(self.sensor_id, t)

positions = self.lowpass_filter.get_value(t)

if positions is None:

return None

positions = self.positions_buf.copy_(positions)

if highpass_position is not None:

# Inverse of the camera

positions[..., :3] -= highpass_position

return positions

@stable_checker

def get_position_screen_space(self, t):

# This is CUDA

point_in_world_frame = self.get_position_global(t)

if point_in_world_frame is None:

return None, self.alive_mask

return self.transform_global_point_to_screen_space_(t, point_in_world_frame), self.alive_mask

def transform_global_point_to_screen_space_(self, frame_timestamp_sec, point_in_world_frame, include_c=True):

# We use CUDA here

imu_to_world_rotation = self.head_mounted_sensor.rotation_matrix_when(

frame_timestamp_sec, device=self.main_device

)

if imu_to_world_rotation is None:

return None

# Transform: World → IMU → Camera

point_in_world_frame[..., :3] = point_in_world_frame[..., :3] @ imu_to_world_rotation

# Apply the flip fix we discovered

point_in_world_frame[..., :2] *= -1 # Flip XY

pixel_point = self.camera_transform_dsp.camera_to_pixels_(point_in_world_frame, include_c=include_c)

return pixel_point

def get_rotation(self, t):

if t == 'now':

return self.head_mounted_sensor.rotation_matrix_now(device=self.main_device)

# The inverse is the true camera rotation

src = torch.from_numpy(self.head_mounted_sensor.rotation_matrix_when(t, device=self.main_device))

self.rotation_gpu.copy_(src, non_blocking=True)

"""Maximum performance rendering with BGR-native path and pinned memory."""

def __init__(self, H, W, num_landmarks, dtype=torch.float32, device='cuda', thickness=2, lut_size=4096):

super().__init__(dtype=dtype, device=device, lut_size=lut_size)

self.H = H

self.W = W

self.num_landmarks = num_landmarks

self.dtype = dtype

self.device = device

self.thickness = thickness

# Pre-compute circle mask

r = int(thickness)

yy, xx = torch.meshgrid(

torch.arange(-r, r + 1, device=device),

torch.arange(-r, r + 1, device=device),

indexing='ij'

)

circle_mask = (xx * xx + yy * yy) <= (r * r)

self.offsets = torch.stack([xx[circle_mask], yy[circle_mask]], dim=1)

self.num_offsets = self.offsets.shape[0]

self.color_lut_bgr = self.color_lut_rgb[:, [2, 1, 0]]

# Pre-allocate rendering buffers

max_rendered = num_landmarks * self.num_offsets

self.expanded_buffer = torch.zeros((max_rendered, 2), dtype=torch.int64, device=device)

self.color_buffer = torch.zeros((max_rendered, 3), dtype=torch.uint8, device=device)

print(f"✓ HighPerformanceRenderer: BGR-native, {self.lut_size} LUT entries, pinned buffers")

@torch.no_grad()

def get_colors_from_lut(self, z_values):

"""Fast color lookup using pre-computed BGR LUT."""

z_clamped = torch.clamp(z_values, 0.01, 50000.0)

log_z = torch.log10(z_clamped)

t = (log_z - self.log_min) / self.log_range

indices = (t * (self.lut_size - 1)).long()

indices = torch.clamp(indices, 0, self.lut_size - 1)

return self.color_lut_bgr[indices]

@torch.no_grad()

def render(self, frame, points, z_values=None, override_color=None):

# """Ultra-fast rendering with BGR output."""

if points is None:

return frame

pts = points[:, :2].long()

if pts.shape[0] == 0:

return frame

valid = (pts[:, 0] >= 0) & (pts[:, 0] < self.W) & (pts[:, 1] >= 0) & (pts[:, 1] < self.H)

pts = pts[valid]

if pts.shape[0] == 0:

return frame

# Get BGR colors

if override_color is not None:

colors = override_color.expand(pts.shape[0], 3).byte()

elif z_values is not None:

z_vals = z_values[valid]

colors = self.get_colors_from_lut(z_vals)

else:

colors = torch.full((pts.shape[0], 3), 255, dtype=frame.dtype, device=frame.device)

# Vectorized expansion

expanded = pts.unsqueeze(1) + self.offsets.unsqueeze(0)

expanded[..., 0].clamp_(0, self.W - 1)

expanded[..., 1].clamp_(0, self.H - 1)

flat_pts = expanded.reshape(-1, 2)

flat_colors = colors.repeat_interleave(self.num_offsets, dim=0)

frame[flat_pts[:, 1], flat_pts[:, 0]] = flat_colors

"""Pinned memory pool for fast CPU-GPU transfers."""

def __init__(self, OH, OW, H, W, device='cuda'):

self.device = device

# Pre-allocate pinned memory buffers

# For RGB uint8 input frame

self.frame_buffer_rgb = torch.zeros((OH, OW, 3), dtype=torch.uint8, device=device)

# For BGR uint8 output frame

self.frame_buffer_bgr = torch.zeros((H, W, 3), dtype=torch.uint8).pin_memory()

self.frame_buffer_bgr_np = self.frame_buffer_bgr.numpy()

print(f"✓ PinnedMemoryPool: {H}x{W} buffers allocated")

def upload_frame(self, frame_np):

"""Fast upload from numpy to GPU via pinned memory."""

# Copy numpy data into pinned buffer (fast mem cpy)

self.frame_buffer_rgb.copy_(torch.from_numpy(frame_np))

# Non-blocking transfer to GPU

return self.frame_buffer_rgb

def download_frame(self, frame_gpu):

"""Fast download from GPU to numpy via pinned memory."""

# Non-blocking transfer to CPU pinned memory

self.frame_buffer_bgr.copy_(frame_gpu, non_blocking=False)

# Return numpy view (no copy)

"""Optimized CPU text rendering."""

@staticmethod

def render_stats_cpu(frame_np, fps, num_alive, total_points, num_fresh,

frame_count, total_fresh, regen_count, frames_until_regen):