实战:基于视觉 Transformer 的目标检测

开发 机器视觉
在本文中,我们将详细探讨目标检测,介绍视觉Transformer的强大功能,并通过一个实际项目逐步演示如何使用ViT进行目标检测。

目标检测是计算机视觉中的一项核心任务,推动了从自动驾驶汽车到实时视频监控等技术的发展。它涉及在图像中检测和定位物体,而深度学习的最新进展使这一任务更加准确和高效。推动目标检测的最新创新之一是视觉Transformer(ViT),该模型通过其比传统方法更好地捕捉全局上下文的能力,改变了图像处理的格局。

在本文中,我们将详细探讨目标检测,介绍视觉Transformer的强大功能,并通过一个实际项目逐步演示如何使用ViT进行目标检测。为了使项目更具吸引力,我们将创建一个交互式界面,允许用户上传图像并查看实时目标检测结果。

一、目标检测简介

目标检测是一种用于识别和定位图像或视频中物体的计算机视觉技术。可以将其视为教计算机识别猫、汽车甚至人等物体。通过在图像中绘制这些物体的边界框,我们可以确定每个物体在图像中的位置。

目标检测的重要性:

  • 自动驾驶汽车:实时识别行人、交通信号灯和其他车辆。
  • 监控:检测和跟踪视频流中的可疑活动。
  • 医疗保健:识别医学扫描中的肿瘤和异常。

二、什么是视觉Transformer?

ViT最初由谷歌的研究人员提出。视觉Transformer(ViT)是一种前沿技术,它使用最初为自然语言处理设计的Transformer架构来理解和处理图像。想象一下,将图像分解成小块(如拼图),然后使用智能算法来识别这些小块代表什么以及它们如何组合在一起。

ViT与CNN的区别:

  • CNN:通过卷积层高效识别局部模式(如边缘、纹理)。
  • ViT:从一开始就捕捉全局模式,使其更适合需要理解整个图像上下文的任务。

三、Transformer架构详解

Transformer架构最初是为机器翻译等基于序列的自然语言处理任务设计的,现已被ViT用于视觉数据。以下是其工作原理的分解:

Transformer架构的关键组件:

Vision Transformers 怎么处理图像:

  • Patch Embedding:将图像分割成小块(例如16x16像素),并将每个块线性嵌入为向量。这些块的处理方式类似于NLP任务中的单词。
  • 位置编码:由于Transformer本身不理解空间信息,因此添加位置编码以保留每个块的相对位置。
  • 自注意力机制:该机制允许模型同时关注图像(或块)的不同部分。每个块学习与其他块的关系权重,从而实现对图像的全局理解。
  • 分类:聚合输出通过分类头传递,模型预测图像中存在哪些物体。

ViT相对于CNN的优势:

  • 更好地捕捉全局上下文:ViT可以建模长距离依赖关系,使其更好地理解复杂场景。
  • 适应不同输入尺寸:与CNN需要固定尺寸输入不同,ViT可以适应不同的图像尺寸。

以下是一张比较视觉Transformer(ViT)与卷积神经网络(CNN)架构的图表:

四、项目设置

我们将使用PyTorch和预训练的视觉Transformer设置一个简单的目标检测项目。确保已安装以下必要的库:

pip install torch torchvision matplotlib pillow ipywidgets

这些库的作用:

  • PyTorch:加载并与预训练模型交互。
  • torchvision:预处理图像并应用变换。
  • matplotlib:可视化图像和结果。
  • pillow:图像处理。
  • ipywidgets:创建交互式UI以上传和处理图像。

五、使用ViT逐步实现目标检测

步骤1:加载并显示图像

我们将从加载网络图像并使用matplotlib显示开始。


import torch
from torchvision import transforms
from PIL import Image
import requests
from io import BytesIO
import matplotlib.pyplot as plt

# Load an image from a URL
image_url = "https://upload.wikimedia.org/wikipedia/commons/2/26/YellowLabradorLooking_new.jpg"

# Use a user agent to avoid being blocked by the website
headers = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36"
}

response = requests.get(image_url, headers=headers)

# Check if the request was successful
if response.status_code == 200:
    image = Image.open(BytesIO(response.content))

    # Display the image
    plt.imshow(image)
    plt.axis('off')
    plt.title('Original Image')
    plt.show()

