使用PyTorch进行物体定位——第一部分(附代码实例)

350 阅读9分钟

随着图像分类模型的普及,深度学习领域得到了关注,剩下的就是历史了。今天,我们仅仅从单一的图像输入就能生成未来的技术。GAN模型生成图像,图像分割模型准确地标记区域,而物体检测模型则检测一切正在发生的事情,比如仅仅从一个普通的闭路电视摄像头中就能识别出繁忙街道上的人。这些成就确实为进一步的研究铺平了道路,而且随着可解释人工智能的到来,基本的建模实践也被接受到医疗领域等严肃的领域。对我来说,图像定位是一个有趣的应用,因为它正好处于图像分类和物体检测之间。每当人们需要在一个项目的检测类别上获得边界框时,人们经常跳到物体检测和YOLO,这很好,但在我看来,如果你只是做分类和定位单一类别,那就太过了。让我们深入了解一下图像定位:概念、PyTorch中的实现以及可能性。

图像分类

图像分类是使用卷积神经网络解决的基本深度学习问题之一,建立在卷积神经网络上的架构,如ResNet、SqueezeNet等,随着时间的推移,此类模型的性能得到了很大程度的提高。根据不同的项目,并不总是要求建立一个超过250MB大小的巨大模型,并给出95%以上的验证精度。有时在开发过程中,比如Lens Studio这样的移动平台,重要的是在尺寸和计算方面得到一个较小的模型,具有合理的准确性,推理时间非常小。在这个较小的模型空间中,我们有像MobileNet、SqueezeNet等模型。

这种架构和用例的多样性导致了分类的改进,也导致了物体的定位。 一个模型不仅可以将图像分类为x或y,而且有能力在图像中定位x或y,并有一个强大的边界框,这种想法成倍地增加了围绕它可以建立的用例。定位的方法与图像分类很相似,所以如果你不熟悉CNN,可能很难跳进去。让我们在进一步了解本地化领域之前,先对其进行一番探讨。

卷积神经网络

如果你是神经网络的新手,请读一下这个

假设你了解深度学习模型是如何在一个基本的全连接层模型上工作的,随着参数的不断增加和全连接层的巨大规模的增加,它们会变得计算量很大。

在这里,如果你看上面的图片,专注于红色和第一个蓝色层,它们之间的计算将采取(不准确,只是为了说明)3*4的操作。现在,这对CPU来说也是很轻的。但是当我们处理图像时,让我们拿一个典型的64X64大小的图像。图像的大小为64X64,每个有3个通道。这就意味着每张图片的RGB通道总共有12288个像素。现在,这听起来很庞大,但当你把处理图像的方式与上面的全连接层相同时,它的计算量就会越来越大。当你把12288个神经元中的每个神经元与下一层的其他神经元相乘时,从计算的角度来看,这真的很昂贵,成本会爆炸性地上升。

为了弥补这一点,我们可以采用CNN:它以一种比许多更重的模型架构更轻的计算方式来完成任务,同时给予更好的准确性。在卷积神经网络中,我们用通道的内核在张量上进行计算,并不断减少张量的大小,同时增加通道的数量,每个通道都有助于学习一些独特的东西,直到我们最终将通道的堆栈平铺到一个全连接层,如下图所示。这只是对CNN的简单解释,我们必须深入研究其架构,以掌握它的每一个方面。

在进行物体定位之前,请关注上面的红点层。那是作为图像分类问题的最后一层,其结果是具有最大值的神经元的指数。这就是预测的类别指数。在物体定位中,我们将不只有一个输出,而是有两个不同的输出,一个是类别,另一个是有4个边界框坐标值的回归输出。

物体定位

与物体检测模型相比,定位的迷人之处在于它的简单性。当然,物体检测模型处理的问题要大得多,在某些架构中必须对初始区域选择层所建议的所有区域进行分类。不过,我们可以像回归问题一样预测坐标的想法还是很有趣的。因此,我们的一般架构,类似于一个图像分类模型,将有额外的1个回归输出(有4个坐标),为获得坐标而平行训练。

