colab上比较DINO

!git clone https://github.com/facebookresearch/dino.git
!pip install timm
import torch
import timm
from PIL import Image
from torchvision import transforms

# 加载模型
model = torch.hub.load('facebookresearch/dino:main', 'dino_vits8')
model.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 定义图像预处理
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 图像特征提取函数
def extract_features(img_path):
    image = Image.open(img_path).convert('RGB')
    x = transform(image).unsqueeze(0).to(device)
    with torch.no_grad():
        features = model(x)
    return features[0]

# 加载图像并提取特征
features1 = extract_features("img1.jpg")
features2 = extract_features("img2.jpg")

# 计算特征之间的余弦相似度
cos = torch.nn.CosineSimilarity(dim=1, eps=1e-6)
similarity = cos(features1.unsqueeze(0), features2.unsqueeze(0))
print("Similarity:", similarity.item())

 

posted @ 2024-07-01 16:11  Anonytt  阅读(9)  评论(0编辑  收藏  举报