| from dataclasses import dataclass |
| from typing import Any, Dict, Optional, Tuple |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from src.models.fusion import StructuralFusionConfig, StructuralVisualFusion |
|
|
|
|
| @dataclass |
| class VisionEncoderConfig: |
| """ |
| Cấu hình cho vision encoder wrapper. |
| |
| Args: |
| backbone_name: Tên backbone. Ưu tiên timm nếu có, ví dụ |
| "vit_base_patch16_224". |
| pretrained: Dùng pretrained weights nếu backend hỗ trợ. |
| image_size: Kích thước ảnh input HxW. |
| patch_size: Patch size của ViT. Với 224/16 -> 14x14 -> 196 tokens. |
| d_model: Hidden dimension output patch tokens. |
| freeze_backbone: Nếu True, freeze ViT backbone để train fusion trước. |
| backend: "auto", "timm", hoặc "torchvision". |
| drop_cls_token: Nếu True, bỏ CLS token để chỉ giữ 196 patch tokens. |
| """ |
|
|
| backbone_name: str = "vit_base_patch16_224" |
| pretrained: bool = True |
| image_size: Tuple[int, int] = (224, 224) |
| patch_size: int = 16 |
| d_model: int = 768 |
| freeze_backbone: bool = True |
| backend: str = "auto" |
| drop_cls_token: bool = True |
|
|
| @property |
| def num_patches(self) -> int: |
| return (self.image_size[0] // self.patch_size) * (self.image_size[1] // self.patch_size) |
|
|
|
|
| class PatchVisionEncoder(nn.Module): |
| """ |
| Wrapper lấy patch tokens từ pretrained ViT. |
| |
| Input: |
| image: [B, 3, H, W] |
| |
| Output: |
| patch_tokens: [B, N, D] |
| |
| Backend: |
| - timm: ưu tiên vì hỗ trợ forward_features tiện. |
| - torchvision: fallback cho vit_b_16 nếu không có timm. |
| """ |
|
|
| def __init__(self, config: VisionEncoderConfig): |
| super().__init__() |
| self.config = config |
| self.backbone, self.backend, inferred_dim = self._build_backbone(config) |
|
|
| self.output_projection = nn.Identity() |
| if inferred_dim != config.d_model: |
| self.output_projection = nn.Linear(inferred_dim, config.d_model) |
|
|
| if config.freeze_backbone: |
| self.freeze_backbone() |
|
|
| def forward(self, image: torch.Tensor) -> torch.Tensor: |
| self._validate_image(image) |
|
|
| if self.backend == "timm": |
| tokens = self._forward_timm(image) |
| elif self.backend == "torchvision": |
| tokens = self._forward_torchvision(image) |
| else: |
| raise RuntimeError(f"Backend không hỗ trợ: {self.backend}") |
|
|
| tokens = self._ensure_patch_tokens(tokens) |
| tokens = self.output_projection(tokens) |
| return tokens |
|
|
| def freeze_backbone(self) -> None: |
| for param in self.backbone.parameters(): |
| param.requires_grad = False |
|
|
| def unfreeze_backbone(self) -> None: |
| for param in self.backbone.parameters(): |
| param.requires_grad = True |
|
|
| def unfreeze_last_blocks(self, num_blocks: int = 2) -> None: |
| """Unfreeze vài block cuối nếu backend expose blocks/encoder.layers.""" |
| self.freeze_backbone() |
| if num_blocks <= 0: |
| return |
|
|
| if hasattr(self.backbone, "blocks"): |
| blocks = self.backbone.blocks |
| for block in blocks[-num_blocks:]: |
| for param in block.parameters(): |
| param.requires_grad = True |
| elif hasattr(self.backbone, "encoder") and hasattr(self.backbone.encoder, "layers"): |
| layers = self.backbone.encoder.layers |
| for layer in list(layers)[-num_blocks:]: |
| for param in layer.parameters(): |
| param.requires_grad = True |
|
|
| @staticmethod |
| def _validate_image(image: torch.Tensor) -> None: |
| if not torch.is_tensor(image): |
| raise TypeError("image phải là torch.Tensor") |
| if image.ndim != 4: |
| raise ValueError(f"image phải có shape [B, 3, H, W], nhận {tuple(image.shape)}") |
| if image.shape[1] != 3: |
| raise ValueError(f"image channel phải là 3 RGB, nhận C={image.shape[1]}") |
|
|
| def _build_backbone(self, config: VisionEncoderConfig): |
| backend = config.backend.lower() |
| if backend not in {"auto", "timm", "torchvision"}: |
| raise ValueError("backend phải là 'auto', 'timm', hoặc 'torchvision'") |
|
|
| if backend in {"auto", "timm"}: |
| try: |
| import timm |
|
|
| backbone = timm.create_model( |
| config.backbone_name, |
| pretrained=config.pretrained, |
| num_classes=0, |
| ) |
| inferred_dim = getattr(backbone, "num_features", config.d_model) |
| return backbone, "timm", int(inferred_dim) |
| except ImportError: |
| if backend == "timm": |
| raise ImportError( |
| "Bạn chọn backend='timm' nhưng chưa cài timm. Cài bằng: pip install timm" |
| ) |
| except Exception: |
| if backend == "timm": |
| raise |
|
|
| try: |
| from torchvision.models import ViT_B_16_Weights, vit_b_16 |
|
|
| weights = ViT_B_16_Weights.DEFAULT if config.pretrained else None |
| backbone = vit_b_16(weights=weights) |
| inferred_dim = backbone.hidden_dim |
| return backbone, "torchvision", int(inferred_dim) |
| except Exception as exc: |
| raise RuntimeError( |
| "Không build được ViT backend. Hãy cài timm bằng `pip install timm` " |
| "hoặc kiểm tra torchvision." |
| ) from exc |
|
|
| def _forward_timm(self, image: torch.Tensor) -> torch.Tensor: |
| features = self.backbone.forward_features(image) |
|
|
| if isinstance(features, dict): |
| for key in ("x", "tokens", "last_hidden_state"): |
| if key in features: |
| features = features[key] |
| break |
| else: |
| raise RuntimeError(f"Không nhận diện được output dict từ timm: {features.keys()}") |
|
|
| return features |
|
|
| def _forward_torchvision(self, image: torch.Tensor) -> torch.Tensor: |
| |
| x = self.backbone._process_input(image) |
| batch_size = x.shape[0] |
| cls_token = self.backbone.class_token.expand(batch_size, -1, -1) |
| x = torch.cat([cls_token, x], dim=1) |
| x = self.backbone.encoder(x) |
| return x |
|
|
| def _ensure_patch_tokens(self, tokens: torch.Tensor) -> torch.Tensor: |
| if tokens.ndim != 3: |
| raise ValueError(f"ViT output phải có shape [B, N, D], nhận {tuple(tokens.shape)}") |
|
|
| expected_patches = self.config.num_patches |
| num_tokens = tokens.shape[1] |
|
|
| if self.config.drop_cls_token and num_tokens == expected_patches + 1: |
| tokens = tokens[:, 1:, :] |
| elif num_tokens == expected_patches: |
| pass |
| elif num_tokens > expected_patches: |
| tokens = tokens[:, -expected_patches:, :] |
| else: |
| raise ValueError( |
| f"Số token ViT={num_tokens} nhỏ hơn expected patches={expected_patches}. " |
| "Kiểm tra image_size/patch_size/backbone." |
| ) |
|
|
| return tokens |
|
|
|
|
| class GlobalFeatureProjector(nn.Module): |
| """ |
| Chiếu global structural features [B, G] thành global token [B, 1, D]. |
| """ |
|
|
| def __init__( |
| self, |
| global_feature_dim: int, |
| d_model: int, |
| hidden_dim: Optional[int] = None, |
| dropout: float = 0.10, |
| use_layer_norm: bool = True, |
| ): |
| super().__init__() |
| if hidden_dim is None: |
| hidden_dim = max(d_model // 2, global_feature_dim * 4) |
|
|
| self.input_norm = nn.LayerNorm(global_feature_dim) if use_layer_norm else nn.Identity() |
| self.projector = nn.Sequential( |
| nn.Linear(global_feature_dim, hidden_dim), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(hidden_dim, d_model), |
| nn.Dropout(dropout), |
| ) |
| self.output_norm = nn.LayerNorm(d_model) if use_layer_norm else nn.Identity() |
|
|
| def forward(self, global_features: torch.Tensor) -> torch.Tensor: |
| if not torch.is_tensor(global_features): |
| global_features = torch.as_tensor(global_features) |
| if global_features.ndim != 2: |
| raise ValueError( |
| f"global_features phải có shape [B, G], nhận {tuple(global_features.shape)}" |
| ) |
|
|
| token = self.projector(self.input_norm(global_features.float())) |
| token = self.output_norm(token) |
| return token.unsqueeze(1) |
|
|
|
|
| class StructuralVisionEncoder(nn.Module): |
| """ |
| End-to-end vision side: |
| |
| image -> ViT -> patch_tokens |
| patch_tokens + prior_mask + topo_features -> StructuralVisualFusion |
| global_features -> global structural token |
| concat -> visual_context |
| |
| Output: |
| visual_context: [B, 197, D] nếu dùng global token |
| fused_tokens: [B, 196, D] |
| patch_tokens: [B, 196, D] |
| """ |
|
|
| def __init__( |
| self, |
| vision_config: VisionEncoderConfig, |
| topo_feature_dim: int = 12, |
| global_feature_dim: int = 8, |
| use_global_token: bool = True, |
| fusion_kwargs: Optional[Dict[str, Any]] = None, |
| ): |
| super().__init__() |
| fusion_kwargs = fusion_kwargs or {} |
|
|
| self.vision_config = vision_config |
| self.use_global_token = use_global_token |
| self.patch_encoder = PatchVisionEncoder(vision_config) |
|
|
| fusion_config = StructuralFusionConfig( |
| d_model=vision_config.d_model, |
| topo_feature_dim=topo_feature_dim, |
| **fusion_kwargs, |
| ) |
| self.fusion = StructuralVisualFusion(fusion_config) |
|
|
| self.global_projector = None |
| if use_global_token: |
| self.global_projector = GlobalFeatureProjector( |
| global_feature_dim=global_feature_dim, |
| d_model=vision_config.d_model, |
| dropout=fusion_config.dropout, |
| use_layer_norm=fusion_config.use_layer_norm, |
| ) |
|
|
| def forward( |
| self, |
| image: torch.Tensor, |
| prior_mask: Optional[torch.Tensor] = None, |
| topo_features: Optional[torch.Tensor] = None, |
| global_features: Optional[torch.Tensor] = None, |
| return_diagnostics: bool = True, |
| ): |
| patch_tokens = self.patch_encoder(image) |
|
|
| fusion_out = self.fusion( |
| visual_tokens=patch_tokens, |
| prior_mask=prior_mask, |
| topo_features=topo_features, |
| return_diagnostics=True, |
| ) |
| fused_tokens = fusion_out["fused_tokens"] |
|
|
| global_token = None |
| visual_context = fused_tokens |
| if self.use_global_token: |
| if global_features is None: |
| raise ValueError("use_global_token=True nhưng global_features=None") |
| global_features = global_features.to(device=fused_tokens.device, dtype=fused_tokens.dtype) |
| global_token = self.global_projector(global_features) |
| visual_context = torch.cat([global_token, fused_tokens], dim=1) |
|
|
| if not return_diagnostics: |
| return visual_context |
|
|
| return { |
| "visual_context": visual_context, |
| "patch_tokens": patch_tokens, |
| "fused_tokens": fused_tokens, |
| "global_token": global_token, |
| **fusion_out, |
| } |
|
|
|
|
| def build_structural_vision_encoder( |
| d_model: int = 768, |
| topo_feature_dim: int = 12, |
| global_feature_dim: int = 8, |
| backbone_name: str = "vit_base_patch16_224", |
| pretrained: bool = True, |
| freeze_backbone: bool = True, |
| backend: str = "auto", |
| use_global_token: bool = True, |
| **fusion_kwargs, |
| ) -> StructuralVisionEncoder: |
| vision_config = VisionEncoderConfig( |
| backbone_name=backbone_name, |
| pretrained=pretrained, |
| d_model=d_model, |
| freeze_backbone=freeze_backbone, |
| backend=backend, |
| ) |
| return StructuralVisionEncoder( |
| vision_config=vision_config, |
| topo_feature_dim=topo_feature_dim, |
| global_feature_dim=global_feature_dim, |
| use_global_token=use_global_token, |
| fusion_kwargs=fusion_kwargs, |
| ) |
|
|
