voc數據集轉換成coco數據集

image

前言

作爲本系列第一篇文章,我分享一個模型訓練過程中常用到的工具,voc數據集轉coco數據集。
在我做一些算法學習的時候,需要將voc數據集轉coco放到yolo當中訓練,但是在網上找了很多個都不是很好用,要不是會報錯,要不是根本不能跑起來。爲了節省在學習算法小夥伴的時間,我分享我在工作常常用的voc轉coco的腳本。

voc 格式分析

爲了能夠更好理解腳本,首先對voc數據集的格式做一個簡單分析。
voc 全稱 The PASCAL Visual Object Classes,它由Visual Object Classes(可視對象類)和挑戰(Challenge)等競賽項目開發, 開始於2005年,結束於2012年最後一屆 。
VOC數據集包含許多不同類型的圖像,每個圖像都標註了一些可視對象,如人,汽車,狗等。這些標註包括每個對象的位置,大小和類別等信息。
常見的voc數據集是voc2007 和voc 2012,當然在模型訓練過程肯定都會自己標註數據集,導出爲voc格式。
image

voc 數據集的格式:

<annotation>
  <folder>17</folder> # 圖片所處文件夾
  <filename>77258.bmp</filename> # 圖片名
  <path>~/frcnn-image/61/ADAS/image/frcnn-image/17/77258.bmp</path>
  <source>  #圖片來源相關信息
    <database>Unknown</database>  
  </source>
  <size> #圖片尺寸
    <width>640</width>
    <height>480</height>
    <depth>3</depth>
  </size>
  <segmented>0</segmented>  #是否有分割label
  <object> 包含的物體
    <name>car</name>  #物體類別
    <pose>Unspecified</pose>  #物體的姿態
    <truncated>0</truncated>  #物體是否被部分遮擋(>15%)
    <difficult>0</difficult>  #是否爲難以辨識的物體, 主要指要結體背景才能判斷出類別的物體。雖有標註, 但一般忽略這類物體
    <bndbox>  #物體的bound box
      <xmin>2</xmin>     #左
      <ymin>156</ymin>   #上
      <xmax>111</xmax>   #右
      <ymax>259</ymax>   #下
    </bndbox>
  </object>
</annotation>

重要的信息包括:filename, size, object 等。除此之外,還有一個主要注意的點就是標註的座標,xmin,ymin,xmax,ymax是標註的四個角,分別代表:

  • xmin: 左上角x軸座標
  • ymin:左上角y周座標
  • xmax: 右下角x軸座標
  • ymax:右下角y軸座標

coco 格式分析

COCO的 全稱是Common Objects in COntext,是微軟團隊提供的一個可以用來進行圖像識別的數據集。MS COCO數 據集中的圖像分爲訓練、驗證和測試集。

假設有以下兩個圖像文件:

  • image1.jpg
  • image2.jpg

coco格式數據集:annotations.json

{
    "images": [
        {
            "id": 1,
            "file_name": "image1.jpg",
            "width": 640,
            "height": 480
        },
        {
            "id": 2,
            "file_name": "image2.jpg",
            "width": 800,
            "height": 600
        }
    ],
    "annotations": [
        {
            "id": 1,
            "image_id": 1,
            "category_id": 1,
            "bbox": [50, 50, 100, 100],
            "area": 10000,
            "segmentation": [ 
                [
                    50, 50, 50, 150, 150, 50
                ]
            ],
            "iscrowd": 0
        },
        {
            "id": 2,
            "image_id": 2,
            "category_id": 2,
            "bbox": [150, 200, 200, 150],
            "area": 30000,
            "segmentation": [
                [
                    150, 200, 150, 350, 350, 200
                ]
            ],
            "iscrowd": 0
        }
    ],
    "categories": [
        {
            "id": 1,
            "name": "cat",
            "supercategory": "animal"
        },
        {
            "id": 2,
            "name": "dog",
            "supercategory": "animal"
        }
    ]
}

coco 數據集字段解析

coco 數據集是一個json文件,一共包括5個部分。