定位的数据集

对于物体的定位,我们不仅需要图像和标签,还需要包含物体的边界盒的坐标。我们必须使用图像注释/物体标签工具来准备这样的数据集,而XML是一种流行的文件格式,用于存储每张图像的坐标值。

比较流行的对象标记工具之一是微软的开源对象标记工具,VOTT。它是轻量级和简单的,这使得它非常适合于本地化问题。这里有该项目和文档的链接

并不强制要求你将数据存储为XML,你可以选择任何你能自如处理的文件类型来读取数据,但XML的工作方式很有魅力。我们的输入将是图像、它的标签和与我们要标记的坐标相对应的四个值。在这个项目中,我提取了流行的狗与猫数据集的一部分,并使用VOTT和生成的XML文件创建了边界框。我们将详细讨论如何利用python和一些库的力量读取这些XML文件以获得相应图像的输入坐标。

PyTorch的实现

像往常一样,我们首先要导入包括pandas、CSV和XML在内的库,这将有助于我们对输入的工作:

import numpy as np

from sklearn.model_selection import train_test_split

import matplotlib.pyplot as plt
import cv2
import random
import os

from PIL import Image

import pandas as pd
from xml.dom import minidom
import csv

你可以从这里下载用于演示这个物体定位实现的数据集。在为你的个人项目实施时,改变数据集,用更多的图像和注释来填充它,以获得可靠的模型性能。由于我必须为演示的目的创建边界框,所以我将样本保持在最低限度。

首先,我们将解压数据集以开始工作。数据集的链接在上面给出:

!unzip localization_dataset.zip

在阅读和处理输入之前,让我们分析一下XML文件。了解其结构将有助于你在遇到新的XML时了解如何读取它:

<annotation verified="yes">
    <folder>Annotation</folder>
    <filename>cat.0.jpg</filename>
    <path>Cat-PascalVOC-export/Annotations/cat.0.jpg</path>
    <source>
        <database>Unknown</database>
    </source>
    <size>
        <width>256</width>
        <height>256</height>
        <depth>3</depth>
    </size>
    <segmented>0</segmented>
    <object>
    <name>cat</name>
    <pose>Unspecified</pose>
    <truncated>0</truncated>
    <difficult>0</difficult>
    <bndbox>
        <xmin>55.35820533192091</xmin>
        <ymin>10.992090947210452</ymin>
        <xmax>197.38757944915255</xmax>
        <ymax>171.24521098163842</ymax>
    </bndbox>
</object>
</annotation>

接下来,让我们跳到阅读和处理图像和注释的环节。我们将使用XML minidom来读取xml文件:

def extract_xml_contents(annot_directory, image_dir):
        
        file = minidom.parse(annot_directory)

        # Get the height and width for our image
        height, width = cv2.imread(image_dir).shape[:2]

        # Get the bounding box co-ordinates 
        xmin = file.getElementsByTagName('xmin')
        x1 = float(xmin[0].firstChild.data)

        ymin = file.getElementsByTagName('ymin')
        y1 = float(ymin[0].firstChild.data)

        xmax = file.getElementsByTagName('xmax')
        x2 = float(xmax[0].firstChild.data)

        ymax = file.getElementsByTagName('ymax')
        y2 = float(ymax[0].firstChild.data)

        class_name = file.getElementsByTagName('name')

        if class_name[0].firstChild.data == "cat":
          class_num = 0
        else:
          class_num = 1

        files = file.getElementsByTagName('filename')
        filename = files[0].firstChild.data

        # Return the extracted attributes
        return filename,  width, height, class_num, x1,y1,x2,y2

让我们创建一个num_to_labels字典,从预测中检索标签:

num_to_labels= {0:'cat',1:'dog'}

