CV (1) пользовательский набор данных

компьютерное зрение
CV (1) пользовательский набор данных

В этой статье в качестве примера для иллюстрации используется набор данных PASCAL VOC2012. (ссылка для скачивания:PASCAL VOC2012)

Пользовательский набор данных Pytorch см. документацию:TorchVision Object Detection Finetuning Tutorial

В этой статье будет настроен набор данных на основе PASCAL VOC.VOCDataset, и случайным образом выберите пять изображений, чтобы преобразовать их соответствующие аннотации в прямоугольные блоки и нарисовать их на изображениях.

image

Подробный код этой статьи см.pytorch-tutorial/01-common/custom_dataset at main · simo-an/pytorch-tutorial (github.com)

определить некоторые служебные классы

Определить данные категории, всего 20 целевых категорий.

class_dict = {
    "aeroplane": 1,
    "bicycle": 2,
    "bird": 3,
    "boat": 4,
    "bottle": 5,
    "bus": 6,
    "car": 7,
    "cat": 8,
    "chair": 9,
    "cow": 10,
    "diningtable": 11,
    "dog": 12,
    "horse": 13,
    "motorbike": 14,
    "person": 15,
    "pottedplant": 16,
    "sheep": 17,
    "sofa": 18,
    "train": 19,
    "tvmonitor": 20
}

преобразовать xml в json

def parse_xml_to_dict(xml):
    if len(xml) == 0:
        return {xml.tag: xml.text}

    result = {}
    for child in xml:
        child_result = parse_xml_to_dict(child)
        if child.tag != 'object':
            result[child.tag] = child_result[child.tag]
        else:
            if child.tag not in result:
                result[child.tag] = []
            result[child.tag].append(child_result[child.tag])

    return {xml.tag: result}

Нарисуйте прямоугольную рамку на картинке (код ссылки:vision/utils.py at main · pytorch/vision (github.com))

def draw_bounding_boxes(
    image,
    boxes: torch.Tensor,
    labels: torch.Tensor,
) -> torch.Tensor:
    img_to_draw = Image.fromarray(image)
    img_boxes = boxes.to(torch.int64).tolist()
    draw = ImageDraw.Draw(img_to_draw)

    class_map = [k for k, v in class_dict.items()]

    for i, bbox in enumerate(img_boxes):
        draw.rectangle(bbox, width=2, outline='red')
        margin = 2
        draw.text((bbox[0] + margin, bbox[1] + margin),
                  class_map[labels[i] - 1], fill='red')

    return np.array(img_to_draw)

Создание пользовательских наборов данных

Некоторые основные библиотеки, которые необходимо импортировать

import torch
import utils
from torch.utils.data import Dataset
from PIL import Image
from os import path
from lxml import etree

Как того требует документация, вVOCDatasetреализовать три метода в__len__,__getitem__,а такжеget_height_and_width.

Инициализировать класс VOCDataset

Конструктор определяется следующим образом

'''
voc_root: voc 数据集的根目录
year: 哪一个年份的数据集
transforms: 数据预处理
text_name: train.txt or val.txt 该txt文件在数据集的 VOCdevkit\VOC2012\ImageSets\Main 文件夹下
'''
def __init__(self, voc_root, year='2012', transforms=None, text_name='train.txt'):

В конструкторе мы в основном выполняем следующие три функции

  1. установить путь к изображениюimage_rootи пометить путьanno_root
  2. Задайте список путей ко всем файлам аннотаций для обучаемого образца на этот раз.xml_list
  3. Установите информацию о целевой категории для обнаруженияclass_dict

установить путь к изображениюimage_rootи пометить путьanno_root 

        # 设置数据集、图片、标注的根目录
        self.root = path.join(voc_root, 'VOCdevkit', f'VOC{year}')
        self.image_root = path.join(self.root, 'JPEGImages')
        self.anno_root = path.join(self.root, 'Annotations')

Задайте список путей ко всем файлам аннотаций для обучаемого образца на этот раз.xml_list 

        # 根据 text_name 拿到对应的标注xml文件路径
        text_path = path.join(self.root, 'ImageSets','Main', text_name)
        # 读取txt文件的每一行并生成xml标注文件路径存放在xml_list中
        with open(text_path) as file_reader:
            self.xml_list = [
                path.join(self.anno_root, f'{line.strip()}.xml')
                for line in file_reader.readlines() if len(line.strip()) > 0
            ]

Установите информацию о целевой категории для обнаруженияclass_dict 

        self.class_dict = utils.class_dict

