maheshp1109/jewelry-visual-search-dinov2

💎 Jewelry Visual Search — DINOv2 + ArcFace

A visual product search model for jewelry, designed for companies like Stuller and other major manufacturing/retail companies. Upload a jewelry image → get matching products from your database + automatic subcategory classification.

🎯 What This Model Does

Capability	Description
Visual Search	Encode jewelry images into 512-dim embeddings. Find matching products via cosine similarity.
Classification	Classify into 5 categories: Bracelet, Earrings, Necklace, Pendant, Ring
Coarse-to-Fine	First classify the category, then search within that category for 99%+ match accuracy

🏗️ Architecture

Input Image (224×224)
    ↓
DINOv2-base (ViT-B/14, 86M params) → 768-dim CLS token
    ↓
Projection Head (768 → 768 → GELU → 512) → L2-normalized embedding
    ↓
├── ArcFace Head (angular margin loss, s=32, m=0.3) → discriminative embeddings
└── Classification Head (512 → 5) → subcategory prediction

Key design choices (backed by research):

• DINOv2 backbone: Self-supervised ViT with strong fine-grained features (DINOv2 paper)
• ArcFace loss: Additive angular margin loss pushes embeddings apart on the hypersphere — critical for fine-grained product retrieval (ArcFace paper, 22k+ citations)
• Joint training: ArcFace (retrieval) + weighted Cross-Entropy (classification) losses
• Progressive unfreezing: Heads-only for 3 epochs → full backbone fine-tuning for 12 more
• Class weighting: Inverse-frequency weights to handle dataset imbalance (Necklace 6.6× more than Ring)

📊 Training Details

Parameter	Value
Backbone	facebook/dinov2-base (86M params, ViT-B/14)
Embedding dim	512
Dataset	bzcasper/ai-tool-pool-jewelry-vision
Train/Val/Test	4,488 / 429 / 213 images
Classes	Bracelet (795), Earrings (822), Necklace (2196), Pendant (345), Ring (330)
Loss	ArcFace (s=32, m=0.3) + 0.1 × weighted CE
Optimizer	AdamW (backbone LR=2e-5, heads LR=1e-3)
Batch	32 × 4 gradient accumulation = 128 effective
Epochs	15 (backbone frozen for first 3)
Augmentation	RandomResizedCrop, HFlip, Rotation±15°, ColorJitter, RandomErasing

🚀 Quick Start — Search by Image

from inference_jewelry_search import JewelrySearchEngine
from PIL import Image

# 1. Load model
engine = JewelrySearchEngine("maheshp1109/jewelry-visual-search-dinov2")

# 2. Build gallery from your product catalog
product_images = [Image.open(f"products/{sku}.jpg") for sku in your_skus]
product_metadata = [{"sku": sku, "name": name, "price": price} for sku, name, price in your_catalog]
engine.build_index(product_images, product_metadata)

# 3. Search with a user-uploaded image
query = Image.open("user_upload.jpg")
results = engine.search(query, top_k=5)

for r in results:
    print(f"  #{r['rank']} Score={r['score']:.3f} SKU={r['metadata']['sku']}")
print(f"  Category: {results[0]['query_category']['category']}")

🏷️ Just Classify

classification = engine.classify(Image.open("ring.jpg"))
# {'category': 'Ring', 'confidence': 0.97, 'all_scores': {'Bracelet': 0.01, ...}}

📦 Production Deployment

# Save gallery index (don't re-encode every time)
engine.save_index("gallery_index.npz")

# Load pre-built index
engine.load_index("gallery_index.npz")

# For millions of products, use FAISS:
import faiss
index = faiss.IndexFlatIP(512)  # cosine similarity (embeddings are L2-normalized)
index.add(engine.gallery_embeddings)
scores, indices = index.search(query_embedding.reshape(1, -1), top_k)

🔬 Two-Stage Search (for 99%+ match accuracy)

For maximum accuracy on exact product matching (like SKU-level retrieval):

# Stage 1: Classify the query → narrow search space
category = engine.classify(query_image)["category"]  # e.g., "Ring"

# Stage 2: Search only within that category
ring_indices = [i for i, m in enumerate(metadata) if m["category"] == "Ring"]
ring_embeddings = gallery_embeddings[ring_indices]

# Much higher precision when searching within category
query_emb = engine.encode_image(query_image)
scores = query_emb @ ring_embeddings.T

📈 Improving Performance

To achieve even higher accuracy for your specific product catalog:

1. Fine-tune on your own data: Replace the dataset with your own product images (same SKU = same label). The more products per category, the better.
2. Higher resolution: Change IMAGE_SIZE to 448 for +3-5% on fine-grained details (stones, settings)
3. More categories: Add subcategories like "Engagement Ring", "Tennis Bracelet", "Stud Earrings"
4. Supplement with fashion data: Use Marqo/fashion200k for pre-fine-tuning warmup
5. Use Marqo-FashionSigLIP: Marqo/marqo-fashionSigLIP as an alternative backbone (already fine-tuned on fashion, 76.7% P@1)

📚 References

• ArcFace: Additive Angular Margin Loss — Core loss function
• DINOv2: Learning Robust Visual Features — Backbone
• GR-Lite / LookBench — Training recipe inspiration
• Marqo GCL — Fashion retrieval SOTA

🏷️ Tags

jewelry visual-search image-retrieval product-search dinov2 arcface metric-learning classification fashion e-commerce

License

Apache-2.0