lib/utils/multi_scale_feature_extractor.py

"""
Multi-Scale Feature Extractor for YOLO Detection

This module extracts raw 116-dimensional features from all 4 scales (P2, P3, P4, P5)
for each detection box, providing the most comprehensive visual representation.

Key features:
- Extracts features before DFL processing (116-dim vs 56-dim)
- Maintains exact correspondence with original spatial layout
- Preserves scale ordering (P2 -> P3 -> P4 -> P5)
- Handles exact 2D reconstruction without rotation/flipping
"""

import torch
import numpy as np
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass


@dataclass
class MultiScaleDetectionFeatures:
    """
    Comprehensive features for a single detection across all 4 scales.
    """
    # Detection information
    bbox: torch.Tensor  # [x1, y1, x2, y2, score, class_id]
    center_point: Tuple[float, float]  # (center_x, center_y)
    
    # Multi-scale features (4 scales × 116 dimensions)
    p2_features: torch.Tensor  # [116] - P2 scale (136×136, stride=4)
    p3_features: torch.Tensor  # [116] - P3 scale (68×68, stride=8)  
    p4_features: torch.Tensor  # [116] - P4 scale (34×34, stride=16)
    p5_features: torch.Tensor  # [116] - P5 scale (17×17, stride=32)
    
    # Feature positions (for verification/debugging)
    p2_position: Tuple[int, int]  # (x, y) in P2 feature map
    p3_position: Tuple[int, int]  # (x, y) in P3 feature map
    p4_position: Tuple[int, int]  # (x, y) in P4 feature map
    p5_position: Tuple[int, int]  # (x, y) in P5 feature map
    
    # Metadata
    confidence: float
    class_id: int
    image_idx: int
    
    # 多模态特征 (可选)
    text_features: Optional[torch.Tensor] = None  # [116] - ClinicalBERT文本特征
    multimodal_features: Optional[torch.Tensor] = None  # [580] - 视觉+文本融合特征
    
    @property
    def concatenated_features(self) -> torch.Tensor:
        """Get all scale features concatenated: [4×116] = [464]"""
        return torch.cat([self.p2_features, self.p3_features, self.p4_features, self.p5_features], dim=0)
    
    @property
    def feature_matrix(self) -> torch.Tensor:
        """Get features as matrix: [4, 116]"""
        return torch.stack([self.p2_features, self.p3_features, self.p4_features, self.p5_features], dim=0)