步骤2:预处理图像

ViT期望在将图像输入模型之前对其进行归一化处理。

from torchvision import transforms

preprocess = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

input_tensor = preprocess(image)
input_batch = input_tensor.unsqueeze(0)

步骤3:加载预训练的视觉Transformer模型

现在,我们将从PyTorch的torchvision中加载一个预训练的视觉Transformer模型。

from torchvision.models import vit_b_16

# Step 3: Load a pre-trained Vision Transformer model
model = vit_b_16(pretrained=True)
model.eval()  # Set the model to evaluation mode (no training happening here)

# Forward pass through the model
with torch.no_grad():  # No gradients are needed, as we are only doing inference
    output = model(input_batch)

# Output: This will be a classification result (e.g., ImageNet classes)

步骤4:解释输出

让我们从ImageNet数据集中获取预测的标签。

# Step 4: Interpret the output
from torchvision import models

# Load ImageNet labels for interpretation
imagenet_labels = requests.get("https://raw.githubusercontent.com/anishathalye/imagenet-simple-labels/master/imagenet-simple-labels.json").json()

# Get the index of the highest score
_, predicted_class = torch.max(output, 1)

# Display the predicted class
predicted_label = imagenet_labels[predicted_class.item()]
print(f"Predicted Label: {predicted_label}")

# Visualize the result
plt.imshow(image)
plt.axis('off')
plt.title(f"Predicted: {predicted_label}")
plt.show()
Predicted Label: Labrador Retriever

六、构建交互式图像分类器

我们可以通过创建一个交互式工具使该项目更加用户友好,用户可以在该工具中上传图像或选择样本图像进行分类。为了使项目更具交互性,我们将使用ipywidgets创建一个用户界面,用户可以在其中上传自己的图像或选择样本图像进行目标检测。


import ipywidgets as widgets
from IPython.display import display, HTML, clear_output
from PIL import Image
import torch
import matplotlib.pyplot as plt
from io import BytesIO
import requests
from torchvision import transforms


# Preprocessing for the image
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]),
])

# Create header with glowing title
header = HTML("""
    <div style='text-align:center; margin-bottom:20px;'>
        <h1 style='font-family: Arial, sans-serif; color: #ffe814; font-size: 40px; text-shadow: 0 0 8px #39FF14;'>
            Vision Transformer Object Detection
        </h1>
        <p style='font-family: Arial, sans-serif; color: #ff14b5; font-size:20px'>Upload an image or select a sample image from the cards below</p>
    </div>
""")

# Footer with signature
footer = HTML("""
    <div style='text-align:center; margin-top:20px;'>
        <p style='font-family: Arial, sans-serif; color: #f3f5f2; font-size:25px'>Powered by Vision Transformers | PyTorch | ipywidgets and Create by AI Innovators</p>
    </div>
""")

# Make upload button bigger and centered
upload_widget = widgets.FileUpload(accept='image/*', multiple=False)
upload_widget.layout = widgets.Layout(width='100%', height='50px')
upload_widget.style.button_color = '#007ACC'
upload_widget.style.button_style = 'success'

# Sample images (as cards)
sample_images = [
    ("Dog", "https://upload.wikimedia.org/wikipedia/commons/2/26/YellowLabradorLooking_new.jpg"),
    ("Cat", "https://upload.wikimedia.org/wikipedia/commons/b/b6/Felis_catus-cat_on_snow.jpg"),
    ("Car", "https://upload.wikimedia.org/wikipedia/commons/f/fc/Porsche_911_Carrera_S_%287522427256%29.jpg"),
    ("Bird", "https://upload.wikimedia.org/wikipedia/commons/3/32/House_sparrow04.jpg"),
    ("Laptop", "https://upload.wikimedia.org/wikipedia/commons/c/c9/MSI_Gaming_Laptop_on_wood_floor.jpg")
]

