借助 OpenCV 和 PyTorch 库,利用卷积神经网络提取图像边缘特征

此 Python 代码借助 OpenCV 和 PyTorch 库,实现了实时获取摄像头图像,利用卷积神经网络提取图像边缘特征,并将原始图像和提取的边缘特征图像实时显示出来的功能。

代码详细说明

1. 导入必要的库

python

复制代码
import cv2
import torch
import torch.nn as nn
import numpy as np
  • cv2:OpenCV 库,用于摄像头图像的读取、处理和显示。
  • torchtorch.nn:PyTorch 深度学习框架相关库,用于构建和运行卷积神经网络。
  • numpy:用于数值计算和数组操作。
2. 定义卷积神经网络类 EdgeDetector

python

复制代码
# 定义一个简单的卷积神经网络用于边缘特征提取
class EdgeDetector(nn.Module):
    def __init__(self):
        super(EdgeDetector, self).__init__()
        # 定义一个卷积层,使用Sobel算子的近似卷积核来提取边缘
        self.conv = nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1, bias=False)
        # Sobel算子的x方向卷积核
        sobel_kernel = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0)
        self.conv.weight.data = sobel_kernel

    def forward(self, x):
        return self.conv(x)
  • EdgeDetector 类继承自 nn.Module,这是 PyTorch 中所有神经网络模块的基类。
    • __init__ 方法:
      • self.conv = nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1, bias=False):定义一个二维卷积层,输入通道数为 1(灰度图像),输出通道数为 1,卷积核大小为 3x3,步长为 1,填充为 1,不使用偏置。
      • sobel_kernel = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0):创建 Sobel 算子的 x 方向卷积核,并将其转换为 PyTorch 张量,同时增加维度以匹配卷积层权重的形状。
      • self.conv.weight.data = sobel_kernel:将自定义的 Sobel 卷积核赋值给卷积层的权重。
    • forward 方法:定义了网络的前向传播过程,将输入 x 通过卷积层 self.conv 进行计算并返回结果。
3. 初始化卷积神经网络

python

复制代码
# 初始化卷积神经网络
model = EdgeDetector()

创建 EdgeDetector 类的实例 model,用于后续的边缘特征提取。

4. 打开摄像头

python

复制代码
# 打开摄像头
cap = cv2.VideoCapture(0)

if not cap.isOpened():
    print("无法打开摄像头")
    exit()
  • cv2.VideoCapture(0):尝试打开默认摄像头(设备索引为 0)。
  • if not cap.isOpened():检查摄像头是否成功打开,若未成功则打印错误信息并退出程序。
5. 主循环,实时处理图像

python

复制代码
while True:
    # 读取摄像头的一帧图像
    ret, frame = cap.read()

    if not ret:
        print("无法读取帧")
        break

    # 将图像转换为灰度图
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # 将图像转换为PyTorch张量
    img_tensor = torch.from_numpy(gray).float().unsqueeze(0).unsqueeze(0)

    # 通过卷积神经网络提取边缘特征
    with torch.no_grad():
        edges = model(img_tensor)

    # 将输出转换为NumPy数组
    edges_np = edges.squeeze().numpy()

    # 对边缘特征进行归一化处理
    edges_np = (edges_np - edges_np.min()) / (edges_np.max() - edges_np.min()) * 255
    edges_np = edges_np.astype(np.uint8)

    # 显示原始图像和边缘特征图像
    cv2.imshow('Original', frame)
    cv2.imshow('Edges', edges_np)

    # 按 'q' 键退出循环
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break
  • ret, frame = cap.read():从摄像头读取一帧图像,ret 表示是否成功读取,frame 为读取到的图像数据。
  • cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY):将彩色图像转换为灰度图像,因为后续的卷积层输入要求为单通道图像。
  • torch.from_numpy(gray).float().unsqueeze(0).unsqueeze(0):将 NumPy 数组形式的灰度图像转换为 PyTorch 张量,并增加维度以匹配卷积层输入的形状(批量大小为 1,通道数为 1)。
  • with torch.no_grad():在推理过程中不计算梯度,以节省内存和计算资源。
  • edges = model(img_tensor):将图像张量输入到卷积神经网络中进行边缘特征提取。
  • edges.squeeze().numpy():将输出的张量转换为 NumPy 数组,并去除多余的维度。
  • (edges_np - edges_np.min()) / (edges_np.max() - edges_np.min()) * 255:对边缘特征数组进行归一化处理,将其像素值范围映射到 0 - 255 之间。
  • edges_np.astype(np.uint8):将归一化后的数组转换为无符号 8 位整数类型,以便使用 OpenCV 进行显示。
  • cv2.imshow('Original', frame)cv2.imshow('Edges', edges_np):分别显示原始图像和提取的边缘特征图像。
  • cv2.waitKey(1) & 0xFF == ord('q'):等待 1 毫秒,检查是否按下了 'q' 键,若按下则退出循环。
