-
Notifications
You must be signed in to change notification settings - Fork 1
/
coco2voc.py
executable file
·134 lines (118 loc) · 4.8 KB
/
coco2voc.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
'''
把coco数据集合的所有标注转换到voc格式,不改变图片命名方式,
注意,原来有一些图片是黑白照片,检测出不是 RGB 图像,这样的图像不会被放到新的文件夹中
更新日期:2019-11-19
'''
from pycocotools.coco import COCO # 这个包可以从git上下载https://github.com/cocodataset/cocoapi/tree/master/PythonAPI,也可以直接用修改后的coco.py
import os, cv2, shutil
from lxml import etree, objectify
from tqdm import tqdm
from PIL import Image
CKimg_dir = '/home/gao/datasets/tt100k_2021/images' # coco图片路径
CKanno_dir = '/home/gao/datasets/tt100k_2021/dataset/annotations' # coco注释路径
# 若模型保存文件夹不存在,创建模型保存文件夹,若存在,删除重建
def mkr(path):
if os.path.exists(path):
shutil.rmtree(path)
os.mkdir(path)
else:
os.mkdir(path)
def save_annotations(filename, objs, filepath):
annopath = CKanno_dir + "/" + filename[:-3] + "xml" # 生成的xml文件保存路径
dst_path = CKimg_dir + "/" + filename
img_path = filepath
img = cv2.imread(img_path)
im = Image.open(img_path)
if im.mode != "RGB":
print(filename + " not a RGB image")
im.close()
return
im.close()
shutil.copy(img_path, dst_path) # 把原始图像复制到目标文件夹
E = objectify.ElementMaker(annotate=False)
anno_tree = E.annotation(
E.folder('1'),
E.filename(filename),
E.source(
E.database('CKdemo'),
E.annotation('VOC'),
E.image('CK')
),
E.size(
E.width(img.shape[1]),
E.height(img.shape[0]),
E.depth(img.shape[2])
),
E.segmented(0)
)
for obj in objs:
E2 = objectify.ElementMaker(annotate=False)
anno_tree2 = E2.object(
E.name(obj[0]),
E.pose(),
E.truncated("0"),
E.difficult(0),
E.bndbox(
E.xmin(obj[2]),
E.ymin(obj[3]),
E.xmax(obj[4]),
E.ymax(obj[5])
)
)
anno_tree.append(anno_tree2)
etree.ElementTree(anno_tree).write(annopath, pretty_print=True)
def showbycv(coco, dataType, img, classes, origin_image_dir, verbose=False):
filename = img['file_name']
filepath = os.path.join(origin_image_dir, dataType, filename)
I = cv2.imread(filepath)
annIds = coco.getAnnIds(imgIds=img['id'], iscrowd=None)
anns = coco.loadAnns(annIds)
objs = []
for ann in anns:
name = classes[ann['category_id']]
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 = [name, 1.0, xmin, ymin, xmax, ymax]
objs.append(obj)
if verbose:
cv2.rectangle(I, (xmin, ymin), (xmax, ymax), (255, 0, 0))
cv2.putText(I, name, (xmin, ymin), 3, 1, (0, 0, 255))
save_annotations(filename, objs, filepath)
if verbose:
cv2.imshow("img", I)
cv2.waitKey(0)
def catid2name(coco): # 将名字和id号建立一个字典
classes = dict()
for cat in coco.dataset['categories']:
classes[cat['id']] = cat['name']
# print(str(cat['id'])+":"+cat['name'])
return classes
def get_CK5(origin_anno_dir, origin_image_dir, verbose=False):
dataTypes = ['train2017', 'val2017', 'val2017']
for dataType in dataTypes:
annFile = 'instances_{}.json'.format(dataType)
annpath = os.path.join(origin_anno_dir, annFile)
coco = COCO(annpath)
classes = catid2name(coco)
imgIds = coco.getImgIds()
# imgIds=imgIds[0:1000]#测试用,抽取10张图片,看下存储效果
for imgId in tqdm(imgIds):
img = coco.loadImgs(imgId)[0]
showbycv(coco, dataType, img, classes, origin_image_dir, verbose=False)
def main():
base_dir = '/home/jerry/data/data/object_detect/coco2017_voc' # step1 这里是一个新的文件夹,存放转换后的图片和标注
image_dir = os.path.join(base_dir, 'images') # 在上述文件夹中生成images,annotations两个子文件夹
anno_dir = os.path.join(base_dir, 'annotations')
mkr(image_dir)
mkr(anno_dir)
origin_image_dir = '/home/jerry/data/data/object_detect/coco2017/images' # step 2原始的coco的图像存放位置
origin_anno_dir = '/home/jerry/data/data/object_detect/coco2017/annotations' # step 3 原始的coco的标注存放位置
verbose = True # 是否需要看下标记是否正确的开关标记,若是true,就会把标记展示到图片上
get_CK5(origin_anno_dir, origin_image_dir, verbose)
if __name__ == "__main__":
main()
# split_traintest()