首页 > > 网络编程 > 其它 >

coco数据集标注图转为二值图python（附代码）

2019-07-24 09:16:16来源：博客园阅读 ()

coco数据集大概有8w张以上的图片，而且每幅图都有精确的边缘mask标注。

后面后分享一个labelme标注的json或xml格式转二值图的源码（以备以后使用）

而我现在在研究显著性目标检测，需要的是边缘mask的二值图像。搜了很久，并没有人做过这种工作，只能得到如下的掩膜图

而我需要的图像为二值图，如下

说下我的过程并附上代码：

首先，coco数据集将所有的8w多张图片标注信息整合到一个json文件中，所以我们需要将单张图片标注信息json文件提取出来，以下是批量提取脚本。

注：需要改动地方 1）第6行：将json_file改为原coco数据集json文件的地址（coco/annotations/xxxxx.json）

2) 第13行：设置需要提取的图片数量我是提取82000张

3）第37行：设置存储json文件的目录需要新建空文件夹我是放在./coco_single_object下

4）第33-35行：可选将图片的名称写入data.txt中不需要的话可以注释掉

 1 # -*- coding:utf-8 -*-
 2 from __future__ import print_function
 3 import json
 4 
 5 #json文件的地址 需要手动设置
 6 json_file='../pycocotools/instances_train2014.json' # # Object Instance 类型的标注
 7 # person_keypoints_val2017.json  # Object Keypoint 类型的标注格式
 8 # captions_val2017.json  # Image Caption的标注格式
 9 
10 data=json.load(open(json_file,'r'))
11 
12 #设置需要提取的图片数量 我设置提取82000张
13 for i in range(82000):
14     data_2 = {}
15     data_2['info'] = data['info']
16     data_2['licenses'] = data['licenses']
17     data_2['images'] = [data['images'][i]]  # 只提取第一张图片
18     data_2['categories'] = data['categories']
19     annotation = []
20 
21     # 通过imgID 找到其所有对象
22     imgID = data_2['images'][0]['id']
23     for ann in data['annotations']:
24         if ann['image_id'] == imgID:
25             annotation.append(ann)
26 
27     data_2['annotations'] = annotation
28     # 保存到新的JSON文件，便于查看数据特点
29     #img_file 获取图片名称
30     img_file=data_2['images'][0]['file_name']
31     img_first=img_file.split(".")[0]
32     #将提取出的图片写入data.txt文件中并换行 (optional)
33     # with open('./coco_single_object/data.txt',mode='a') as f:
34     #         f.write(img_file)
35     #         f.write("\n")
36     #设置存储目录 我的是存在当前目录下coco_single_object文件夹下 需要手动创建空文件夹
37     json.dump(data_2, open('./coco_single_object/'+img_first+'.json', 'w'), indent=4)  # indent=4 更加美观显示

最后的结果是82000张 json文件

---------------------------------------------------------------------------------------------------------------------------------------

有了单张json文件之后，就是将mask掩膜提取出二值图片的过程了

注：函数传入4个参数 json_path,img_path,color_img_save,binary_img_save

分别对应 json_path: 上一步提取出的json文件的文件夹路径

img_path: coco数据集下载时原图目录

color_img_save: 存放原图的目录（需要新建此文件夹）

binary_img_save: 存放二值图的目录（需要新建此文件夹）

 1 from __future__ import print_function
 2 from pycocotools.coco import COCO
 3 import os, sys, zipfile
 4 import urllib.request
 5 import shutil
 6 import numpy as np
 7 import skimage.io as io
 8 import matplotlib.pyplot as plt
 9 import pylab
10 pylab.rcParams['figure.figsize'] = (8.0, 10.0)
11 import os
12 def get_single_binaryImg(json_path,img_path,color_img_save,binary_img_save):
13     # json_path json文件路径  从coco数据集的annotations标注json文件中提取出的单个json文件
14     #  img_path 原图目录   下载coco数据集时的原图目录
15     # color_img_save 原图存放目录
16     # binary_img_save 二值图存放目录
17     dir=os.listdir(json_path)
18     for jfile in dir:
19         annFile =os.path.join(json_path,jfile)
20         coco = COCO(annFile)
21         imgIds = coco.getImgIds()
22         img = coco.loadImgs(imgIds[0])[0]
23         dataDir = img_path
24         shutil.copy(os.path.join(dataDir, img['file_name']), color_img_save)
25 
26         # load and display instance annotations
27         # 加载实例掩膜
28         catIds = []
29         for ann in coco.dataset['annotations']:
30             if ann['image_id'] == imgIds[0]:
31                 catIds.append(ann['category_id'])
32 
33         annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
34         width = img['width']
35         height = img['height']
36         anns = coco.loadAnns(annIds)
37         mask_pic = np.zeros((height, width))
38         for single in anns:
39             mask_single = coco.annToMask(single)
40             mask_pic += mask_single
41 
42         for row in range(height):
43             for col in range(width):
44                 if (mask_pic[row][col] > 0):
45                     mask_pic[row][col] = 255
46 
47         imgs = np.zeros(shape=(height, width, 3), dtype=np.float32)
48         imgs[:, :, 0] = mask_pic[:, :]
49         imgs[:, :, 1] = mask_pic[:, :]
50         imgs[:, :, 2] = mask_pic[:, :]
51         imgs = imgs.astype(int)
52         img_name = img['file_name'].split(".")[0]
53         plt.imsave(binary_img_save + "/" + img_name + ".png", imgs)
54 
55 if __name__ == '__main__':
56 
57     json_path =r"G:\jianfeng\code\dataset\cocoapi-master\PythonAPI\get_json\test"
58     img_path=r"G:\jianfeng\code\dataset\coco\train2014"
59     color_img_save = r"G:\jianfeng\code\dataset\cocoapi-master\PythonAPI\get_json\color_img"
60     binary_img_save = r"G:\jianfeng\code\dataset\cocoapi-master\PythonAPI\get_json\binary_img"
61 
62     get_single_binaryImg(json_path,img_path,color_img_save,binary_img_save)