Spaces:
Sleeping
Sleeping
| """ | |
| SatCLIP-compatible image encoder using OpenAI CLIP ViT-L/14. | |
| Replaces the custom SatCLIP ViT with OpenAI's CLIP (openai/clip-vit-large-patch14) | |
| which produces actually discriminative embeddings for land-cover retrieval. | |
| Interface preserved: .encode(tensor, normalize=True) -> (N, 768) tensor | |
| """ | |
| import torch | |
| import torch.nn.functional as F | |
| from transformers import CLIPModel, CLIPProcessor | |
| from torchvision import transforms | |
| class SatCLIPEncoder: | |
| """ | |
| Image encoder for satellite image retrieval using OpenAI CLIP ViT-L/14. | |
| Handles multi-channel input by converting to 3-channel RGB internally. | |
| """ | |
| def __init__(self, device: str = None): | |
| self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") | |
| self.embed_dim = 768 | |
| print("Loading OpenAI CLIP ViT-L/14...") | |
| self.model = CLIPModel.from_pretrained( | |
| "openai/clip-vit-large-patch14").to(self.device) | |
| self.processor = CLIPProcessor.from_pretrained( | |
| "openai/clip-vit-large-patch14") | |
| self.model.eval() | |
| # For direct tensor input (bypass processor) | |
| self.normalize = transforms.Normalize( | |
| mean=[0.48145466, 0.4578275, 0.40821073], | |
| std=[0.26862954, 0.26130258, 0.27577711], | |
| ) | |
| def _to_3ch(self, tensor: torch.Tensor) -> torch.Tensor: | |
| """Convert any channel-count tensor to 3 channels for CLIP.""" | |
| n = tensor.shape[1] | |
| if n == 3: | |
| return tensor | |
| if n == 1: | |
| return tensor.repeat(1, 3, 1, 1) | |
| if n == 2: | |
| return tensor.repeat(1, 3, 1, 1)[:, :3] | |
| # n >= 3: take first 3 channels | |
| return tensor[:, :3] | |
| def encode(self, image_tensor: torch.Tensor, | |
| normalize: bool = True) -> torch.Tensor: | |
| """ | |
| Encode image tensor to embedding. | |
| Args: | |
| image_tensor: (N, C, 224, 224) tensor, values in [0, 1] | |
| normalize: L2-normalize output | |
| Returns: | |
| (N, 768) embedding tensor | |
| """ | |
| # Convert to 3 channels | |
| x = self._to_3ch(image_tensor.to(self.device)) | |
| # Apply CLIP normalization (ImageNet stats) | |
| x = self.normalize(x) | |
| # Pass through CLIP vision encoder | |
| vision_outputs = self.model.vision_model(x) | |
| features = vision_outputs.pooler_output | |
| # Apply the visual projection | |
| features = self.model.visual_projection(features) | |
| if normalize: | |
| features = F.normalize(features, dim=-1) | |
| return features | |
| # -- Self-check -- | |
| if __name__ == "__main__": | |
| print("Testing SatCLIPEncoder (CLIP backend)...") | |
| encoder = SatCLIPEncoder() | |
| print(f"Embed dim: {encoder.embed_dim}") | |
| dummy = torch.randn(2, 3, 224, 224) | |
| emb = encoder.encode(dummy) | |
| print(f"Output shape: {emb.shape}") | |
| print(f"L2 norm: {emb.norm(dim=-1).tolist()}") | |
| # Test multi-channel handling | |
| dummy_1ch = torch.randn(2, 1, 224, 224) | |
| emb_1ch = encoder.encode(dummy_1ch) | |
| print(f"1ch -> 768: {emb_1ch.shape}, norm={emb_1ch.norm(dim=-1).tolist()}") | |
| dummy_13ch = torch.randn(2, 13, 224, 224) | |
| emb_13ch = encoder.encode(dummy_13ch) | |
| print(f"13ch -> 768: {emb_13ch.shape}, norm={emb_13ch.norm(dim=-1).tolist()}") | |
| print("OK") | |