Spaces:
Sleeping
Sleeping
File size: 3,310 Bytes
f343f06 | 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 | """
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]
@torch.no_grad()
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")
|