You Don't Need a Million Labeled Images: Solving Vision Problems with DINOv3

pip install transformers torch torchvision scikit-learn matplotlib einops accelerate timm
  

import torch
from transformers import AutoImageProcessor, AutoModel
from transformers.image_utils import load_image
import numpy as np
# Use a smaller DINOv3 model for faster demos
model_name = "facebook/dino-vits16"
# Load the processor and model
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)
model.eval() # Set the model to evaluation mode

def extract_pooled_features(images):
    """Extracts the global (class token) features for a batch of images."""
    inputs = processor(images=images, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    return outputs.pooler_output.cpu().numpy()

def extract_patch_features(image):
    """Extracts the dense patch features for a single image."""
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    # The first token is the class token, so we skip it
    patch_features = outputs.last_hidden_state[:, 1:, :].squeeze(0).cpu().numpy()
    return patch_features

# Load some sample images from public URLs
image_urls = [
    "http://images.cocodataset.org/val2017/000000039769.jpg",  # Cat image
    "https://images.unsplash.com/photo-1583337130417-2346a1e724c2", # Dog image
    "http://images.cocodataset.org/val2017/000000004134.jpg" # Bird image
]
images = [load_image(url) for url in image_urls]

from sklearn.linear_model import LogisticRegression
print("\n--- 2. Linear Classification (Combined Features) ---")

def extract_combined_features(images):
    """Extracts combined global and average patch features."""
    inputs = processor(images=images, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    
    global_features = outputs.pooler_output
    patch_features = outputs.last_hidden_state[:, 1:, :].mean(dim=1)
    combined = torch.cat([global_features, patch_features], dim=1)
    return combined.cpu().numpy()

# Extract combined features
combined_features = extract_combined_features(images)

# Train a logistic regression classifier
log_reg = LogisticRegression(max_iter=1000)
log_reg.fit(combined_features, labels)

# Predict with the combined features of the dog image
test_combined_features = extract_combined_features([test_image])
prediction = log_reg.predict(test_combined_features)

print(f"Linear classifier prediction for the dog image: Class {prediction[0]}")

from sklearn.metrics.pairwise import cosine_similarity
print("\n--- 3. Image Retrieval and Similarity Search ---")

# Use the cat image as the query
query_image = images[0]
query_features = extract_pooled_features([query_image])

# Calculate cosine similarity between the query and all images
similarities = cosine_similarity(query_features, features).flatten()

# Sort images by similarity
sorted_indices = np.argsort(similarities)[::-1]

print("Image indices sorted by similarity to the cat image:", sorted_indices)
print("Similarity scores:", similarities[sorted_indices])

import torch.nn as nn
print("\n--- 4. Depth Estimation (with a simple head) ---")

class DepthHead(nn.Module):
    def __init__(self, input_dim=384): # For ViT-S/16 model
        super().__init__()
        self.head = nn.Linear(input_dim, 1)
    def forward(self, x):
        return self.head(x)

# Create a dummy depth estimation head
depth_estimator = DepthHead()
depth_estimator.eval()

# "Estimate" depth for each image
with torch.no_grad():
    depth_predictions = depth_estimator(torch.from_numpy(features))

print("Estimated relative depth values:", depth_predictions.numpy().flatten())

print("\n--- 5. Semantic Segmentation (with a simple head) ---")

class SegmentationHead(nn.Module):
    def __init__(self, input_dim=384, num_classes=10):
        super().__init__()
        self.head = nn.Linear(input_dim, num_classes)
    def forward(self, x):
        return self.head(x)

# Create a dummy segmentation head
seg_head = SegmentationHead()
seg_head.eval()

# Get patch features for the cat image
patch_features = extract_patch_features(images[0])

# Get segmentation map
with torch.no_grad():
    segmentation_logits = seg_head(torch.from_numpy(patch_features))
    segmentation_map = torch.argmax(segmentation_logits, dim=-1)

print("Shape of the segmentation map:", segmentation_map.shape)

print("\n--- 6. 3D Keypoint Correspondences ---")

# Use two different images for matching
image1 = images[0]
image2 = load_image("http://images.cocodataset.org/val2017/000000483108.jpg") # Another cat

# Get patch features
features1 = torch.from_numpy(extract_patch_features(image1))
features2 = torch.from_numpy(extract_patch_features(image2))

# Normalize features
features1 = features1 / features1.norm(dim=-1, keepdim=True)
features2 = features2 / features2.norm(dim=-1, keepdim=True)

# Find mutual nearest neighbors
similarity_matrix = torch.matmul(features1, features2.T)
nn1 = torch.argmax(similarity_matrix, dim=1)
nn2 = torch.argmax(similarity_matrix, dim=0)

mutual_matches = []
for i in range(len(nn1)):
    if nn2[nn1[i]] == i:
        mutual_matches.append((i, nn1[i].item()))

print(f"Found {len(mutual_matches)} mutual keypoint matches between the two cat images.")

from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
print("\n--- 7. Unsupervised Object Discovery ---")

# Get patch features for the cat image
patch_features = extract_patch_features(images[0])

# Cluster the patch features
kmeans = KMeans(n_clusters=3, random_state=0, n_init=10)
clusters = kmeans.fit_predict(patch_features)

# Reshape to a 14x14 grid to visualize the clusters
cluster_map = clusters.reshape(14, 14)
plt.imshow(cluster_map, cmap='viridis')
plt.title("Unsupervised Object Clusters")
plt.show()

print("\n--- 8. Video Segmentation Tracking ---")
# Create a dummy video (3 frames of the same cat image)
video_frames = [images[0]] * 3
frame_features = [torch.from_numpy(extract_patch_features(f)) for f in video_frames]

# Track by finding the best match for each patch in the next frame
tracked_patches = []
for i in range(len(frame_features) - 1):
    similarity = torch.matmul(frame_features[i], frame_features[i+1].T)
    best_matches = torch.argmax(similarity, dim=1)
    tracked_patches.append(best_matches)
print("Tracked patch indices between frame 0 and 1:", tracked_patches[0].tolist())

print("\n--- 9. Video Classification ---")
# A simple approach is to average the features of all frames
video_features_pooled = extract_pooled_features(video_frames)
average_video_feature = np.mean(video_features_pooled, axis=0)

# Use this average feature with a classifier like k-NN or Logistic Regression
# prediction = knn.predict([average_video_feature])
print("Shape of the average video feature vector:", average_video_feature.shape)