一旦我们得到所有这些数据,我们可以将其存储为CSV文件,这样我们就不必在每次训练时读取XML文件。我们将使用Pandas数据框架来存储数据,随后我们将把它保存为CSV:

# Function to convert XML files to CSV
def xml_to_csv():

  # List containing all our attributes regarding each image
  xml_list = []

  # We loop our each class and its labels one by one to preprocess and augment 
  image_dir = 'dataset/images'
  annot_dir = 'dataset/annot'

  # Get each file in the image and annotation directory
  mat_files = os.listdir(annot_dir)
  img_files = os.listdir(image_dir)

  # Loop over each of the image and its label
  for mat, image_file in zip(mat_files, img_files):
      
      # Full mat path
      mat_path = os.path.join(annot_dir, mat)

      # Full path Image
      img_path = os.path.join(image_dir, image_file)

      # Get Attributes for each image 
      value = extract_xml_contents(mat_path, img_path)

      # Append the attributes to the mat_list
      xml_list.append(value)

  # Columns for Pandas DataFrame
  column_name = ['filename', 'width', 'height', 'class_num', 'xmin', 'ymin', 
                 'xmax', 'ymax']

  # Create the DataFrame from mat_list
  xml_df = pd.DataFrame(xml_list, columns=column_name)

  # Return the dataframe
  return xml_df

# The Classes we will use for our training
classes_list = sorted(['cat',  'dog'])

# Run the function to convert all the xml files to a Pandas DataFrame
labels_df = xml_to_csv()

# Saving the Pandas DataFrame as CSV File
labels_df.to_csv(('dataset.csv'), index=None)

现在我们得到了CSV中的数据,让我们加载相应的图像和标签来构建数据加载器。我们将对图像进行预处理,包括对像素进行标准化处理。这里需要注意的一件事是我们如何存储坐标。我们通过将其与图像大小(这里是256)相除,将其存储为0到1的范围值,这样模型的输出就可以在以后被缩放为任何图像(我们将在可视化部分看到这个):

def preprocess_dataset():
  # Lists that will contain the whole dataset
  labels = []
  boxes = []
  img_list = []

  h = 256
  w = 256
  image_dir = 'dataset/images'

  with open('dataset.csv') as csvfile:
      rows = csv.reader(csvfile)
      columns = next(iter(rows))
      for row in rows:
        labels.append(int(row[3]))
        #Scaling Coordinates to the range of [0,1] by dividing the coordinate with image size, 256 here.
        arr = [float(row[4])/256,  
               float(row[5])/256,
               float(row[6])/256,
               float(row[7])/256]
        boxes.append(arr)
        img_path = row[0]
        # Read the image
        img  = cv2.imread(os.path.join(image_dir,img_path))

        # Resize all images to a fix size
        image = cv2.resize(img, (256, 256))

        # # Convert the image from BGR to RGB as NasNetMobile was trained on RGB images
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # Normalize the image by dividing it by 255.0 
        image = image.astype("float") / 255.0

        # Append it to the list of images
        img_list.append(image)

  return labels, boxes, img_list

现在这个函数已经准备好了,让我们调用它并加载输入的数据。作为数据预处理的最后一步,我们将对数据进行洗牌:

# All images will resized to 300, 300 
image_size = 256

# Get Augmented images and bounding boxes
labels, boxes, img_list = preprocess_dataset()

# Now we need to shuffle the data, so zip all lists and shuffle
combined_list = list(zip(img_list, boxes, labels))
random.shuffle(combined_list)

# Extract back the contents of each list
img_list, boxes, labels = zip(*combined_list)

到目前为止,我们已经完成了数据预处理,包括从XML文件中读取数据注释,现在我们已经准备好在PyTorch和模型架构设计中构建自定义数据加载器。实施工作将在本文的下一部分继续进行,届时我们将详细讨论上述方法以完成实施工作。谢谢你的阅读 :)

今天就为你的机器学习工作流程增加速度和简单性吧