src/features/hybrid.py

"""
Hybrid feature extractor: CLIP + SAR Adapter + DINOv2.

Combines CLIP global semantics, DINOv2 patch features,
and SAR-specific preprocessing into a single retrieval-ready module.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from typing import Optional
from dataclasses import dataclass
import numpy as np
from torchvision import transforms

from .sar_adapter import SARAdapter


@dataclass
class HybridConfig:
    clip_model: str = "openai/clip-vit-large-patch14"
    dinov2_model: str = "facebook/dinov2-base"
    clip_weight: float = 0.7
    dinov2_weight: float = 0.3
    embed_dim: int = 768
    device: Optional[str] = None


class HybridExtractor(nn.Module):
    """
    Unified hybrid extractor combining CLIP, DINOv2, and SAR adapter.

    Fusion: embedding = w_clip * CLIP(img) + w_dino * DINOv2(img)
    SAR path: SAR -> adapter(2ch->3ch) -> CLIP+DINOv2
    """

    def __init__(self, config: Optional[HybridConfig] = None):
        super().__init__()
        self.config = config or HybridConfig()
        self.device = self.config.device or ("cuda" if torch.cuda.is_available() else "cpu")

        self.sar_adapter = SARAdapter().to(self.device)
        self._clip_model = None
        self._clip_processor = None
        self._dinov2_model = None
        self._loaded = False

        self.dino_transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

        self.fusion_proj = nn.Sequential(
            nn.Linear(self.config.embed_dim, self.config.embed_dim),
            nn.GELU(),
            nn.Linear(self.config.embed_dim, self.config.embed_dim),
        )

    def load(self):
        if self._loaded:
            return
        from transformers import CLIPProcessor, CLIPModel, AutoModel

        print(f"Loading CLIP: {self.config.clip_model} ...")
        self._clip_processor = CLIPProcessor.from_pretrained(self.config.clip_model)
        self._clip_model = CLIPModel.from_pretrained(self.config.clip_model).to(self.device)
        self._clip_model.eval()

        print(f"Loading DINOv2: {self.config.dinov2_model} ...")
        try:
            self._dinov2_model = AutoModel.from_pretrained(self.config.dinov2_model).to(self.device)
            self._dinov2_model.eval()
            self._has_dino = True
            print("DINOv2 loaded")
        except Exception as e:
            self._has_dino = False
            print(f"DINOv2 unavailable: {e}")

        self._loaded = True
        print(f"Hybrid extractor ready on {self.device}")

    @torch.no_grad()
    def _clip_features(self, img: Image.Image) -> np.ndarray:
        inputs = self._clip_processor(images=img, return_tensors="pt").to(self.device)
        out = self._clip_model.vision_model(**inputs)
        pooled = out.last_hidden_state[:, 0, :]
        feat = self._clip_model.visual_projection(pooled).squeeze(0)
        return torch.nn.functional.normalize(feat, dim=-1).cpu().numpy()

    @torch.no_grad()
    def _dinov2_features(self, img: Image.Image) -> Optional[np.ndarray]:
        if not self._has_dino:
            return None
        t = self.dino_transform(img).unsqueeze(0).to(self.device)
        out = self._dinov2_model(t)
        patch_feat = out.last_hidden_state[:, 1:, :].mean(dim=1)
        return torch.nn.functional.normalize(patch_feat.squeeze(0), dim=-1).cpu().numpy()

    def _preprocess_sar(self, img: Image.Image) -> Image.Image:
        arr = np.array(img).astype(np.float32) / 255.0
        t = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
        with torch.no_grad():
            adapted = self.sar_adapter(t)
        arr_out = (adapted.squeeze(0).permute(1, 2, 0).numpy() * 255).clip(0, 255).astype(np.uint8)
        return Image.fromarray(arr_out)

    def extract(
        self,
        img: Image.Image,
        modality: str = "optical",
        normalize: bool = True,
    ) -> np.ndarray:
        if not self._loaded:
            self.load()

        if modality == "sar":
            img = self._preprocess_sar(img)

        clip_feat = self._clip_features(img)
        dino_feat = self._dinov2_features(img)

        if dino_feat is not None:
            w_c, w_d = self.config.clip_weight, self.config.dinov2_weight
            hybrid = w_c * clip_feat + w_d * dino_feat
        else:
            hybrid = clip_feat

        if normalize:
            norm = np.linalg.norm(hybrid)
            if norm > 0:
                hybrid = hybrid / norm

        return hybrid.astype(np.float32)

    def extract_batch(
        self,
        images: list,
        modalities: list = None,
        normalize: bool = True,
    ) -> np.ndarray:
        if modalities is None:
            modalities = ["optical"] * len(images)
        return np.array([
            self.extract(img, mod, normalize)
            for img, mod in zip(images, modalities)
        ])


def create_hybrid_extractor(**kwargs) -> HybridExtractor:
    config = HybridConfig(**kwargs)
    return HybridExtractor(config)


if __name__ == "__main__":
    ext = create_hybrid_extractor()
    ext.load()

    dummy = Image.fromarray(np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8))
    feat = ext.extract(dummy, modality="optical")
    print(f"Feature dim: {feat.shape}, norm: {np.linalg.norm(feat):.4f}")