6. 释放资源

python

复制代码
# 释放摄像头并关闭所有窗口
cap.release()
cv2.destroyAllWindows()
  • cap.release():释放摄像头资源。
  • cv2.destroyAllWindows():关闭所有由 OpenCV 创建的窗口。

总结

此代码通过结合 OpenCV 和 PyTorch,实现了一个简单的实时图像边缘特征提取系统。利用自定义的 Sobel 卷积核的卷积神经网络对摄像头捕获的图像进行处理,最终将原始图像和提取的边缘特征图像实时显示出来。

完整代码

python 复制代码
import cv2
import torch
import torch.nn as nn
import numpy as np

# 定义一个简单的卷积神经网络用于边缘特征提取
class EdgeDetector(nn.Module):
    def __init__(self):
        super(EdgeDetector, self).__init__()
        # 定义一个卷积层,使用Sobel算子的近似卷积核来提取边缘
        self.conv = nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1, bias=False)
        # Sobel算子的x方向卷积核
        sobel_kernel = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0)
        self.conv.weight.data = sobel_kernel

    def forward(self, x):
        return self.conv(x)

# 初始化卷积神经网络
model = EdgeDetector()

# 打开摄像头
cap = cv2.VideoCapture(0)

if not cap.isOpened():
    print("无法打开摄像头")
    exit()

while True:
    # 读取摄像头的一帧图像
    ret, frame = cap.read()

    if not ret:
        print("无法读取帧")
        break

    # 将图像转换为灰度图
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # 将图像转换为PyTorch张量
    img_tensor = torch.from_numpy(gray).float().unsqueeze(0).unsqueeze(0)

    # 通过卷积神经网络提取边缘特征
    with torch.no_grad():
        edges = model(img_tensor)

    # 将输出转换为NumPy数组
    edges_np = edges.squeeze().numpy()

    # 对边缘特征进行归一化处理
    edges_np = (edges_np - edges_np.min()) / (edges_np.max() - edges_np.min()) * 255
    edges_np = edges_np.astype(np.uint8)

    # 显示原始图像和边缘特征图像
    cv2.imshow('Original', frame)
    cv2.imshow('Edges', edges_np)

    # 按 'q' 键退出循环
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放摄像头并关闭所有窗口
cap.release()
cv2.destroyAllWindows()
    
相关推荐
青松@FasterAI33 分钟前
【程序员 NLP 入门】词嵌入 - 上下文中的窗口大小是什么意思? (★小白必会版★)
人工智能·自然语言处理
AIGC大时代1 小时前
高效使用DeepSeek对“情境+ 对象 +问题“型课题进行开题!
数据库·人工智能·算法·aigc·智能写作·deepseek
硅谷秋水1 小时前
GAIA-2:用于自动驾驶的可控多视图生成世界模型
人工智能·机器学习·自动驾驶
偶尔微微一笑1 小时前
AI网络渗透kali应用(gptshell)
linux·人工智能·python·自然语言处理·编辑器
深度之眼1 小时前
2025时间序列都有哪些创新点可做——总结篇
人工智能·深度学习·机器学习·时间序列
晓数2 小时前
【硬核干货】JetBrains AI Assistant 干货笔记
人工智能·笔记·jetbrains·ai assistant
jndingxin2 小时前
OpenCV 图形API(60)颜色空间转换-----将图像从 YUV 色彩空间转换为 RGB 色彩空间函数YUV2RGB()
人工智能·opencv·计算机视觉
Sherlock Ma2 小时前
PDFMathTranslate:基于LLM的PDF文档翻译及双语对照的工具【使用教程】
人工智能·pytorch·语言模型·pdf·大模型·机器翻译·deepseek
知舟不叙2 小时前
OpenCV中的SIFT特征提取
人工智能·opencv·计算机视觉
kadog3 小时前
PubMed PDF下载 cloudpmc-viewer-pow逆向
前端·javascript·人工智能·爬虫·pdf