基于CNN与视觉Transformer(ViT)的图像处理 | AI生成和翻译

Home 2025.03

以下是使用 卷积神经网络（CNNs） 和 视觉变换器（ViT） 在 Python 中通过 PyTorch 处理图像的代码示例。

1. 使用 CNN（卷积神经网络）进行图像处理

CNN 广泛用于图像分类、目标检测和特征提取。

示例：使用预训练的 CNN（ResNet）

import torch
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image

# 加载预训练的 ResNet 模型
model = models.resnet18(pretrained=True)
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]),
])

# 加载并预处理图像
image = Image.open("example.jpg")  # 替换为你的图像路径
input_tensor = preprocess(image)
input_batch = input_tensor.unsqueeze(0)  # 添加批次维度

# 如果可用则移至 GPU
if torch.cuda.is_available():
    input_batch = input_batch.to('cuda')
    model.to('cuda')

# 提取特征（在最终分类层之前）
with torch.no_grad():
    features = model(input_batch)

print("特征向量形状:", features.shape)  # 例如：torch.Size([1, 1000])

说明：

  1. ResNet18 是在 ImageNet 上预训练的 CNN 架构。
  2. 图像经过预处理（调整大小、标准化）。
  3. 模型将图像转换为特征向量（例如，ResNet18 为 1000 维）。

2. 使用视觉变换器（ViT）进行图像处理

ViT 将图像视为补丁序列，并使用自注意力机制（类似于 NLP）。

示例：使用预训练的 ViT（Hugging Face）

from transformers import ViTFeatureExtractor, ViTModel
from PIL import Image
import torch

# 加载预训练的视觉变换器（ViT）
model_name = "google/vit-base-patch16-224-in21k"
feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
model = ViTModel.from_pretrained(model_name)

# 加载图像
image = Image.open("example.jpg")  # 替换为你的图像路径

# 预处理图像（转换为补丁）
inputs = feature_extractor(images=image, return_tensors="pt")

# 提取特征（CLS 标记或补丁嵌入）
with torch.no_grad():
    outputs = model(**inputs)

# 获取特征向量（CLS 标记）
features = outputs.last_hidden_state[:, 0, :]  # 形状：[1, 768]

print("特征向量形状:", features.shape)  # 例如：torch.Size([1, 768])

说明：

  1. ViT 将图像分割为 16x16 补丁，并像 NLP 中的标记一样处理它们。
  2. CLS 标记（第一个标记）代表整个图像的特征向量。
  3. 输出是一个 768 维向量（对于 vit-base）。

3. CNN 与 ViT 特征提取对比

模型 方法 特征向量大小 库
CNN (ResNet18) 卷积层 + 池化 1000（ImageNet 类别数） torchvision
ViT (Google ViT-Base) 补丁嵌入 + 变换器 768（隐藏维度） transformers

4. 图像特征向量的应用

# 示例：计算两张图像的相似度
from sklearn.metrics.pairwise import cosine_similarity

# 假设 features1 和 features2 是提取的向量
similarity = cosine_similarity(features1, features2)
print("相似度得分:", similarity[0][0])

关键要点

  1. CNN（例如 ResNet）使用卷积层提取层次化特征。
  2. ViT（例如 Google 的 ViT）将图像视为补丁序列并使用自注意力。
  3. 两者都将图像转换为特征向量以用于下游任务。

需要关于从头训练自定义 CNN/ViT 的示例吗？告诉我吧！🚀

Back Donate