python使用RegNetY 模型加载图片特征

最新推荐文章于 2025-08-04 22:31:11 发布

susu1083018911

最新推荐文章于 2025-08-04 22:31:11 发布

阅读量113

点赞数 1

CC 4.0 BY-SA版权

文章标签： python

本文链接：https://blog.csdn.net/susu1083018911/article/details/149603885

import torch
import timm
from torchvision import transforms
from PIL import Image
import torch.nn.functional as F
from scipy.spatial.distance import cosine
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

# 1. 加载预训练的 RegNetY 模型
model = timm.create_model('regnety_1280', pretrained=True)  # 选择不同模型 'regnety_800mf', 'regnety_1_6gf'
model.eval()  # 设置为评估模式

# 移除分类头，仅保留特征提取部分
model.head = torch.nn.Identity()

# 2. 图像预处理函数
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 调整图像大小
    transforms.ToTensor(),  # 转换为 Tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
])

# 3. 加载和预处理图像
def preprocess_image(image_path):
    img = Image.open(image_path).convert('RGB')  # 打开并转换为 RGB
    return transform(img).unsqueeze(0)  # 添加批量维度

# 4. 提取特征
def extract_features(image_path):
    img_tensor = preprocess_image(image_path)  # 预处理图像
    with torch.no_grad():  # 禁用梯度计算
        features = model(img_tensor)  # 获取特征
    return features.squeeze()  # 返回去除批量维度后的特征

# 5. 计算余弦相似度
def cosine_similarity(features1, features2):
    # 将特征展平为 1D 向量
    features1 = features1.flatten()
    features2 = features2.flatten()
    # 计算余弦相似度
    return 1 - cosine(features1.numpy(), features2.numpy())  # 余弦相似度越接近1表示越相似


# 6. 计算欧氏距离
def euclidean_distance(features1, features2):
    return torch.norm(features1 - features2).item()

# 7. 加载两张图片并提取特征
image1_path = '1.jpg'
image2_path = '2.jpg'

features1 = extract_features(image1_path)
features2 = extract_features(image2_path)

# 8. 计算相似度
similarity_cosine = cosine_similarity(features1, features2)
similarity_euclidean = euclidean_distance(features1, features2)

# 9. 打印结果
print(f"Cosine Similarity: {similarity_cosine:.4f}")
print(f"Euclidean Distance: {similarity_euclidean:.4f}")