本流程主要是在mmdet上使用,前五个步骤为数据集制作步骤,后面为在mmdet版本使用步骤

1.在网上下载coco数据集

2.安装COCOAPI(Linux版本)

git clone https://github.com/cocodataset/cocoapi.git $COCOAPI
cd $COCOAPI/PythonAPI
make
python3.7 setup.py install --user #Python版本改为自己对应的

3.提取特定类别(此处借鉴其他博主内容,可以使用此代码提取多个类别,提取类别时,测试集和训练集是根据图片中的类别标定类别id的,因此后续需要进行更改,此处是一个小坑!https://blog.csdn.net/weixin_40922744/article/details/111180137:)

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 = "/data1/bhuang/4/"
img_dir = savepath + 'images/'
anno_dir = savepath + 'annotations/'
datasets_list = ['val2017']

# coco有80类,这里写要提取类的名字,以person为例
classes_names = ['person']
# 提取多个类别示例:classes_names = ['person','car','train']
# 包含所有类别的原coco数据集路径
# 包含所有类别的原coco数据集路径
'''
目录格式如下:
$COCO_PATH
----|annotations
----|train2017
----|val2017
----|test2017
'''
dataDir = '/data1/bhuang/coco/annotations_trainval2017/'

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>
        <bndbox>
            <xmin>%d</xmin>
            <ymin>%d</ymin>
            <xmax>%d</xmax>
            <ymax>%d</ymax>
        </bndbox>
    </object>
"""

tailstr = '''\
</annotation>
'''


# 检查目录是否存在,如果存在,先删除再创建,否则,直接创建
def mkr(path):
    if not os.path.exists(path):
        os.makedirs(path)  # 可以创建多级目录


def id2name(coco):
    classes = dict()
    for cls in coco.dataset['categories']:
        classes[cls['id']] = cls['name']
    return classes


def write_xml(anno_path, head, objs, tail):
    f = open(anno_path, "w")
    f.write(head)
    for obj in objs:
        f.write(objstr % (obj[0], obj[1], obj[2], obj[3], obj[4]))
    f.write(tail)


def save_annotations_and_imgs(coco, dataset, filename, objs):
    # 将图片转为xml,例:COCO_train2017_000000196610.jpg-->COCO_train2017_000000196610.xml
    dst_anno_dir = os.path.join(anno_dir, dataset)
    mkr(dst_anno_dir)
    anno_path = dst_anno_dir + '/' + filename[:-3] + 'xml'
    img_path = dataDir + dataset + '/' + filename
    print("img_path: ", img_path)
    dst_img_dir = os.path.join(img_dir, dataset)
    mkr(dst_img_dir)
    dst_imgpath = dst_img_dir + '/' + filename
    print("dst_imgpath: ", dst_imgpath)
    img = cv2.imread(img_path)
    # if (img.shape[2] == 1):
    #    print(filename + " not a RGB image")
    #   return
    shutil.copy(img_path, dst_imgpath)

    head = headstr % (filename, img.shape[1], img.shape[0], img.shape[2])
    tail = tailstr
    write_xml(anno_path, head, objs, tail)


def showimg(coco, dataset, img, classes, cls_id, show=True):
    global dataDir
    I = Image.open('%s/%s/%s' % (dataDir, dataset, img['file_name']))
    # 通过id,得到注释的信息
    annIds = coco.getAnnIds(imgIds=img['id'], catIds=cls_id, iscrowd=None)
    # print(annIds)
    anns = coco.loadAnns(annIds)
    # print(anns)
    # coco.showAnns(anns)
    objs = []
    for ann in anns:
        class_name = classes[ann['category_id']]
        if class_name in classes_names:
            print(class_name)
            if 'bbox' in ann:
                bbox = ann['bbox']
                xmin = int(bbox[0])
                ymin = int(bbox[1])
                xmax = int(bbox[2] + bbox[0])
                ymax = int(bbox[3] + bbox[1])
                obj = [class_name, xmin, ymin, xmax, ymax]
                objs.append(obj)
                draw = ImageDraw.Draw(I)
                draw.rectangle([xmin, ymin, xmax, ymax])
    if show:
        plt.figure()
        plt.axis('off')
        plt.imshow(I)
        plt.show()

    return objs


for dataset in datasets_list:
    # ./COCO/annotations/instances_train2017.json
    annFile = '{}/annotations/instances_{}.json'.format(dataDir, dataset)

    # 使用COCO API用来初始化注释数据
    coco = COCO(annFile)

    # 获取COCO数据集中的所有类别
    classes = id2name(coco)
    print(classes)
    # [1, 2, 3, 4, 6, 8]
    classes_ids = coco.getCatIds(catNms=classes_names)
    print(classes_ids)
    for cls in classes_names:
        # 获取该类的id
        cls_id = coco.getCatIds(catNms=[cls])
        img_ids = coco.getImgIds(catIds=cls_id)
        print(cls, len(img_ids))
        # imgIds=img_ids[0:10]
        for imgId in tqdm(img_ids):
            img = coco.loadImgs(imgId)[0]
            filename = img['file_name']
            # print(filename)
            objs = showimg(coco, dataset, img, classes, classes_ids, show=False)
            print(objs)
            save_annotations_and_imgs(coco, dataset, filename, objs)

4.提取完成后得到的是xml格式文件,划分数据集后再进行转换。按照6:2:2或者8:2划分数据集(有两个数据集划分代码,第二个无法同时划分图片和标注文件)

# 将图片和标注数据按比例切分为 训练集和测试集

import os
import random
from shutil import copy2

# 原始路径
image_original_path = "/data1/bhuang/3/images/train2017/"
label_original_path = "/data1/bhuang/3/annotations/train2017/"
# 上级目录
# parent_path = os.path.dirname(os.getcwd())
# parent_path = "D:\\AI_Find"
# 训练集路径
# train_image_path = os.path.join(parent_path, "image_data/seed/train/images/")
# train_label_path = os.path.join(parent_path, "image_data/seed/train/labels/")
train_image_path = os.path.join("/data1/bhuang/5/train/images/")
train_label_path = os.path.join("/data1/bhuang/5/train/annotations/")
# 测试集路径
test_image_path = os.path.join("/data1/bhuang/5/test/images/")
test_label_path = os.path.join("/data1/bhuang/5/test/annotations/")


# test_image_path = os.path.join(parent_path, 'image_data/seed/val/images/')
# test_label_path = os.path.join(parent_path, 'image_data/seed/val/labels/')


# 检查文件夹是否存在
def mkdir():
    if not os.path.exists(train_image_path):
        os.makedirs(train_image_path)
    if not os.path.exists(train_label_path):
        os.makedirs(train_label_path)

    if not os.path.exists(test_image_path):
        os.makedirs(test_image_path)
    if not os.path.exists(test_label_path):
        os.makedirs(test_label_path)


def main():
    mkdir()
    # 复制移动图片数据
    all_image = os.listdir(image_original_path)
    for i in range(len(all_image)):
        num = random.randint(1,5)#随机给图片赋值,每五个随机赋值一次,抽取不为2的图片
        if num != 2:
            copy2(os.path.join(image_original_path, all_image[i]), train_image_path)
            train_index.append(i)
        else:
            copy2(os.path.join(image_original_path, all_image[i]), test_image_path)
            val_index.append(i)

    # 复制移动标注数据
    all_label = os.listdir(label_original_path)
    for i in train_index:
            copy2(os.path.join(label_original_path, all_label[i]), train_label_path)
    for i in val_index:
            copy2(os.path.join(label_original_path, all_label[i]), test_label_path)


if __name__ == '__main__':
    train_index = []
    val_index = []
    main()
# *_*coding: utf-8 *_*
# Author --LiMing--

import os
import random
import shutil
import time

def copyFile(fileDir, class_name):
 image_list = os.listdir(fileDir) # 获取图片的原始路径
 image_number = len(image_list)

 train_number = int(image_number * train_rate)
 train_sample = random.sample(image_list, train_number) # 从image_list中随机获取0.8比例的图像.
 test_sample = list(set(image_list) - set(train_sample))
 sample = [train_sample, test_sample]

 # 复制图像到目标文件夹
 for k in range(len(save_dir)):
     if os.path.isdir(save_dir[k] + class_name):
         for name in sample[k]:
             shutil.copy(os.path.join(fileDir, name), os.path.join(save_dir[k] + class_name+'/', name))
     else:
         os.makedirs(save_dir[k] + class_name)
         for name in sample[k]:
             shutil.copy(os.path.join(fileDir, name), os.path.join(save_dir[k] + class_name+'/', name))

if __name__ == '__main__':
 time_start = time.time()

 # 原始数据集路径
 origion_path = '/data1/bhuang/3/'
 #origion_path = '/data1/bhuang/3/images/'

 # 保存路径
 # save_train_dir = '/data1/bhuang/5/images/train/images/'
 # save_test_dir = '/data1/bhuang/5/images/val/images/'
 save_train_dir = '/data1/bhuang/5/'
 save_test_dir = '/data1/bhuang/5/'
 save_dir = [save_train_dir, save_test_dir]

 # 训练集比例
 train_rate = 0.8

 # 数据集类别及数量
 file_list = os.listdir(origion_path)
 num_classes = len(file_list)

 for i in range(num_classes):
     class_name = file_list[i]
     image_Dir = os.path.join(origion_path, class_name)
     copyFile(image_Dir, class_name)
     print('%s划分完毕!' % class_name)

 time_end = time.time()
 print('---------------')
 print('训练集和测试集划分共耗时%s!' % (time_end - time_start))

5.划分完数据集对标注文件转换成json格式(转换完成后特制小型coco数据集制作完成)

import xml.etree.ElementTree as ET
import os
import json

coco = dict()
coco['images'] = []
coco['type'] = 'instances'
coco['annotations'] = []
coco['categories'] = []

category_set = dict()
image_set = set()

category_item_id = 0
image_id = 20180000000
annotation_id = 0


def addCatItem(name):
    global category_item_id
    category_item = dict()
    category_item['supercategory'] = 'none'
    category_item_id += 1
    category_item['id'] = category_item_id
    category_item['name'] = name
    coco['categories'].append(category_item)
    category_set[name] = category_item_id
    return category_item_id


def addImgItem(file_name, size):
    global image_id
    if file_name is None:
        raise Exception('Could not find filename tag in xml file.')
    if size['width'] is None:
        raise Exception('Could not find width tag in xml file.')
    if size['height'] is None:
        raise Exception('Could not find height tag in xml file.')
    image_id += 1
    image_item = dict()
    image_item['id'] = image_id
    image_item['file_name'] = file_name
    image_item['width'] = size['width']
    image_item['height'] = size['height']
    coco['images'].append(image_item)
    image_set.add(file_name)
    return image_id


def addAnnoItem(object_name, image_id, category_id, bbox):
    global annotation_id
    annotation_item = dict()
    annotation_item['segmentation'] = []
    seg = []
    # bbox[] is x,y,w,h
    # left_top
    seg.append(bbox[0])
    seg.append(bbox[1])
    # left_bottom
    seg.append(bbox[0])
    seg.append(bbox[1] + bbox[3])
    # right_bottom
    seg.append(bbox[0] + bbox[2])
    seg.append(bbox[1] + bbox[3])
    # right_top
    seg.append(bbox[0] + bbox[2])
    seg.append(bbox[1])

    annotation_item['segmentation'].append(seg)

    annotation_item['area'] = bbox[2] * bbox[3]
    annotation_item['iscrowd'] = 0
    annotation_item['ignore'] = 0
    annotation_item['image_id'] = image_id
    annotation_item['bbox'] = bbox
    annotation_item['category_id'] = category_id
    annotation_id += 1
    annotation_item['id'] = annotation_id
    coco['annotations'].append(annotation_item)


def parseXmlFiles(xml_path):
    for f in os.listdir(xml_path):
        if not f.endswith('.xml'):
            continue

        bndbox = dict()
        size = dict()
        current_image_id = None
        current_category_id = None
        file_name = None
        size['width'] = None
        size['height'] = None
        size['depth'] = None

        xml_file = os.path.join(xml_path, f)
        print(xml_file)

        tree = ET.parse(xml_file)
        root = tree.getroot()
        if root.tag != 'annotation':
            raise Exception('pascal voc xml root element should be annotation, rather than {}'.format(root.tag))

        # elem is <folder>, <filename>, <size>, <object>
        for elem in root:
            current_parent = elem.tag
            current_sub = None
            object_name = None

            if elem.tag == 'folder':
                continue

            if elem.tag == 'filename':
                file_name = elem.text
                if file_name in category_set:
                    raise Exception('file_name duplicated')

            # add img item only after parse <size> tag
            elif current_image_id is None and file_name is not None and size['width'] is not None:
                if file_name not in image_set:
                    current_image_id = addImgItem(file_name, size)
                    print('add image with {} and {}'.format(file_name, size))
                else:
                    raise Exception('duplicated image: {}'.format(file_name))
                    # subelem is <width>, <height>, <depth>, <name>, <bndbox>
            for subelem in elem:
                bndbox['xmin'] = None
                bndbox['xmax'] = None
                bndbox['ymin'] = None
                bndbox['ymax'] = None

                current_sub = subelem.tag
                if current_parent == 'object' and subelem.tag == 'name':
                    object_name = subelem.text
                    if object_name not in category_set:
                        current_category_id = addCatItem(object_name)
                    else:
                        current_category_id = category_set[object_name]

                elif current_parent == 'size':
                    if size[subelem.tag] is not None:
                        raise Exception('xml structure broken at size tag.')
                    size[subelem.tag] = int(subelem.text)

                # option is <xmin>, <ymin>, <xmax>, <ymax>, when subelem is <bndbox>
                for option in subelem:
                    if current_sub == 'bndbox':
                        if bndbox[option.tag] is not None:
                            raise Exception('xml structure corrupted at bndbox tag.')
                        bndbox[option.tag] = int(option.text)

                # only after parse the <object> tag
                if bndbox['xmin'] is not None:
                    if object_name is None:
                        raise Exception('xml structure broken at bndbox tag')
                    if current_image_id is None:
                        raise Exception('xml structure broken at bndbox tag')
                    if current_category_id is None:
                        raise Exception('xml structure broken at bndbox tag')
                    bbox = []
                    # x
                    bbox.append(bndbox['xmin'])
                    # y
                    bbox.append(bndbox['ymin'])
                    # w
                    bbox.append(bndbox['xmax'] - bndbox['xmin'])
                    # h
                    bbox.append(bndbox['ymax'] - bndbox['ymin'])
                    print('add annotation with {},{},{},{}'.format(object_name, current_image_id, current_category_id,
                                                                   bbox))
                    addAnnoItem(object_name, current_image_id, current_category_id, bbox)


if __name__ == '__main__':
    # 需要自己设定的地址,一个是已生成的是xml文件的父目录;一个是要生成的json文件的目录
    xml_dir = r'/data1/bhuang/5/test/'
    json_dir = r'/data1/bhuang/4/annotations/test/'
    # dataset_lists = ['train2017', 'val2017']
    dataset_lists = ['annotations']
    for dataset in dataset_lists:
        xml_path = os.path.join(xml_dir, dataset)
        json_file = json_dir + '/instances_{}.json'.format(dataset)
        parseXmlFiles(xml_path)
        json.dump(coco, open(json_file, 'w'))

6.开始在mmdetection上使用,以Faster-rcnn为例

首先在configs/_base_/datasets/coco_detection.py中更改数据集路径

data = dict(
    samples_per_gpu=1,
    workers_per_gpu=1,
    train=dict(
        type=dataset_type,
        ann_file=data_root + '5/annotations/train/instances_annotations.json',
        img_prefix=data_root + '5/images/train/images/',
        #CLASSES = CLASSES,
        pipeline=train_pipeline),
    val=dict(
        type=dataset_type,
        ann_file=data_root + '5/annotations/test/instances_annotations.json',
        img_prefix=data_root + '5/images/test/images/',
        #CLASSES = CLASSES,
        pipeline=test_pipeline),
    test=dict(
        type=dataset_type,
        ann_file=data_root + '5/annotations/test/instances_annotations.json',
        img_prefix=data_root + '5/images/test/images/',
        #CLASSES = CLASSES,
        pipeline=test_pipeline))

其次在mmdet/datasets/coco.py中更改coco类别(改成自己的类别)

class CocoDataset(CustomDataset):

    CLASSES = ('person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
                'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
                'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog')

    # CLASSES = ('person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
    #            'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
    #            'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
    #            'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
    #            'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
    #            'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',
    #            'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
    #            'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
    #            'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
    #            'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
    #            'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',
    #            'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
    #            'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
    #            'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush')

接下来在mmdet/core/evaluation/class_names.py中更改coco类别,必须和上面类别一一对应,第二个小坑

def coco_classes():
    return [
         'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train',
         'truck', 'boat', 'traffic_light', 'fire_hydrant', 'stop_sign',
         'parking_meter', 'bench', 'bird', 'cat', 'dog'
    ]
    # return [
    #     'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train',
    #     'truck', 'boat', 'traffic_light', 'fire_hydrant', 'stop_sign',
    #     'parking_meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep',
    #     'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella',
    #     'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard',
    #     'sports_ball', 'kite', 'baseball_bat', 'baseball_glove', 'skateboard',
    #     'surfboard', 'tennis_racket', 'bottle', 'wine_glass', 'cup', 'fork',
    #     'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange',
    #     'broccoli', 'carrot', 'hot_dog', 'pizza', 'donut', 'cake', 'chair',
    #     'couch', 'potted_plant', 'bed', 'dining_table', 'toilet', 'tv',
    #     'laptop', 'mouse', 'remote', 'keyboard', 'cell_phone', 'microwave',
    #     'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase',
    #     'scissors', 'teddy_bear', 'hair_drier', 'toothbrush'
    # ]

最后在faster_rcnn_r50_fpn.py中更改

num_classes=12 #自己coco数据集的类别

最后一个小坑,由于之前提取的数据集的测试集和训练集json文件中类别id不同,需要将test中的类别id改为train中的类别id,上面涉及到的类别顺序必须相同一一对应。

Logo

更多推荐