File size: 13,578 Bytes
e168a4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 | """
Detection Vector Tracker for EnergySnake
This module provides utilities to track high-dimensional vectors for each detection box
before YOLO decoder, maintaining correspondence with detection boxes.
"""
import torch
import numpy as np
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass, field
@dataclass
class DetectionVector:
"""
Represents a single detection with its high-dimensional vector and corresponding box.
"""
# High-dimensional vector before YOLO decoder (flattened from feature map)
vector: torch.Tensor # Shape: [C] where C is the feature dimension
# Detection box information [x1, y1, x2, y2, score, class_id]
bbox: torch.Tensor # Shape: [6]
# Position information for correspondence
grid_pos: Tuple[int, int] # (h_idx, w_idx) position on feature grid
# Feature map location info
feature_map_idx: int # Index in the flattened feature map (0-based)
# Additional metadata
confidence: float
class_id: int
image_idx: int # Batch index
@dataclass
class DetectionVectorBatch:
"""
Container for all detection vectors in a batch.
"""
detections: List[DetectionVector] = field(default_factory=list)
# Raw tensors for batch processing
raw_vectors: torch.Tensor = None # [N, C] all vectors concatenated
raw_bboxes: torch.Tensor = None # [N, 6] all bboxes concatenated
raw_yolo_output: torch.Tensor = None # [B, no, HW] original YOLO output
# Feature map dimensions
feature_h: int = 0
feature_w: int = 0
batch_size: int = 0
# Vector dimension per detection
vector_dim: int = 0
def add_detection(self, detection: DetectionVector):
"""Add a detection to the batch."""
self.detections.append(detection)
def finalize(self):
"""Convert list of detections to batch tensors."""
if not self.detections:
return
# Stack vectors and bboxes
self.raw_vectors = torch.stack([det.vector for det in self.detections])
self.raw_bboxes = torch.stack([det.bbox for det in self.detections])
def get_vectors_by_class(self, class_id: int) -> torch.Tensor:
"""Get all vectors for a specific class."""
if not self.detections:
return torch.empty(0, self.vector_dim)
class_vectors = [det.vector for det in self.detections if det.class_id == class_id]
return torch.stack(class_vectors) if class_vectors else torch.empty(0, self.vector_dim)
def get_bboxes_by_class(self, class_id: int) -> torch.Tensor:
"""Get all bboxes for a specific class."""
if not self.detections:
return torch.empty(0, 6)
class_bboxes = [det.bbox for det in self.detections if det.class_id == class_id]
return torch.stack(class_bboxes) if class_bboxes else torch.empty(0, 6)
class DetectionVectorTracker:
"""
Tracks high-dimensional vectors from YOLO detection head before decoder.
The YOLO detection head outputs high-dimensional features that are then decoded
into bounding boxes. This tracker captures those intermediate features.
"""
def __init__(self):
self.current_batch: Optional[DetectionVectorBatch] = None
self.history: List[DetectionVectorBatch] = []
# YOLO configuration (will be determined from actual model output)
self.reg_max = 16 # Default for YOLOv8
self.num_classes = 52 # Spinal structures
def extract_vectors_from_yolo_output(self,
yolo_output: torch.Tensor,
feature_maps: List[torch.Tensor],
detection_bboxes: torch.Tensor,
image_size: Tuple[int, int] = (544, 544)) -> DetectionVectorBatch:
"""
Extract high-dimensional vectors from YOLO output before decoder.
Args:
yolo_output: Raw YOLO output tensor [B, no, HW]
feature_maps: List of feature maps from YOLO backbone
detection_bboxes: Final detection boxes after NMS [B, M, 6]
image_size: Input image size (H, W)
Returns:
DetectionVectorBatch containing extracted vectors and correspondence
"""
batch_size, output_dim, hw = yolo_output.shape
# For YOLOv8-p2 model, the output is already flattened from multiple scales
# We need to reconstruct the multi-scale feature information
# Create new batch container
batch = DetectionVectorBatch(
feature_h=int(np.sqrt(hw)), # Will be updated below
feature_w=int(np.sqrt(hw)), # Will be updated below
batch_size=batch_size,
vector_dim=output_dim,
raw_yolo_output=yolo_output.clone()
)
# Get feature map dimensions from the actual feature maps
if feature_maps:
# Use the first (largest) feature map size
p2_shape = feature_maps[0].shape # [B, C, H, W]
batch.feature_h = p2_shape[2]
batch.feature_w = p2_shape[3]
# Reshape YOLO output to feature grid format
yolo_reshaped = yolo_output.permute(0, 2, 1).contiguous() # [B, HW, no]
# For YOLOv8-p2, HW includes contributions from multiple scales
# We'll work with the flattened format and map detections accordingly
# Process each image in the batch
for batch_idx in range(batch_size):
# Get detections for this image
img_detections = detection_bboxes[batch_idx]
valid_detections = img_detections[img_detections[:, 4] > 0] # Filter by confidence
# Process each valid detection
for det_idx, detection in enumerate(valid_detections):
bbox = detection # [x1, y1, x2, y2, score, class_id]
confidence = float(detection[4])
class_id = int(detection[5])
# Find the best matching position in YOLO output
# Since YOLO output is flattened across multiple scales, we find the closest match
center_x = (bbox[0] + bbox[2]) / 2
center_y = (bbox[1] + bbox[3]) / 2
# Try to map to the flattened YOLO output positions
# For YOLOv8-p2, we need to find the position that best matches the detection
best_match_idx = self._find_best_yolo_position(
bbox, yolo_reshaped[batch_idx], image_size, feature_maps
)
if best_match_idx is not None:
vector = yolo_reshaped[batch_idx, best_match_idx] # [output_dim]
# For multi-scale models, estimate the grid position
# This is an approximation since we can't perfectly reconstruct multi-scale mappings
grid_x = best_match_idx % batch.feature_w
grid_y = best_match_idx // batch.feature_w
# Create detection vector
det_vector = DetectionVector(
vector=vector,
bbox=bbox,
grid_pos=(int(grid_y), int(grid_x)),
feature_map_idx=best_match_idx,
confidence=confidence,
class_id=class_id,
image_idx=batch_idx
)
batch.add_detection(det_vector)
def _find_best_yolo_position(self, bbox, yolo_flat_output, image_size, feature_maps):
"""
Find the best matching position in YOLO output for a given detection.
For multi-scale YOLO, this is an approximation that finds the closest spatial match.
"""
center_x = (bbox[0] + bbox[2]) / 2
center_y = (bbox[1] + bbox[3]) / 2
# Simple strategy: map center to flattened position
# This is approximate but works for tracking purposes
h, w = image_size
flat_positions = yolo_flat_output.shape[0]
# Estimate spatial grid size (approximate)
grid_size = int(np.sqrt(flat_positions))
if grid_size * grid_size != flat_positions:
# Use closest square number
grid_size = int(np.sqrt(flat_positions))
# Map bbox center to grid position
grid_x = int(center_x * grid_size / w)
grid_y = int(center_y * grid_size / h)
# Clamp to valid range
grid_x = max(0, min(grid_x, grid_size - 1))
grid_y = max(0, min(grid_y, grid_size - 1))
# Convert to flat index
flat_idx = grid_y * grid_size + grid_x
# Ensure index is within bounds
if flat_idx >= flat_positions:
flat_idx = flat_positions - 1
return flat_idx
# Finalize batch processing
batch.finalize()
# Store current batch
self.current_batch = batch
self.history.append(batch)
return batch
def get_vector_shape_info(self) -> Dict[str, any]:
"""
Get information about the vector shapes and dimensions.
Returns:
Dictionary containing shape information
"""
if not self.current_batch:
return {"error": "No batch processed yet"}
batch = self.current_batch
yolo_output = batch.raw_yolo_output
# YOLO output shape: [B, no, HW]
# where no = reg_max * 4 + num_classes
# reg_max = 16, so no = 16 * 4 + 52 = 116
reg_max = self.reg_max
num_classes = self.num_classes
output_channels = reg_max * 4 + num_classes # 64 + 52 = 116
return {
"yolo_output_shape": list(yolo_output.shape),
"output_channels": output_channels,
"regression_channels": reg_max * 4, # 64 channels for bbox regression
"classification_channels": num_classes, # 52 channels for classification
"feature_map_size": (batch.feature_h, batch.feature_w),
"total_positions": batch.feature_h * batch.feature_w,
"high_dim_vector_shape": [output_channels], # Shape of each detection's vector
"vector_breakdown": {
"bbox_regression": [reg_max * 4], # DFL channels for 4 bbox coordinates
"class_logits": [num_classes] # Raw class scores before sigmoid
}
}
def get_batch_summary(self) -> Dict[str, any]:
"""Get summary statistics for the current batch."""
if not self.current_batch:
return {"error": "No batch processed yet"}
batch = self.current_batch
# Count detections by class
class_counts = {}
for det in batch.detections:
class_counts[det.class_id] = class_counts.get(det.class_id, 0) + 1
# Confidence statistics
confidences = [det.confidence for det in batch.detections]
return {
"total_detections": len(batch.detections),
"detections_by_class": class_counts,
"confidence_stats": {
"mean": np.mean(confidences) if confidences else 0,
"min": np.min(confidences) if confidences else 0,
"max": np.max(confidences) if confidences else 0
},
"vector_stats": {
"shape": list(batch.raw_vectors.shape) if batch.raw_vectors is not None else None,
"mean_norm": float(torch.mean(torch.norm(batch.raw_vectors, dim=1))) if batch.raw_vectors is not None else 0
}
}
def clear_history(self):
"""Clear processing history."""
self.history.clear()
def save_vectors_to_file(self, filepath: str, batch_idx: int = -1):
"""
Save detection vectors to file.
Args:
filepath: Path to save the vectors
batch_idx: Index of batch to save (-1 for current/latest)
"""
if batch_idx == -1 and self.current_batch:
batch = self.current_batch
elif 0 <= batch_idx < len(self.history):
batch = self.history[batch_idx]
else:
raise ValueError(f"Invalid batch_idx: {batch_idx}")
save_data = {
"batch_info": {
"batch_size": batch.batch_size,
"feature_map_size": (batch.feature_h, batch.feature_w),
"vector_dim": batch.vector_dim,
"total_detections": len(batch.detections)
},
"detections": [
{
"vector": det.vector.cpu().numpy().tolist(),
"bbox": det.bbox.cpu().numpy().tolist(),
"grid_pos": det.grid_pos,
"confidence": det.confidence,
"class_id": det.class_id,
"image_idx": det.image_idx
}
for det in batch.detections
]
}
import json
with open(filepath, 'w') as f:
json.dump(save_data, f, indent=2) |