Обычно используйте 0, чтобы указать, что текущая категория является фоновой.

Получить все образцы

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

Количество семплов — это длина отмеченного списка файлов.

Получить указанный образец по индексу

Функция определяется следующим образом

    def __getitem__(self, idx):

Входными данными является значение индекса выборки, а диапазон его значений равен0 ~ len(xml_list)

Получение указанного образца необходимо разделить на следующие два этапа.

  1. получить фотографии
  2. Получить информацию об изображении (информация о метке, индекс, площадь области и т. д.)

получить фотографии

Во-первых, нам нужно получить соответствующую аннотационную информацию по индексу и преобразовать ее вjsonФорматопределить получениеjsonформатированныйannotationМетоды

    def get_annotation(self, idx):
        xml_path = self.xml_list[idx]
        assert path.exists(xml_path), f'file {xml_path} not found'

        xml_reader = open(xml_path)
        xml_text = xml_reader.read()
        xml = etree.fromstring(xml_text)
        annotation = utils.parse_xml_to_dict(xml)['annotation']

Получатьannotation

        annotation = self.get_annotation(idx)

Тогда мы можем получить отannotationполучить имя файла и получить файл

        image_path = path.join(self.image_root, annotation['filename'])
        image = Image.open(image_path)

Получить информацию о картинке

Декларировать всю информацию, которую необходимо получить

        # 生成 target
        target = {
            'boxes': [], # 标注的左上、右下坐标(xmin, ymin, xmax, ymax)
            'labels': [],# 标注类别
            'image_id': [], # 图片索引
            'area': [], # 含有目标区域的面积 (xmax-xmin) * (ymax-ymin)
            'iscrowd': [], # 是不是一堆密集的东西在一起
        }

облегчить всеobject


        for obj in annotation['object']:
            bndbox = obj['bndbox']
            xmin = float(bndbox['xmin'])
            ymin = float(bndbox['ymin'])
            xmax = float(bndbox['xmax'])
            ymax = float(bndbox['ymax'])
            target['boxes'].append([xmin, ymin, xmax, ymax]) # 设置有目标的坐标信息
            target['labels'].append(self.class_dict[obj['name']]) # 获取对应的标签
            target['area'].append((xmax - xmin) * (ymax - ymin)) # 计算面积

            # 使用 difficult(当前目标是否难以识别) 字段来设置 iscrowd
            if 'difficult' in obj:
                target['iscrowd'].append(int(obj['difficult']))
            else:
                target['iscrowd'].append(0)

преобразовать всю информацию вTensor

        # Convert to tensor
        target['boxes'] = torch.as_tensor(target['boxes'])
        target['labels'] = torch.as_tensor(target['labels'])
        target['iscrowd'] = torch.as_tensor(target['iscrowd'])
        target['area'] = torch.as_tensor(target['area'])
        target['image_id'] = torch.tensor([idx])

Если установлен препроцессор данных, он вызывается перед возвратом данных.

        if self.transforms is not None:
            image = self.transforms(image)

Возврат фотографий и соответствующей информации

        return image, target

Получить ширину и высоту текущего изображения на основе индекса

Информация на этикетке содержит информацию о ширине и высоте изображения, поэтому ее можно легко получить.

    def get_height_and_width(self, idx):
        annotation = annotation = self.get_annotation(idx)
        # 从 annotation 中取出宽高并返回
        width = int(annotation['size']['width'])
        height = int(annotation['size']['height'])

        return height, width

Мы завершили определение набора данных выше, и мы будем использовать пример кода для использования этого набора данных ниже.

Используйте пользовательский набор данных и нарисуйте поле выноски

Импорт некоторых основных библиотек

import os
import random
import utils
import main
import matplotlib.pyplot as plt
import numpy as np

определениеtransformer, преобразовать данные вTensor

data_transform = ts.Compose([ts.ToTensor()])

так какToTensorДанные будут стандартизированы, для простоты кода они здесь использоваться не будут.

Получите набор данных и нарисуйте целевое поле и категорию

train_data_set = VOCDataset(os.getcwd(), '2012', None, 'train.txt')

for index in random.sample(range(0, len(train_data_set)), k=5):
    image, target = train_data_set[index]
    image = utils.draw_bounding_boxes(
        np.array(image),
        target['boxes'],
        target['labels'],
    )
    plt.imshow(image)
    plt.show()

На этом весь процесс завершен!

запустить и проверить

image

Видно, что результат операции правильный!

Категории