{
    "info": info,               # 數據集的基本信息
    "licenses": [license],      # 許可證
    "images": [image],          #  圖片信息,名字和寬高
    "annotations": [annotation],  # 標註信息 
    "categories": [category]    # 標籤信息
}
info{                           # 數據集信息描述
    "year": int,                # 數據集年份
    "version": str,             # 數據集版本
    "description": str,         # 數據集描述
    "contributor": str,         # 數據集提供者
    "url": str,                 # 數據集下載鏈接
    "date_created": datetime,   # 數據集創建日期
}
license{
    "id": int,
    "name": str,
    "url": str,
} 
image{      # images是一個list,存放所有圖片(dict)信息。image是一個dict,存放單張圖片信息 
    "id": int,                  # 圖片的ID編號(每張圖片ID唯一)
    "width": int,               # 圖片寬
    "height": int,              # 圖片高
    "file_name": str,           # 圖片名字
    "license": int,             # 協議
    "flickr_url": str,          # flickr鏈接地址
    "coco_url": str,            # 網絡連接地址
    "date_captured": datetime,  # 數據集獲取日期
}
annotation{ # annotations是一個list,存放所有標註(dict)信息。annotation是一個dict,存放單個目標標註信息。
    "id": int,                  # 目標對象ID(每個對象ID唯一),每張圖片可能有多個目標
    "image_id": int,            # 對應圖片ID
    "category_id": int,         # 對應類別ID,與categories中的ID對應
    "segmentation": RLE or [polygon],   # 實例分割,對象的邊界點座標[x1,y1,x2,y2,....,xn,yn]
    "area": float,              # 對象區域面積
    "bbox": [xmin,ymin,width,height], # 目標檢測,對象定位邊框[x,y,w,h]
    "iscrowd": 0 or 1,          # 表示是否是人羣
}
categories{                     # 類別描述
    "id": int,                  # 類別對應的ID(0默認爲背景)
    "name": str,                # 子類別名字
    "supercategory": str,       # 主類別名字
}

需要注意的是coco數據集標註的座標。xmin ymin width height和voc有很大差異,分別代表:

  • xmin 左上角x軸座標
  • ymin 左上角y軸座標
  • width 圖片像素寬
  • heidht 圖片像素高

腳本使用

通常在yolo模型檢測訓練時需要的數據集是coco格式或者yolo格式,那麼就需要將voc轉成coco。通常在生成任務中實用的voc數據集的文件和官方數據集格式略有差異。所以首先需要說明,使用該腳本之前需要將voc文件調整成如下格式:
image

數據集包含兩個文件夾,包括gt和images。gt是xml文件保存的目錄,images是圖片保存的目錄。而且xml文件和images同名,只是後綴不一樣。

import os
import json
from xml.etree import ElementTree as ET
from collections import defaultdict


class VocToCoco:

    def __init__(self, voc_gt_dir: str, output_coco_path: str) -> None:
        self.voc_gt_dir = voc_gt_dir
        self.output_coco_path = output_coco_path
        self.categories_count = 1
        self.images = []
        self.categories = {}
        self.annotations = []
        self.data = defaultdict(list)

    # 圖片處理
    def images_handle(self, root: ET.Element, img_id: int) -> None:
        filename = root.find('filename').text.strip()
        width = int(root.find('size').find('width').text)
        height = int(root.find('size').find('height').text)

        self.images.append({
            'id': int(img_id),
            'file_name': filename,
            'height': height,
            'width': width,
        })

    # 標籤轉換
    def categories_handle(self, category: str) -> None:
        if category not in self.categories:
            self.categories[category] = {'id': len(self.categories) + 1, 'name': category}

    # 標註轉換
    def annotations_handle(self, bbox: ET.Element, img_id: int, category: str) -> None:
        x1 = int(bbox.find('xmin').text)
        y1 = int(bbox.find('ymin').text)
        x2 = int(bbox.find('xmax').text)
        y2 = int(bbox.find('ymax').text)

        self.annotations.append({
            'id': self.categories_count,
            'image_id': int(img_id),
            'category_id': self.categories[category].get('id'),
            'bbox': [x1, y1, x2 - x1, y2 - y1],
            'iscrowd': 0
        })
        self.categories_count += 1

    def parse_voc_annotation(self) -> None:

        for img_id, filename in enumerate(os.listdir(self.voc_gt_dir), 1):
            xml_file = os.path.join(self.voc_gt_dir, filename)
            tree = ET.parse(xml_file)
            root = tree.getroot()

            self.images_handle(root, img_id)

            for obj in root.iter('object'):
                category = obj.find('name').text
                self.categories_handle(category)

                bbox = obj.find('bndbox')
                self.annotations_handle(bbox, img_id, category)

        self.data['images'] = self.images
        self.data['categories'] = list(self.categories.values())
        self.data['annotations'] = self.annotations

        with open(self.output_coco_path, 'w') as f:
            json.dump(self.data, f)


if __name__ == "__main__":
    # Example usage
    voc_gt_dir = 'person_驗證集_voc/gt'
    img_dir = 'person_驗證集_voc/images'
    output_coco_path = 'person_驗證集_voc/annocation_coco.json'

    voc2coco = VocToCoco(voc_gt_dir, output_coco_path)
    voc2coco.parse_voc_annotation()

腳本的主要邏輯:

  1. 遍歷所有voc數據集
  2. 獲取所有標籤信息,去重保存在self.categories
  3. 獲取所有圖片元數據,保存在self.images
  4. 獲取所有標註信息,保存在self.annotations
  5. 將以上三個容器保存到字典中,並將字典保存爲一個json文件
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章