在MS-COCO数据集中提取边界框和类别标签

时间:2017-10-24 12:08:39

标签: python dictionary machine-learning data-processing mscoco

我正在使用MS-COCO数据集,我想提取边界框以及与背包(类别ID:27)和笔记本电脑对应的图像的标签(类别ID:73)类别,并将它们存储到不同的文本文件中以便稍后训练基于神经网络的模型。

我已经提取了与上述两个类别相对应的图像,并在单独的文件夹中创建了空注释文件,其中我希望将注释与标签一起存储(注释文件的格式如下: label xywh 其中w和h表示检测到的类别的宽度和高度。我建立在COCO-API(准确地说是coco.py)上来提取图像并创建空文本注释文件。

以下是我在coco.py之上编写的主要功能:

if __name__ == "__main__":
    littleCo = COCO('/home/r.bohare/coco_data/annotations/instances_train2014.json')
    #id_laptop = littleCo.getCatIds('laptop')

"""Extracting image ids corresponding to backpack and laptop images."""
    bag_img_ids = littleCo.getImgIds(catIds=[27])
    laptop_img_ids = littleCo.getImgIds(catIds=[73])
    #print "IDs of bag images:", bag_img_ids
    #print "IDs of laptop imgs:", laptop_img_ids

"""Extracting annotation ids corresponding to backpack and laptop images."""
    bag_ann_ids = littleCo.getAnnIds(catIds=[27])
    laptop_ann_ids = littleCo.getAnnIds(catIds=[73])
    #print "Annotation IDs of bags:", bag_ann_ids
    #print "Annotation IDs of laptops:", laptop_ann_ids

"""Extracting image names corresponding to bag and laptop categories."""
    bag_imgs = littleCo.loadImgs(ids=bag_img_ids)
    laptop_imgs = littleCo.loadImgs(ids=laptop_img_ids)
    #print "Bag images:", bag_imgs
    #print "Laptop images:", laptop_imgs

    bag_img_names = [image['file_name'] for image in bag_imgs]
    laptop_img_names = [image['file_name'] for image in laptop_imgs]
    print "Bag Images:", len(bag_img_names), bag_img_names[:5]
    print "Laptop Images:", len(laptop_img_names), laptop_img_names[:5]

"""Extracting annotations corresponding to bag and laptop images."""
    bag_ann = littleCo.loadAnns(ids=bag_ann_ids)
    laptop_ann = littleCo.loadAnns(ids=laptop_ann_ids)
    bag_bbox = [ann['bbox'] for ann in bag_ann]
    laptop_bbox = [ann['bbox'] for ann in laptop_ann]
    print "Bags' bounding boxes:", len(bag_ann), bag_bbox[:5]
    print "Laptops' bounding boxes:", len(laptop_bbox), laptop_bbox[:5]

"""Saving files corresponding to bags and laptop category in a directory."""
    import shutil
    #path_to_imgs = "/export/work/Data Pool/coco_data/train2014/"
    #path_to_subset_imgs = "/export/work/Data Pool/coco_subset_data/"
    path_to_ann = "/export/work/Data Pool/coco_subset_data/annotations/"
    dirs_list = [("/export/work/Data Pool/coco_data/train2014/", "/export/work/Data Pool/coco_subset_data/")]

    for source_folder, destination_folder in dirs_list:
        for img in bag_img_names:
            shutil.copy(source_folder + img, destination_folder + img)
        print "Bag images copied!"
        for img in laptop_img_names:
            shutil.copy(source_folder + img, destination_folder + img)
        print "Laptop images copied!" 

"""Creating empty files for annotation."""
    for f in os.listdir("/export/work/Data Pool/coco_subset_data/images/"):
        if f.endswith('.jpg'):
            open(os.path.join(path_to_ann, f.replace('.jpg', '.txt')), 'w+').close()
    print "Done creating empty annotation files." 

我在这里仅提供了main函数,因为其余的代码是COCO-API中的coco.py文件。

我调试了代码,发现有不同的数据结构:

  • cats,一个将类别ID映射到其超类别和类别名称(标签)的字典。
  • imgToAnns,也是一个字典,它将每个图像ID映射到它的分割基础事实,边界框基础事实,类别ID等。从我到目前为止所知道的,我想我需要使用这个字典以某种方式将我在 bag_img_names laptop_img_names 列表中的图像名称映射到他们的标签和边界框,但我无法思考正确的方向,如何访问这个字典(coco.py中没有方法直接返回它)。
  • imgs,另一个字典,提供有关所有图像的元信息,例如图像名称,图像网址,捕获日期等。

最后,我知道这是一个非常具体的问题。如果这属于stackoverflow以外的社区(例如stats.stackexchange.com),请随时告诉我,我将删除它。此外,我可能错过了一些重要信息。如果我能想到它,或者有人问,我会提供它。

我只是Python的初学者,所以如果我可能错过了一些明显的东西,请原谅我。

任何帮助都非常感谢。谢谢。

1 个答案:

答案 0 :(得分:0)

2年过去了。现在,coco.py已经可以做您想做的事情,只需添加一些函数以将注释转换为RLE格式,然后将其映射到图像即可。看看这个cocoapi.