python 为暗网生成coco标签

Posted 2021-05-09

import os
import os.path as osp
import sys

sys.path.insert(0, 'coco/PythonAPI')

from pycocotools.coco import COCO


def coco_label(data_type):
    coco = COCO('annotations/instances_%s.json' % data_type)

    if not osp.exists('labels'):
        os.makedirs('labels')

    open('%s.txt' % data_type, 'w')

    img_fname_tmpl = osp.join(data_type, 'COCO_%s_{:012}.jpg' % data_type)
    cat_ids = coco.getCatIds(catNms=CLASSES);
    img_ids = coco.getImgIds(catIds=cat_ids)
    for img_id in img_ids:
        img_fname = img_fname_tmpl.format(img_id)
        with open('%s.txt' % data_type, 'a') as f:
            f.write('%s\n' % osp.realpath(img_fname))

        ann_ids = coco.getAnnIds(imgIds=img_id, catIds=cat_ids, iscrowd=None)
        anns = coco.loadAnns(ann_ids)
        basename = osp.splitext(osp.basename(img_fname))[0]
        with open('%s/%s.txt' % (data_type, basename), 'w') as f:
            for ann in anns:
                x, y, width, height = ann['bbox']
                cls = coco.loadCats(ann['category_id'])[0]['name']
                cls_id = CLASSES.index(cls)
                f.write('%d %d %d %d %d\n' % (cls_id, x, y, width, height))


if __name__ == '__main__':
    CLASSES = ['person']
    data_type = 'train2014'
    for data_type in ['train2014', 'val2014']:
        coco_label(data_type)