torchvision resnet18 计算相似度

import torch import torchvision.models as models import torchvision.transforms as transforms from PIL import Image import time import os import psutil # 需要安装psutil库 import logging from torchvision.models import ResNet18_Weights # 配置日志 logging.basicConfig(filename='photo_similarity.log', level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S') def get_all_file_paths(directory): file_paths = [] # 用于存储文件路径的数组 for root, directories, files in os.walk(directory): for filename in files: filepath = os.path.join(root, filename) file_paths.append(filepath) # 将文件路径添加到列表中 return file_paths def print_resource_usage(start_cpu_percent, start_memory_usage): current_process = psutil.Process(os.getpid()) end_cpu_percent = current_process.cpu_percent() end_memory_usage = current_process.memory_info().rss / (1024 * 1024) # Convert bytes to MB # print(f"CPU使用率从{start_cpu_percent}%增加到{end_cpu_percent}%") logging.info(f"CPU使用率从{start_cpu_percent}%增加到{end_cpu_percent}%") # print(f"内存使用从{start_memory_usage}MB增加到{end_memory_usage}MB") logging.info(f"内存使用从{start_memory_usage}MB增加到{end_memory_usage}MB") # 记录程序开始时间和资源使用情况 start_time = time.time() start_cpu_percent = psutil.cpu_percent() process = psutil.Process(os.getpid()) start_memory_usage = process.memory_info().rss / (1024 * 1024) # Convert bytes to MB # 加载预训练的深度网络模型 model = models.resnet18(weights=ResNet18_Weights.DEFAULT) # 选择你需要的模型 # model = models.resnet34(weights=ResNet34_Weights.DEFAULT) # 选择你需要的模型 # model = models.resnet50(weights=ResNet50_Weights.DEFAULT) # 选择你需要的模型 # model = models.resnet101(weights=ResNet101_Weights.DEFAULT) # 选择你需要的模型 # model = models.resnet152(weights=ResNet152_Weights.DEFAULT) # 选择你需要的模型 # model = models.efficientnet_v2_l(weights=EfficientNet_V2_L_Weights.DEFAULT) # 选择你需要的模型 model.eval() # 设置为评估模式 # 定义图像转换 preprocess = 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]), ]) # 读取并预处理第一张图像 # originImage = 'LSCABN2E6PG182573-08.jpeg' # originImage = 'LVHTG6851R7000484-01-e.jpg' # originImage = 'LFBGE3069FJD46842-01.jpg' # originImage = 'LVHRV3853P5044817-01.jpeg' # originImage = 'LVHTG4818R7000659-01.jpeg' # originImage = '14681.jpg' originImage = 'LVHTG6851R7000484-02-copy.jpg' # originImage = 'LSCABN2E6PG182573-08-e.jpg' image1 = Image.open(originImage) # image1 = Image.open('14681.jpg') image1_tensor = preprocess(image1).unsqueeze(0) # 使用模型提取第一张图像的特征 with torch.no_grad(): image1_features = model(image1_tensor) # 初始化余弦相似度计算 cos = torch.nn.CosineSimilarity(dim=1, eps=1e-6) # 存储相似度和图片路径 similarities = [] # 指定你想要遍历的文件夹路径 # directory = 'old-arr' directory = 'LVHTG6851R7000484_copy' # directory = 'LSCABN2E6PG182573_e' other_images = get_all_file_paths(directory) # 计算与其他图片的相似度 for img_path in other_images: if '.DS_Store' in img_path: continue image2 = Image.open(img_path) image2_tensor = preprocess(image2).unsqueeze(0) with torch.no_grad(): image2_features = model(image2_tensor) similarity = cos(image1_features, image2_features).item() similarities.append((similarity, img_path)) # 按相似度从大到小排序 similarities.sort(reverse=True, key=lambda x: x[0]) # 打印排序后的相似度 for sim, path in similarities: print(f"Cosine Similarity between {originImage} and {path}: {sim}") logging.info(f"Cosine Similarity between {originImage} and {path}: {sim}") # 记录并打印程序结束时间和资源使用情况 end_time = time.time() print(f"程序执行时间：{end_time - start_time}秒") # logging.info(f"程序执行时间：{end_time - start_time}秒") print_resource_usage(start_cpu_percent, start_memory_usage)