class MultiScaleFeatureExtractor:
    """
    Extracts multi-scale raw features (116-dim) for detection boxes.
    
    This extractor works with the pre-DFL features to capture the richest
    visual information across all 4 scales of YOLOv8-p2.
    """
    
    def __init__(self, input_size: Tuple[int, int] = (544, 544)):
        self.input_size = input_size
        self.feature_sizes = self._calculate_feature_sizes()
        
        # Verify total matches expected HW = 24565
        total_positions = sum(size[0] * size[1] for size in self.feature_sizes.values())
        assert total_positions == 24565, f"Total positions {total_positions} != expected 24565"
        
        # Calculate offsets for each scale in the concatenated tensor
        self._calculate_scale_offsets()

    def _adapt_to_hw(self, hw: int):
        """
        动态适应不同的HW特征图尺寸
        
        Args:
            hw: 实际的特征位置总数
        """
        # 常见的配置映射 - 只保留验证过的配置
        common_configs = {
            21760: {
                'feature_sizes': {
                    'P2': (128, 128),  # 16384 positions
                    'P3': (64, 64),    # 4096 positions
                    'P4': (32, 32),    # 1024 positions  
                    'P5': (16, 16)     # 256 positions
                },
                'scale_offsets': {
                    'P2': 0,
                    'P3': 16384,
                    'P4': 16384 + 4096,
                    'P5': 16384 + 4096 + 1024
                }
            }
        }
        
        if hw in common_configs:
            config = common_configs[hw]
            # 验证配置是否匹配
            total_positions = sum(h*w for h, w in config['feature_sizes'].values())
            if total_positions == hw:
                self.feature_sizes = config['feature_sizes']
                self.scale_offsets = config['scale_offsets']
                # 重新计算scale_ranges
                self._update_scale_ranges()
                return
        
        # 改进的动态计算配置，确保每个尺度都有合理的位置数
        # 使用更保守的比例分配，确保P5至少有一些位置
        
        # 为P5预留最小位置数
        min_p5_positions = 64  # 8×8的最小特征图
        
        # P2获取大约75%的位置
        p2_positions = hw * 3 // 4
        
        # 为P5预留位置
        p5_positions = max(min_p5_positions, hw // 32)  # 至少1/32的位置给P5
        
        remaining_for_p3_p4 = hw - p2_positions - p5_positions
        
        # P3和P4按3:1的比例分配剩余位置
        p3_positions = remaining_for_p3_p4 * 3 // 4
        p4_positions = remaining_for_p3_p4 - p3_positions
        
        # 调整确保总和等于hw
        total = p2_positions + p3_positions + p4_positions + p5_positions
        if total != hw:
            # 微调P2以匹配总数
            p2_positions += (hw - total)
        
        assert p2_positions > 0 and p3_positions > 0 and p4_positions > 0 and p5_positions > 0, "每个尺度都需要有正数位置"
        
        def positions_to_size(positions):
            import math
            side = int(math.sqrt(positions))
            return side, side
        
        p2_size = positions_to_size(p2_positions)
        p3_size = positions_to_size(p3_positions) 
        p4_size = positions_to_size(p4_positions)
        p5_size = positions_to_size(p5_positions)
        
        actual_p2 = p2_size[0] * p2_size[1]
        actual_p3 = p3_size[0] * p3_size[1]
        actual_p4 = p4_size[0] * p4_size[1]
        
        self.feature_sizes = {
            'P2': p2_size,
            'P3': p3_size, 
            'P4': p4_size,
            'P5': p5_size
        }
        
        self.scale_offsets = {
            'P2': 0,
            'P3': actual_p2,
            'P4': actual_p2 + actual_p3,
            'P5': actual_p2 + actual_p3 + actual_p4
        }
        
        # 重新计算scale_ranges
        self._update_scale_ranges()
        
        print(f"[DYNAMIC] 动态计算配置完成 HW={hw}")

    def _update_scale_ranges(self):
        """重新计算scale_ranges以匹配当前配置"""
        self.scale_ranges = {}
        scales = ['P2', 'P3', 'P4', 'P5']
        
        for scale in scales:
            h, w = self.feature_sizes[scale]
            positions = h * w
            offset = self.scale_offsets[scale]
            self.scale_ranges[scale] = (offset, offset + positions)
        
        # Update scale ranges
    
    def _calculate_feature_sizes(self) -> Dict[str, Tuple[int, int]]:
        """Calculate feature map sizes for all scales."""
        h, w = self.input_size
        
        return {
            'P2': (h // 4, w // 4),    # stride=4  -> 136×136 = 18,496
            'P3': (h // 8, w // 8),    # stride=8  -> 68×68   =  4,624  
            'P4': (h // 16, w // 16),  # stride=16 -> 34×34   =  1,156
            'P5': (h // 32, w // 32),  # stride=32 -> 17×17   =    289
        }
    
    def _calculate_scale_offsets(self):
        """Calculate offset positions for each scale in concatenated tensor."""
        self.scale_offsets = {}
        self.scale_ranges = {}
        
        scales = ['P2', 'P3', 'P4', 'P5']
        current_offset = 0
        
        for scale in scales:
            h, w = self.feature_sizes[scale]
            positions = h * w
            self.scale_offsets[scale] = current_offset
            self.scale_ranges[scale] = (current_offset, current_offset + positions)
            current_offset += positions
        
        # Scale offsets calculated
    
    def extract_multi_scale_features(
        self,
        raw_116_features: torch.Tensor,  # [B, 116, HW] - Raw features before DFL
        detection_bboxes: torch.Tensor,  # [B, N, 6] - Detection boxes [x1,y1,x2,y2,score,cls]
        image_size: Optional[Tuple[int, int]] = None,
        confidence_threshold: float = 0.0  # 添加置信度阈值参数，默认为0（不过滤）
    ) -> List[MultiScaleDetectionFeatures]:
        """Extract multi-scale 116-dimensional features for detection boxes."""
        if image_size is None:
            image_size = self.input_size
        
        B, feature_dim, HW = raw_116_features.shape
        assert feature_dim == 116, f"Expected feature_dim=116, got {feature_dim}"
        
        # 动态适应不同的特征图尺寸
        if HW != 24565:
            # 验证当前配置是否匹配HW
            current_positions = sum(h*w for h, w in self.feature_sizes.values())
            if current_positions != HW:
                print(f"[DYNAMIC] 适配特征尺寸: HW={HW}")
                # 尝试基于常见配置重新初始化
                self._adapt_to_hw(HW)
        
        img_h, img_w = image_size
        detections = []
        
        # Extract multi-scale features from raw features
        
        # Process each image in the batch
        for batch_idx in range(B):
            # Get detections for this image
            img_detections = detection_bboxes[batch_idx]
            valid_detections = img_detections[img_detections[:, 4] > confidence_threshold]  # 使用传入的置信度阈值
            
            # Skip images with no valid detections
            
            # Process each 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])
                
                # Calculate center point
                center_x = (bbox[0] + bbox[2]) / 2
                center_y = (bbox[1] + bbox[3]) / 2
                
                # 减少详细输出以提高性能
                # print(f"    Detection {det_idx+1}: center=({center_x:.1f}, {center_y:.1f}), conf={confidence:.3f}")
                
                # Extract features for all 4 scales
                scale_features = {}
                scale_positions = {}
                
                for scale in ['P2', 'P3', 'P4', 'P5']:
                    # Calculate feature map position for this scale
                    feat_pos, feat_coords = self._map_center_to_scale_position(
                        center_x, center_y, scale, img_w, img_h
                    )
                    
                    # Extract 116-dimensional feature from raw features
                    feature = self._extract_raw_feature_at_position(
                        raw_116_features[batch_idx], feat_pos, scale
                    )
                    
                    scale_features[scale] = feature
                    scale_positions[scale] = feat_coords
                
                # Create multi-scale detection features
                multi_features = MultiScaleDetectionFeatures(
                    bbox=bbox,
                    center_point=(float(center_x), float(center_y)),
                    p2_features=scale_features['P2'],
                    p3_features=scale_features['P3'], 
                    p4_features=scale_features['P4'],
                    p5_features=scale_features['P5'],
                    p2_position=scale_positions['P2'],
                    p3_position=scale_positions['P3'],
                    p4_position=scale_positions['P4'],
                    p5_position=scale_positions['P5'],
                    confidence=confidence,
                    class_id=class_id,
                    image_idx=batch_idx
                )
                
                detections.append(multi_features)
        
        return detections
    
    def _map_center_to_scale_position(
        self, 
        center_x: float, 
        center_y: float,
        scale: str,
        img_w: int,
        img_h: int
    ) -> Tuple[int, Tuple[int, int]]:
        """
        Map image center coordinates to feature map position for a specific scale.
        
        Args:
            center_x, center_y: Center coordinates in original image
            scale: Target scale ('P2', 'P3', 'P4', 'P5')
            img_w, img_h: Image dimensions
            
        Returns:
            (flat_position, (feat_x, feat_y)) in feature map
        """
        # Get stride for this scale
        stride = {'P2': 4, 'P3': 8, 'P4': 16, 'P5': 32}[scale]
        
        # Get feature map dimensions
        feat_h, feat_w = self.feature_sizes[scale]
        
        # Map image coordinates to feature coordinates
        # This maintains exact spatial correspondence without rotation/flipping
        feat_x = int(center_x / stride)
        feat_y = int(center_y / stride)
        
        # Clamp to valid range
        feat_x = max(0, min(feat_x, feat_w - 1))
        feat_y = max(0, min(feat_y, feat_h - 1))
        
        # Convert to flat position in this scale's feature map
        flat_position = feat_y * feat_w + feat_x
        
        # Convert to flat position in concatenated tensor
        concat_position = self.scale_offsets[scale] + flat_position
        
        return concat_position, (feat_x, feat_y)
    
    def _extract_raw_feature_at_position(
        self,
        batch_features: torch.Tensor,  # [116, 24565] for one batch
        position: int,
        scale: str
    ) -> torch.Tensor:
        """
        Extract 116-dimensional raw feature at a specific position.
        
        Args:
            batch_features: Features for one image [116, 24565]
            position: Flat position in concatenated tensor
            scale: Scale name (for verification)
            
        Returns:
            Feature tensor [116]
        """
        # Verify position is within expected range for this scale
        start, end = self.scale_ranges[scale]
        assert start <= position < end, f"Position {position} outside {scale} range [{start}:{end})"
        
        # Additional safety check: ensure position is within tensor bounds
        tensor_size = batch_features.shape[1]
        if position >= tensor_size:
            raise AssertionError(f"Position {position} exceeds tensor size {tensor_size} for scale {scale}")
        
        # Extract the 116-dimensional feature
        feature = batch_features[:, position]  # [116]
        
        return feature
    
    def create_synthetic_test_data(self, batch_size: int = 1) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Create synthetic test data for validation.
        
        Returns:
            (raw_features, detection_bboxes)
            raw_features: [B, 116, 24565]
            detection_bboxes: [B, N, 6]
        """
        print("Creating synthetic test data...")
        
        # Create random raw features (116-dimensional)
        raw_features = torch.randn(batch_size, 116, 24565)
        
        # Create synthetic detections at known positions
        detections_list = []
        
        for batch_idx in range(batch_size):
            batch_detections = []
            
            # Create test detections at strategic positions
            test_positions = [
                (100, 100, 'P2'),  # Should map to P2 position
                (200, 200, 'P2'),  # Should map to P2 position  
                (300, 300, 'P3'),  # Should map to P3 position
                (400, 400, 'P4'),  # Should map to P4 position
                (500, 500, 'P5'),  # Should map to P5 position
            ]
            
            for center_x, center_y, target_scale in test_positions:
                # Create a detection box around the center
                box_size = {'P2': 20, 'P3': 30, 'P4': 50, 'P5': 80}[target_scale]
                
                x1 = center_x - box_size // 2
                y1 = center_y - box_size // 2
                x2 = center_x + box_size // 2
                y2 = center_y + box_size // 2
                
                detection = torch.tensor([
                    x1, y1, x2, y2,  # bbox coordinates
                    0.8 + 0.1 * torch.rand(1).item(),  # confidence (0.8-0.9)
                    torch.randint(0, 52, (1,)).item()   # class_id (0-51)
                ], dtype=torch.float32)
                
                batch_detections.append(detection)
                
                print(f"  Synthetic detection: center=({center_x},{center_y}), scale={target_scale}")
            
            detections_list.append(torch.stack(batch_detections))
        
        # Pad to same number of detections per batch
        max_detections = max(len(dets) for dets in detections_list)
        detection_bboxes = torch.zeros(batch_size, max_detections, 6)
        
        for batch_idx, dets in enumerate(detections_list):
            detection_bboxes[batch_idx, :len(dets)] = dets
        
        print(f"Created synthetic data: raw_features={raw_features.shape}, detections={detection_bboxes.shape}")
        
        return raw_features, detection_bboxes

    def extract_features_from_bbox(
        self,
        raw_116_features: Dict[str, torch.Tensor],  # Dictionary with scale keys {'P2': tensor, 'P3': tensor, ...}
        bbox: torch.Tensor,  # [4] - Single bbox [x1, y1, x2, y2]
        center: Tuple[float, float],  # [2] - Center point [center_x, center_y]
        confidence: float,  # Confidence score
        class_id: int,  # Class ID
        image_idx: int,  # Image index in batch
        image_size: Tuple[int, int]  # [H, W]
    ) -> MultiScaleDetectionFeatures:
        """
        为单个指定的边界框提取多尺度特征，主要用于训练时的GT检测框特征提取
        
        这个方法允许直接从任意指定的边界框位置提取特征，而不依赖于YOLO检测结果。
        在训练阶段，可以使用GT检测框的位置和类别信息来提取完全对齐的特征。
        """
        # 从字典中获取设备信息
        if isinstance(raw_116_features, dict):
            device = raw_116_features['P2'].device if 'P2' in raw_116_features else next(iter(raw_116_features.values())).device
        else:
            device = raw_116_features.device
        
        # 处理center坐标
        if isinstance(center, (tuple, list)):
            center_x, center_y = float(center[0]), float(center[1])
        else:
            center_x, center_y = center[0].item(), center[1].item()
        
        # 为每个尺度提取特征
        scale_features = {}
        scale_positions = {}
        
        if isinstance(raw_116_features, dict):
            # 处理字典格式的多尺度特征
            for scale, stride in [('P2', 4), ('P3', 8), ('P4', 16), ('P5', 32)]:
                # 计算在当前尺度下的特征图坐标
                feat_x = int(center_x / stride)
                feat_y = int(center_y / stride)
                
                if scale in raw_116_features:
                    # 获取对应尺度的特征图
                    scale_feat = raw_116_features[scale]  # [B, C, H, W]
                    
                    # 边界检查，确保坐标在有效范围内
                    _, _, feat_h, feat_w = scale_feat.shape
                    feat_x = max(0, min(feat_x, feat_w - 1))
                    feat_y = max(0, min(feat_y, feat_h - 1))
                    
                    # 提取特征向量 [C]
                    feat_vector = scale_feat[image_idx, :, feat_y, feat_x]  # [C]
                    
                    # 标准化为116维（如果需要）
                    if feat_vector.shape[0] != 116:
                        # 如果不是116维，使用线性变换或重复/裁剪
                        if feat_vector.shape[0] > 116:
                            feat_vector = feat_vector[:116]  # 截断
                        else:
                            # 重复填充到116维
                            repeat_times = (116 + feat_vector.shape[0] - 1) // feat_vector.shape[0]
                            feat_vector = feat_vector.repeat(repeat_times)[:116]
                    
                    scale_features[scale.lower()] = feat_vector
                    scale_positions[f"{scale.lower()}_position"] = (feat_x, feat_y)
                else:
                    # 如果该尺度不存在，创建零向量
                    scale_features[scale.lower()] = torch.zeros(116, device=device)
                    scale_positions[f"{scale.lower()}_position"] = (0, 0)
        else:
            # 处理张量格式的原始116维特征（保持原有逻辑）
            for scale, stride in [('P2', 4), ('P3', 8), ('P4', 16), ('P5', 32)]:
                # 计算在当前尺度下的特征图坐标
                feat_x = int(center_x / stride)
                feat_y = int(center_y / stride)
                
                # 边界检查，确保坐标在有效范围内
                h, w = self.feature_sizes[scale]
                feat_x = max(0, min(feat_x, w - 1))
                feat_y = max(0, min(feat_y, h - 1))
                
                # 计算在拼接特征图中的线性索引
                linear_idx = self.get_feature_index(scale, feat_x, feat_y)
                
                # 提取对应的特征向量 [116]
                feat_vector = raw_116_features[image_idx, :, linear_idx]  # [116]
                
                scale_features[scale.lower()] = feat_vector
                scale_positions[f"{scale.lower()}_position"] = (feat_x, feat_y)
        
        # 构建MultiScaleDetectionFeatures对象
        result = MultiScaleDetectionFeatures(
            bbox=torch.cat([bbox, torch.tensor([confidence, float(class_id)], device=device)]),  # [6]
            center_point=(center_x, center_y),
            p2_features=scale_features['p2'],
            p3_features=scale_features['p3'],
            p4_features=scale_features['p4'],
            p5_features=scale_features['p5'],
            p2_position=scale_positions['p2_position'],
            p3_position=scale_positions['p3_position'],
            p4_position=scale_positions['p4_position'],
            p5_position=scale_positions['p5_position'],
            confidence=confidence,
            class_id=class_id,
            image_idx=image_idx
        )
        
        return result


def test_multi_scale_extractor():
    """Test the multi-scale feature extractor."""
    print("=" * 80)
    print("TESTING MULTI-SCALE FEATURE EXTRACTOR")
    print("=" * 80)
    
    # Initialize extractor
    extractor = MultiScaleFeatureExtractor(input_size=(544, 544))
    
    # Create synthetic test data
    raw_features, detection_bboxes = extractor.create_synthetic_test_data(batch_size=1)
    
    print("\n" + "=" * 80)
    print("EXTRACTING MULTI-SCALE FEATURES")
    print("=" * 80)
    
    # Extract features
    multi_features = extractor.extract_multi_scale_features(
        raw_116_features=raw_features,
        detection_bboxes=detection_bboxes
    )
    
    print(f"\n" + "=" * 80)
    print("ANALYZING RESULTS")
    print("=" * 80)
    
    print(f"Total detections processed: {len(multi_features)}")
    
    for i, det_features in enumerate(multi_features):
        print(f"\n--- Detection {i+1} ---")
        print(f"BBox: [{det_features.bbox[0]:.1f}, {det_features.bbox[1]:.1f}, "
              f"{det_features.bbox[2]:.1f}, {det_features.bbox[3]:.1f}]")
        print(f"Center: {det_features.center_point}")
        print(f"Confidence: {det_features.confidence:.3f}")
        print(f"Class ID: {det_features.class_id}")
        
        # Verify feature dimensions
        print(f"Feature dimensions:")
        print(f"  P2: {det_features.p2_features.shape} (pos: {det_features.p2_position})")
        print(f"  P3: {det_features.p3_features.shape} (pos: {det_features.p3_position})")
        print(f"  P4: {det_features.p4_features.shape} (pos: {det_features.p4_position})")
        print(f"  P5: {det_features.p5_features.shape} (pos: {det_features.p5_position})")
        
        # Feature statistics
        concatenated = det_features.concatenated_features
        feature_matrix = det_features.feature_matrix
        
        print(f"Combined features:")
        print(f"  Concatenated: {concatenated.shape} (4×116=464)")
        print(f"  Matrix:       {feature_matrix.shape} (4×116)")
        print(f"  Feature norm: {torch.norm(concatenated):.4f}")
        
        # Scale-specific statistics
        scales = ['P2', 'P3', 'P4', 'P5']
        for j, scale in enumerate(scales):
            scale_feat = feature_matrix[j]
            print(f"  {scale} norm: {torch.norm(scale_feat):.4f}, mean: {torch.mean(scale_feat):.4f}")
    
    # Test reconstruction
    print(f"\n" + "=" * 80)
    print("VALIDATING POSITION MAPPING")
    print("=" * 80)
    
    for i, det_features in enumerate(multi_features[:2]):  # Test first 2 detections
        center_x, center_y = det_features.center_point
        
        print(f"\nDetection {i+1} position mapping validation:")
        print(f"Image center: ({center_x}, {center_y})")
        
        for scale in ['P2', 'P3', 'P4', 'P5']:
            # Get expected position
            stride = {'P2': 4, 'P3': 8, 'P4': 16, 'P5': 32}[scale]
            expected_feat_x = int(center_x / stride)
            expected_feat_y = int(center_y / stride)
            
            # Get actual position
            if scale == 'P2':
                actual_pos = det_features.p2_position
            elif scale == 'P3':
                actual_pos = det_features.p3_position
            elif scale == 'P4':
                actual_pos = det_features.p4_position
            else:  # P5
                actual_pos = det_features.p5_position
            
            actual_feat_x, actual_feat_y = actual_pos
            
            print(f"  {scale} (stride={stride}):")
            print(f"    Expected: ({expected_feat_x}, {expected_feat_y})")
            print(f"    Actual:   ({actual_feat_x}, {actual_feat_y})")
            print(f"    Match:    {'✓' if (actual_feat_x, actual_feat_y) == (expected_feat_x, expected_feat_y) else '✗'}")
    
    print(f"\n✅ Multi-scale feature extraction test completed!")
    print(f"✅ Each detection now has 4×116 = 464 dimensional multi-scale features")
    
    return extractor, multi_features


if __name__ == "__main__":
    extractor, features = test_multi_scale_extractor()