# Function to display and process image
def process_image(image):
    # Clear any previous outputs and predictions
    clear_output(wait=True)

    # Re-display header, upload button, and sample images after clearing
    display(header)
    display(upload_widget)
    display(sample_buttons_box)

    if image.mode == 'RGBA':
        image = image.convert('RGB')

    # Center and display the uploaded image
    plt.imshow(image)
    plt.axis('off')
    plt.title('Uploaded Image')
    plt.show()

    # Preprocess and make prediction
    input_tensor = preprocess(image)
    input_batch = input_tensor.unsqueeze(0)

    with torch.no_grad():
        output = model(input_batch)

    _, predicted_class = torch.max(output, 1)
    predicted_label = imagenet_labels[predicted_class.item()]

    # Display the prediction with space and style
    display(HTML(f"""
        <div style='text-align:center; margin-top:20px; font-size:30px; font-weight:bold; color:#39FF14; text-shadow: 0 0 8px #39FF14;'>
            Predicted: {predicted_label}
        </div>
    """))

    # Display footer after prediction
    display(footer)

# Function triggered by file upload
def on_image_upload(change):
    uploaded_image = Image.open(BytesIO(upload_widget.value[list(upload_widget.value.keys())[0]]['content']))
    process_image(uploaded_image)

# Function to handle sample image selection
def on_sample_image_select(image_url):
    # Define custom headers with a compliant User-Agent
    headers = {
        'User-Agent': 'MyBot/1.0 (your-email@example.com)'  # Replace with your bot's name and contact email
    }

    response = requests.get(image_url, stream=True, headers=headers)  # Added headers
    response.raise_for_status()
    img = Image.open(response.raw)
    process_image(img)

# Add a button for each sample image to the UI (as cards)
sample_image_buttons = [widgets.Button(description=label, layout=widgets.Layout(width='150px', height='150px')) for label, _ in sample_images]

# Link each button to its corresponding image
for button, (_, url) in zip(sample_image_buttons, sample_images):
    button.on_click(lambda b, url=url: on_sample_image_select(url))

# Display buttons horizontally
sample_buttons_box = widgets.HBox(sample_image_buttons, layout=widgets.Layout(justify_content='center'))

# Link the upload widget to the function
upload_widget.observe(on_image_upload, names='value')

# Display the complete UI
display(header)
display(upload_widget)  # Show file upload widget
display(sample_buttons_box)  # Display sample image cards

七、常见问题

Q1:视觉Transformer可以进行微调吗?是的,预训练的视觉Transformer可以在自定义数据集上进行微调,以用于目标检测和分割等任务。

Q2:ViT的计算成本高吗?由于其自注意力机制,ViT的计算成本比CNN更高,尤其是在小型数据集上。

Q3:哪些数据集最适合训练ViT?像ImageNet这样的大型数据集是训练ViT的理想选择,因为与CNN相比,ViT在扩展性方面具有优势。

八、后续步骤

现在你已经学习了视觉Transformer的基础知识,并使用PyTorch实现了目标检测。接下来,你可以尝试在自定义数据集上微调ViT,或者探索其他基于Transformer的模型,例如DETR(Detection Transformer)。

九、结论

视觉Transformer(ViT)代表了计算机视觉领域的一次重大飞跃,为传统的基于CNN的方法提供了一种全新的替代方案。通过利用Transformer架构从一开始就捕捉全局上下文的能力,ViT在大型数据集上展现了令人印象深刻的性能。

责任编辑:赵宁宁 来源: 小白玩转Python
相关推荐

2024-11-21 16:06:02

2023-12-06 09:37:55

模型视觉

2022-12-29 11:57:43

3D框架

2024-11-27 16:06:12

2024-07-04 09:22:24

2023-11-22 13:45:37

计算机视觉数据预处理

2017-10-14 21:24:33

TensorFlow目标检测模型

2024-07-22 13:49:38

YOLOv8目标检测开发

2024-12-16 08:06:42

2023-10-09 09:42:18

自动驾驶模型

2024-06-07 09:22:27

2023-11-22 09:00:00

NLP语言模型LSTM

2023-11-09 23:45:01

Pytorch目标检测

2021-12-01 10:05:12

模型人工智能计算

2022-04-11 09:20:00

模型训练

2024-04-26 10:00:03

自动驾驶模型

2024-07-03 09:39:52

2021-04-15 09:50:41

深度学习编程人工智能

2022-10-21 16:07:10

编码器自然语言模型

2024-08-06 11:01:03

点赞
收藏

51CTO技术栈公众号