提取COCO数据集指定类别

原创

已于 2024-05-23 09:29:36 修改

· 920 阅读

18 ·

版权

文章标签：

#python #计算机视觉

于 2024-05-22 15:17:22 首次发布

安装pycocotools库

Win11：

pycocotools库下载链接：https://pypi.tuna.tsinghua.edu.cn/simple/pycocotools-windows/

根据自己的python版本和操作系统下载相应的pycocotools

在anaconda环境下cd到上述whl文件的下载路径，用pip安装

命令为：

pip install "pycocotools_windows-2.0-cp38-cp38-win_amd64.whl"  -i https://pypi.tuna.tsinghua.edu.cn/simple/

安装完成后提示：

提取COCO数据集指定类别

1、提取指定类别并将标注信息转换为XML格式

from pycocotools.coco import COCO
import os
import shutil
from tqdm import tqdm
import skimage.io as io
import matplotlib.pyplot as plt
import cv2
from PIL import Image, ImageDraw
 
# 需要设置的路径
savepath="/path/to/generate/COCO/" 
img_dir=savepath+'images/'
anno_dir=savepath+'annotations/'
datasets_list=['train2017', 'val2017']

#coco有80类，这里写要提取类的名字，以person为例 
classes_names = ['person'] 
#包含所有类别的原coco数据集路径
'''
目录格式如下：
$COCO_PATH
----|annotations
----|train2017
----|val2017
----|test2017
'''
dataDir= '/path/to/coco_orgi/' 
 
headstr = """\
<annotation>
    <folder>VOC</folder>
    <filename>%s</filename>
    <source>
        <database>My Database</database>
        <annotation>COCO</annotation>
        <image>flickr</image>
        <flickrid>NULL</flickrid>
    </source>
    <owner>
        <flickrid>NULL</flickrid>
        <name>company</name>
    </owner>
    <size>
        <width>%d</width>
        <height>%d</height>
        <depth>%d</depth>
    </size>
    <segmented>0</segmented>
"""
objstr = """\
    <object>
        <name>%s</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <b