Crater Intelligence v1 - Mars Crater Embedding Model

Model Description

This model generates 256-dimensional embeddings for Mars crater images, trained using multi-crop contrastive learning with hard negative mining. It's designed for instance-level crater retrieval and identification tasks.

Architecture

Backbone: Vision Transformer (ViT-B/16) pretrained with Mars MAE
Input: Single-channel grayscale crater images (224×224)
Output: 256-dimensional normalized embeddings
Training: Multi-crop contrastive learning (DINO/SwAV style)

Key Features

Instance-level understanding: Distinguishes individual craters even when visually similar
Part-to-whole matching: Recognizes partial crater views (rims, quadrants)
Scale invariance: Robust to different crater sizes and zoom levels
Mars-specific: Pretrained on Mars imagery for optimal performance

Installation

pip install torch torchvision timm huggingface_hub

Usage

Quick Start

import torch
from huggingface_hub import hf_hub_download

# Download model
model_path = hf_hub_download(
    repo_id="jfang/crater-intelli-v1-vit-b-256-0820",
    filename="pytorch_model.bin"
)

# Load model (simple version)
import timm
import torch.nn as nn

class CraterEmbedder(nn.Module):
    def __init__(self, model_path):
        super().__init__()
        # Create ViT backbone
        self.backbone = timm.create_model(
            'vit_base_patch16_224',
            in_chans=1,
            num_classes=0,
            global_pool='token'
        )
        # Projection head
        self.proj = nn.Sequential(
            nn.Linear(768, 1024),
            nn.GELU(),
            nn.Linear(1024, 256)
        )
        # Load weights
        state_dict = torch.load(model_path, map_location='cpu')
        self.load_state_dict(state_dict)
        
    def forward(self, x):
        # x: [B, 1, 224, 224]
        features = self.backbone(x)
        embeddings = self.proj(features)
        # L2 normalize
        return torch.nn.functional.normalize(embeddings, p=2, dim=-1)

# Initialize model
model = CraterEmbedder(model_path)
model.eval()

# Process crater image
from PIL import Image
import torchvision.transforms as T

transform = T.Compose([
    T.Grayscale(num_output_channels=1),
    T.Resize((224, 224)),
    T.ToTensor(),
    T.Normalize(mean=[0.5], std=[0.25])
])

# Load your crater image
img = Image.open("crater.jpg")
img_tensor = transform(img).unsqueeze(0)  # [1, 1, 224, 224]

# Get embedding
with torch.no_grad():
    embedding = model(img_tensor)  # [1, 256]
    
print(f"Embedding shape: {embedding.shape}")
print(f"Embedding norm: {embedding.norm():.3f}")  # Should be ~1.0

Advanced Usage - Crater Retrieval

import numpy as np
import torch
import torch.nn.functional as F
from typing import List, Tuple

class CraterRetriever:
    def __init__(self, model):
        self.model = model
        self.model.eval()
        self.gallery_embeddings = None
        self.gallery_ids = None
        
    def build_gallery(self, images: List[torch.Tensor], crater_ids: List[str]):
        """Build gallery of crater embeddings."""
        embeddings = []
        with torch.no_grad():
            for img_batch in torch.split(torch.stack(images), 32):
                emb = self.model(img_batch)
                embeddings.append(emb)
        
        self.gallery_embeddings = torch.cat(embeddings, dim=0)
        self.gallery_ids = crater_ids
        
    def retrieve(self, query_image: torch.Tensor, k: int = 10) -> List[Tuple[str, float]]:
        """Retrieve k most similar craters."""
        with torch.no_grad():
            query_emb = self.model(query_image.unsqueeze(0))
            
        # Compute cosine similarities
        similarities = F.cosine_similarity(
            query_emb.unsqueeze(1),
            self.gallery_embeddings.unsqueeze(0),
            dim=2
        ).squeeze(0)
        
        # Get top-k
        topk_sims, topk_indices = similarities.topk(k)
        
        results = []
        for sim, idx in zip(topk_sims, topk_indices):
            results.append((self.gallery_ids[idx], sim.item()))
            
        return results

# Example usage
retriever = CraterRetriever(model)

# Build gallery from your crater database
gallery_images = [...]  # List of preprocessed crater tensors
gallery_ids = [...]     # List of crater IDs

retriever.build_gallery(gallery_images, gallery_ids)

# Query with a new crater
query = transform(Image.open("query_crater.jpg")).unsqueeze(0)
results = retriever.retrieve(query, k=5)

for crater_id, similarity in results:
    print(f"Crater {crater_id}: {similarity:.3f}")

Batch Processing

def process_crater_batch(model, image_paths: List[str], batch_size: int = 32):
    """Process multiple crater images efficiently."""
    
    embeddings = []
    
    for i in range(0, len(image_paths), batch_size):
        batch_paths = image_paths[i:i+batch_size]
        batch_tensors = []
        
        for path in batch_paths:
            img = Image.open(path)
            img_tensor = transform(img)
            batch_tensors.append(img_tensor)
        
        batch = torch.stack(batch_tensors)
        
        with torch.no_grad():
            batch_embeddings = model(batch)
            embeddings.append(batch_embeddings)
    
    return torch.cat(embeddings, dim=0)

# Process large crater catalog
crater_paths = ["crater1.jpg", "crater2.jpg", ...]
all_embeddings = process_crater_batch(model, crater_paths)

Input Requirements

Format: Single-channel grayscale images
Size: 224×224 pixels (will be resized if different)
Normalization: Mean=0.5, Std=0.25
Data type: Float32 tensors

Performance

Retrieval Metrics (Validation Set)

Whole crater R@1: 95%+
Partial views:
- Rim crops: ~35%
- Quadrant crops: ~25%
- Offset crops: ~40%
- Zoom crops: ~90%

Training Details

Method: Multi-crop contrastive learning
Loss: Supervised Contrastive (SupCon)
Temperature: Cosine annealing 0.1 → 0.04
Batch size: 64 × (2 global + 6 local views)
Optimizer: AdamW with discriminative LR
Backbone LR: 5e-5 (10× slower than head)
Projection head LR: 5e-4

Limitations

Mars-specific: Trained on Mars craters, may not generalize to other planets
Resolution: Optimized for 224×224 input, very high-res details may be lost
Single channel: Expects grayscale images, not multi-spectral
Crater-centered: Best performance when crater is roughly centered

Citation

@model{crater_intelligence_v1,
  title={Crater Intelligence v1: Mars Crater Instance Embedding},
  author={Fang, J},
  year={2024},
  publisher={HuggingFace},
  howpublished={\url{https://huggingface.co/jfang/crater-intelli-v1-vit-b-256-0820}}
}

License

Apache 2.0

Acknowledgments

Backbone pretrained with Mars MAE on Mars orbital imagery
Training data from Mars crater catalogs
Contrastive learning approach inspired by DINO/SwAV