300 lines
12 KiB
Python
300 lines
12 KiB
Python
import numpy as np
|
|
import cv2
|
|
from typing import Tuple, Optional, List
|
|
import os
|
|
import torch
|
|
|
|
def create_mask_from_tracks(forward_tracks: np.ndarray,
|
|
forward_visibility: np.ndarray,
|
|
frame_shape: Tuple[int, int],
|
|
radius: int = 10,
|
|
frame_idx: Optional[int] = None) -> np.ndarray:
|
|
|
|
H, W = frame_shape
|
|
|
|
if frame_idx is not None:
|
|
mask = np.zeros((H, W), dtype=np.uint8)
|
|
points = forward_tracks[frame_idx] # shape (N, 2)
|
|
visibility = forward_visibility[frame_idx] # shape (N,)
|
|
|
|
valid_points = points[visibility > 0]
|
|
|
|
for point in valid_points:
|
|
x, y = int(point[0]), int(point[1])
|
|
if 0 <= x < W and 0 <= y < H:
|
|
cv2.circle(mask, (x, y), radius, 1, -1)
|
|
|
|
return mask
|
|
else:
|
|
T = forward_tracks.shape[0]
|
|
masks = np.zeros((T, H, W), dtype=np.uint8)
|
|
|
|
for t in range(T):
|
|
points = forward_tracks[t] # shape (N, 2)
|
|
visibility = forward_visibility[t] # shape (N,)
|
|
|
|
valid_points = points[visibility > 0]
|
|
|
|
for point in valid_points:
|
|
x, y = int(point[0]), int(point[1])
|
|
if 0 <= x < W and 0 <= y < H:
|
|
cv2.circle(masks[t], (x, y), radius, 1, -1)
|
|
|
|
return masks
|
|
|
|
def detect_empty_regions(forward_tracks: np.ndarray,
|
|
forward_visibility: np.ndarray,
|
|
frame_shape: Tuple[int, int],
|
|
frame_idx: int,
|
|
radius: int = 10,
|
|
min_region_size: int = 100) -> List[Tuple[int, int, int, int]]:
|
|
|
|
mask = create_mask_from_tracks(forward_tracks, forward_visibility, frame_shape, radius, frame_idx)
|
|
|
|
empty_mask = 1 - mask
|
|
|
|
contours, _ = cv2.findContours(empty_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
|
|
empty_regions = []
|
|
for contour in contours:
|
|
area = cv2.contourArea(contour)
|
|
if area >= min_region_size:
|
|
x, y, w, h = cv2.boundingRect(contour)
|
|
empty_regions.append((x, y, w, h))
|
|
|
|
return empty_regions
|
|
|
|
def has_data_in_region(backward_tracks: np.ndarray,
|
|
backward_visibility: np.ndarray,
|
|
spatial_region: Tuple[int, int, int, int],
|
|
frame_idx: int,
|
|
min_points: int = 1) -> bool:
|
|
x, y, w, h = spatial_region
|
|
points = backward_tracks[frame_idx] # shape (N, 2)
|
|
visibility = backward_visibility[frame_idx] # shape (N,)
|
|
|
|
valid_points = points[visibility > 0]
|
|
|
|
if len(valid_points) == 0:
|
|
return False
|
|
|
|
in_region = ((valid_points[:, 0] >= x) & (valid_points[:, 0] < x + w) &
|
|
(valid_points[:, 1] >= y) & (valid_points[:, 1] < y + h))
|
|
|
|
return np.sum(in_region) >= min_points
|
|
|
|
def extract_trajectory_with_indices(tracks: np.ndarray,
|
|
visibility: np.ndarray,
|
|
spatial_region: Tuple[int, int, int, int],
|
|
frame_idx: int) -> Tuple[np.ndarray, np.ndarray, List[int]]:
|
|
|
|
x, y, w, h = spatial_region
|
|
points = tracks[frame_idx] # shape (N, 2)
|
|
vis = visibility[frame_idx] # shape (N,)
|
|
|
|
valid_mask = vis > 0
|
|
in_region_mask = ((points[:, 0] >= x) & (points[:, 0] < x + w) &
|
|
(points[:, 1] >= y) & (points[:, 1] < y + h))
|
|
|
|
selected_indices = np.where(valid_mask & in_region_mask)[0]
|
|
|
|
if len(selected_indices) == 0:
|
|
return np.empty((tracks.shape[0], 0, 2)), np.empty((tracks.shape[0], 0)), []
|
|
|
|
extracted_tracks = tracks[:, selected_indices, :]
|
|
extracted_visibility = visibility[:, selected_indices]
|
|
|
|
return extracted_tracks, extracted_visibility, selected_indices.tolist()
|
|
|
|
|
|
def time_reverse(tracks: np.ndarray, visibility: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
|
|
return np.flip(tracks, axis=0), np.flip(visibility, axis=0)
|
|
|
|
|
|
|
|
def integrate_tracking_results(forward_tracks: np.ndarray,
|
|
forward_visibility: np.ndarray,
|
|
backward_tracks: np.ndarray,
|
|
backward_visibility: np.ndarray,
|
|
frame_shape: Tuple[int, int],
|
|
radius: int = 10,
|
|
min_region_size: int = 100,
|
|
min_points: int = 1,
|
|
output_dir: Optional[str] = None) -> Tuple[np.ndarray, np.ndarray]:
|
|
|
|
T = forward_tracks.shape[0]
|
|
|
|
backward_tracks, backward_visibility = time_reverse(backward_tracks, backward_visibility)
|
|
|
|
if output_dir is not None:
|
|
import os
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
print(f"Mask images will be saved to: {output_dir}")
|
|
|
|
integrated_tracks_list = [forward_tracks]
|
|
integrated_visibility_list = [forward_visibility]
|
|
|
|
global_used_indices = set()
|
|
|
|
for frame_idx in range(T):
|
|
print(f"Processing frame {frame_idx}...")
|
|
|
|
current_integrated_tracks = np.concatenate(integrated_tracks_list, axis=1)
|
|
current_integrated_visibility = np.concatenate(integrated_visibility_list, axis=1)
|
|
|
|
if output_dir is not None:
|
|
save_mask_visualization(
|
|
current_integrated_tracks, current_integrated_visibility,
|
|
frame_shape, frame_idx, radius, output_dir,
|
|
suffix="before_integration"
|
|
)
|
|
|
|
empty_regions = detect_empty_regions(
|
|
current_integrated_tracks, current_integrated_visibility,
|
|
frame_shape, frame_idx, radius, min_region_size
|
|
)
|
|
print(f" Found {len(empty_regions)} empty regions")
|
|
|
|
frame_new_tracks = []
|
|
frame_new_visibility = []
|
|
|
|
for region_idx, spatial_region in enumerate(empty_regions):
|
|
available_indices = [i for i in range(backward_tracks.shape[1]) if i not in global_used_indices]
|
|
|
|
if len(available_indices) == 0:
|
|
print(f" Region {region_idx}: No more available backward tracks")
|
|
continue
|
|
|
|
available_backward_tracks = backward_tracks[:, available_indices, :]
|
|
available_backward_visibility = backward_visibility[:, available_indices]
|
|
|
|
if has_data_in_region(available_backward_tracks, available_backward_visibility, spatial_region, frame_idx, min_points):
|
|
backward_trajectory, backward_vis, local_extracted_indices = extract_trajectory_with_indices(
|
|
available_backward_tracks, available_backward_visibility, spatial_region, frame_idx
|
|
)
|
|
|
|
if backward_trajectory.shape[1] > 0:
|
|
global_extracted_indices = [available_indices[i] for i in local_extracted_indices]
|
|
|
|
print(f" Region {region_idx}: Extracted {len(global_extracted_indices)} tracks: {global_extracted_indices}")
|
|
|
|
frame_new_tracks.append(backward_trajectory)
|
|
frame_new_visibility.append(backward_vis)
|
|
|
|
global_used_indices.update(global_extracted_indices)
|
|
|
|
print(f" Total used indices so far: {len(global_used_indices)}")
|
|
else:
|
|
print(f" Region {region_idx}: No backward data found")
|
|
|
|
if frame_new_tracks:
|
|
integrated_tracks_list.extend(frame_new_tracks)
|
|
integrated_visibility_list.extend(frame_new_visibility)
|
|
|
|
|
|
if output_dir is not None:
|
|
final_integrated_tracks = np.concatenate(integrated_tracks_list, axis=1)
|
|
final_integrated_visibility = np.concatenate(integrated_visibility_list, axis=1)
|
|
|
|
save_mask_visualization(
|
|
final_integrated_tracks, final_integrated_visibility,
|
|
frame_shape, frame_idx, radius, output_dir,
|
|
suffix="after_integration"
|
|
)
|
|
|
|
integrated_tracks = np.concatenate(integrated_tracks_list, axis=1)
|
|
integrated_visibility = np.concatenate(integrated_visibility_list, axis=1)
|
|
|
|
|
|
if output_dir is not None:
|
|
print(f"\nCreating comparison images...")
|
|
for frame_idx in range(T):
|
|
create_mask_comparison(
|
|
forward_tracks, forward_visibility,
|
|
integrated_tracks, integrated_visibility,
|
|
frame_shape, frame_idx, radius, output_dir
|
|
)
|
|
|
|
print(f"\nFinal summary:")
|
|
print(f"Used backward track indices: {sorted(global_used_indices)}")
|
|
print(f"Total backward tracks used: {len(global_used_indices)}")
|
|
print(f"Original backward tracks: {backward_tracks.shape[1]}")
|
|
print(f"Final integrated tracks shape: {integrated_tracks.shape}")
|
|
|
|
return integrated_tracks, integrated_visibility
|
|
|
|
def save_mask_visualization(tracks: np.ndarray,
|
|
visibility: np.ndarray,
|
|
frame_shape: Tuple[int, int],
|
|
frame_idx: int,
|
|
radius: int,
|
|
output_dir: str,
|
|
suffix: str = "") -> None:
|
|
|
|
mask = create_mask_from_tracks(tracks, visibility, frame_shape, radius, frame_idx)
|
|
|
|
mask_vis = (mask * 255).astype(np.uint8)
|
|
|
|
mask_colored = cv2.applyColorMap(mask_vis, cv2.COLORMAP_JET)
|
|
|
|
mask_colored[mask == 0] = [0, 0, 0]
|
|
|
|
if suffix:
|
|
filename = f"mask_frame_{frame_idx:04d}_{suffix}.png"
|
|
else:
|
|
filename = f"mask_frame_{frame_idx:04d}.png"
|
|
|
|
filepath = os.path.join(output_dir, filename)
|
|
cv2.imwrite(filepath, mask_colored)
|
|
print(f" Saved mask: {filepath}")
|
|
|
|
def create_mask_comparison(forward_tracks: np.ndarray,
|
|
forward_visibility: np.ndarray,
|
|
integrated_tracks: np.ndarray,
|
|
integrated_visibility: np.ndarray,
|
|
frame_shape: Tuple[int, int],
|
|
frame_idx: int,
|
|
radius: int,
|
|
output_dir: str) -> None:
|
|
|
|
forward_mask = create_mask_from_tracks(forward_tracks, forward_visibility, frame_shape, radius, frame_idx)
|
|
|
|
integrated_mask = create_mask_from_tracks(integrated_tracks, integrated_visibility, frame_shape, radius, frame_idx)
|
|
|
|
added_mask = integrated_mask - forward_mask
|
|
|
|
H, W = frame_shape
|
|
comparison = np.zeros((H, W, 3), dtype=np.uint8)
|
|
|
|
comparison[forward_mask > 0] = [255, 0, 0]
|
|
|
|
comparison[added_mask > 0] = [0, 0, 255]
|
|
|
|
filename = f"comparison_frame_{frame_idx:04d}.png"
|
|
filepath = os.path.join(output_dir, filename)
|
|
cv2.imwrite(filepath, comparison)
|
|
print(f" Saved comparison: {filepath}")
|
|
|
|
|
|
def trajectory_integration(forward_tracks, forward_visibility, backward_tracks, backward_visibility, frame_shape, grid_size):
|
|
|
|
ft = forward_tracks.squeeze(0).cpu().numpy() # (T, N, 2)
|
|
fv = forward_visibility.squeeze(0).cpu().numpy() # (T, N)
|
|
|
|
bt = backward_tracks.squeeze(0).cpu().numpy() # (T, N, 2)
|
|
bv = backward_visibility.squeeze(0).cpu().numpy() # (T, N)
|
|
|
|
radius = min(frame_shape) / grid_size * 1.5
|
|
radius = max(int(round(radius)), 3)
|
|
min_region_size = radius ** 2
|
|
|
|
t, v = integrate_tracking_results(ft,fv,bt,bv,frame_shape,
|
|
radius=radius, min_region_size=min_region_size, min_points=1, output_dir=None)
|
|
|
|
t = torch.from_numpy(t).float().unsqueeze(0)
|
|
v = torch.from_numpy(v).float().unsqueeze(0)
|
|
|
|
return t,v
|
|
|
|
|