图像相似性检测代码报错ValueError：X与Y矩阵维度不兼容

问题重现 #

训练图像相似性检测模型（判断输入图像与训练集图像是否相似）时，训练环节正常生成features.npy文件，但运行相似性计算代码时触发维度不匹配错误：

ValueError: Incompatible dimension for X and Y matrices: X.shape[1] == 150528 while Y.shape[1] == 512

相似性计算代码 #

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from PIL import Image
import torchvision.transforms as transforms
import torch
import os


# Load the pre-trained features
features = np.load('features.npy')

# Define a function to calculate cosine similarity between an image and features
def calculate_similarity(image, features, transformer):
    # Apply the same transformations used for training data
    image = transformer(image)

    # Extract features from the input image
    image = image.unsqueeze(0)  # Add batch dimension
    
    # Flatten both the image and features
    image_flat = image.view(1, -1)
    features_flat = features.reshape(features.shape[0], -1)

    # Now calculate cosine similarity
    similarity = cosine_similarity(image_flat, features_flat)

    return similarity

# Define a function to check if an image is similar to the training set
def is_similar(image, features, transformer, threshold=0.7):
    similarity = calculate_similarity(image, features, transformer)
    return similarity.max() >= threshold

# Example usage for processing images in a directory
if __name__ == '__main__':
    # Define the transformation pipeline for both training and test data
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    # Directory containing the images
    image_dir = 'data/train'

    # List all files in the directory
    image_files = [os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.endswith('.jpeg')]

    for image_path in image_files:
        image = Image.open(image_path)
        is_similar_image = is_similar(image, features, transform)

        if is_similar_image:
            print(f"The image {image_path} is similar to the training set.")
        else:
            print(f"The image {image_path} is not similar to the training set.")

模型训练代码 #

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision import models
from torch.utils.data import DataLoader, Dataset
import os
from PIL import Image
import ssl
ssl._create_default_https_context = ssl._create_unverified_context

# Define the custom dataset for image loading
class ImageDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.image_filenames = os.listdir(root_dir)

    def __len__(self):
        return len(self.image_filenames)

    def __getitem__(self, idx):
        img_name = os.path.join(self.root_dir, self.image_filenames[idx])
        image = Image.open(img_name)
        if self.transform:
            image = self.transform(image)
        return image

# Set the data directory containing your training images
data_dir = 'data/train'

# Define data transformations for the model
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Create a custom dataset and data loader
dataset = ImageDataset(data_dir, transform=transform)
dataloader = DataLoader(dataset, batch_size=64, shuffle=True)

# Load a pre-trained ResNet model
model = models.resnet18(pretrained=True)
model = nn.Sequential(*list(model.children())[:-1])  # Remove the final classification layer

# Extract features from images using the pre-trained model
model.eval()
features = []
with torch.no_grad():
    for inputs in dataloader:
        inputs = inputs.to('cuda') if torch.cuda.is_available() else inputs
        outputs = model(inputs)
        features.extend(outputs.cpu().numpy())

# Save the extracted features to a file
import numpy as np

features = np.array(features)
np.save('features.npy', features)

问题原因 #

训练代码中，你用移除最后一层的ResNet18提取特征，输出的是512维特征向量（ResNet18的avgpool层输出为(batch_size, 512, 1, 1)，展平后是512维）。但相似性计算代码里，你直接把输入图像的张量展平成了224*224*3=150528维的原始像素向量，没有用同一个特征提取模型处理输入图像，导致两者维度完全不匹配，无法计算余弦相似度。

解决方案 #

修改相似性计算代码，加载训练时使用的特征提取模型，用它来处理输入图像，得到和训练集特征维度一致的512维向量，再计算余弦相似度。

修改后的完整相似性计算代码 #

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from PIL import Image
import torchvision.transforms as transforms
import torch
import torch.nn as nn
from torchvision import models
import os


# Load the pre-trained features
features = np.load('features.npy')
# 加载和训练时一致的特征提取模型
model = models.resnet18(pretrained=True)
model = nn.Sequential(*list(model.children())[:-1])  # 移除最后分类层
model.eval()
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)

# Define a function to calculate cosine similarity between an image and features
def calculate_similarity(image, features, transformer, model, device):
    # Apply the same transformations used for training data
    image = transformer(image).unsqueeze(0).to(device)  # 添加batch维度并移到对应设备
    
    # 用特征提取模型提取输入图像的特征
    with torch.no_grad():
        image_feature = model(image)
    image_feature = image_feature.cpu().numpy().reshape(1, -1)  # 展平为1x512的向量
    
    # 训练集特征已经是512维，无需额外展平（如果保存前没展平，这里可以reshape）
    features_flat = features.reshape(features.shape[0], -1)

    # 计算余弦相似度
    similarity = cosine_similarity(image_feature, features_flat)

    return similarity

# Define a function to check if an image is similar to the training set
def is_similar(image, features, transformer, model, device, threshold=0.7):
    similarity = calculate_similarity(image, features, transformer, model, device)
    return similarity.max() >= threshold

# Example usage for processing images in a directory
if __name__ == '__main__':
    # Define the transformation pipeline for both training and test data
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    # Directory containing the images
    image_dir = 'data/train'

    # List all files in the directory
    image_files = [os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.endswith('.jpeg')]

    for image_path in image_files:
        image = Image.open(image_path)
        is_similar_image = is_similar(image, features, transform, model, device)

        if is_similar_image:
            print(f"图像 {image_path} 与训练集图像相似。")
        else:
            print(f"图像 {image_path} 与训练集图像不相似。")

额外说明 #

• 确保训练和测试时使用完全一致的图像预处理流程，包括尺寸、归一化参数等，否则会影响特征一致性。
• 如果训练时保存的features.npy维度是(num_samples, 512, 1, 1)，可以在加载后先展平为(num_samples, 512)，避免重复reshape操作。

内容的提问来源于stack exchange，提问作者Mohammed Abid Nafi