DatasetProcessing/coco2yolo.py at master · LazyShark2001/DatasetProcessing · GitHub

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
135
136
137
138
139
140
141
142
'''https://blog.csdn.net/TracyGC/article/details/137405686'''
'''https://jingjing.blog.csdn.net/article/details/137162505'''
# import os
# import json
# from tqdm import tqdm
# import argparse
#
# parser = argparse.ArgumentParser()
# # 这里根据自己的json文件位置，换成自己的就行
# parser.add_argument('--json_path', default=r'C:\Users\LazyShark\Desktop\data\COCOdataset\annotations\instances_val2017.json', type=str,
#                     help="input: coco format(json)")
# # 这里设置.txt文件保存位置
# parser.add_argument('--save_path', default=r'C:\Users\LazyShark\Desktop\data\COCOdataset\annotations\1', type=str,
#                     help="specify where to save the output dir of labels")
# arg = parser.parse_args()
#
#
# def convert(size, box):
#     dw = 1. / (size[0])
#     dh = 1. / (size[1])
#     x = box[0] + box[2] / 2.0
#     y = box[1] + box[3] / 2.0
#     w = box[2]
#     h = box[3]
#     # round函数确定(xmin, ymin, xmax, ymax)的小数位数
#     x = round(x * dw, 6)
#     w = round(w * dw, 6)
#     y = round(y * dh, 6)
#     h = round(h * dh, 6)
#     return (x, y, w, h)
#
#
# if __name__ == '__main__':
#     json_file = arg.json_path  # COCO Object Instance 类型的标注
#     ana_txt_save_path = arg.save_path  # 保存的路径
#
#     data = json.load(open(json_file, 'r'))
#     if not os.path.exists(ana_txt_save_path):
#         os.makedirs(ana_txt_save_path)
#
#     id_map = {}  # coco数据集的id不连续！重新映射一下再输出！
#     with open(os.path.join(ana_txt_save_path, 'classes.txt'), 'w') as f:
#         # 写入classes.txt
#         for i, category in enumerate(data['categories']):
#             f.write(f"{category['name']}\n")
#             id_map[category['id']] = i
#     # print(id_map)
#     # 这里需要根据自己的需要，更改写入图像相对路径的文件位置。
#     list_file = open(os.path.join(ana_txt_save_path, 'train2017.txt'), 'w')
#     for img in tqdm(data['images']):
#         filename = img["file_name"]
#         img_width = img["width"]
#         img_height = img["height"]
#         img_id = img["id"]
#         head, tail = os.path.splitext(filename)
#         ana_txt_name = head + ".txt"  # 对应的txt名字，与jpg一致
#         f_txt = open(os.path.join(ana_txt_save_path, ana_txt_name), 'w')
#         for ann in data['annotations']:
#             if ann['image_id'] == img_id:
#                 box = convert((img_width, img_height), ann["bbox"])
#                 f_txt.write("%s %s %s %s %s\n" % (id_map[ann["category_id"]], box[0], box[1], box[2], box[3]))
#         f_txt.close()
#         # 将图片的相对路径写入train2017或val2017的路径
#         list_file.write('G:/coco/images/train2017/%s.jpg\n' % (head))
#     list_file.close()
#


"""
https://blog.csdn.net/qiumokucao/article/details/119151978?spm=1001.2014.3001.5502
优化版
COCO 格式的数据集转化为 YOLO 格式的数据集
--json_path 输入的json文件路径
--save_path 保存的文件夹名字，默认为当前目录下的labels。
"""

import os
import json
from tqdm import tqdm
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--json_path', default=r'C:\Users\LazyShark\Desktop\data\COCOdataset\annotations\instances_val2017.json', type=str, help="input: coco format(json)")
parser.add_argument('--save_path', default=r'C:\Users\LazyShark\Desktop\data\COCOdataset\annotations\1', type=str, help="specify where to save the output dir of labels")
arg = parser.parse_args()


def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = box[0] + box[2] / 2.0
    y = box[1] + box[3] / 2.0
    w = box[2]
    h = box[3]

    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)


if __name__ == '__main__':
    json_file = arg.json_path  # COCO Object Instance 类型的标注
    ana_txt_save_path = arg.save_path  # 保存的路径

    data = json.load(open(json_file, 'r'))
    if not os.path.exists(ana_txt_save_path):
        os.makedirs(ana_txt_save_path)

    id_map = {}  # coco数据集的id不连续！重新映射一下再输出！
    with open(os.path.join(ana_txt_save_path, 'classes.txt'), 'w') as f:
        # 写入classes.txt
        for i, category in enumerate(data['categories']):
            f.write(f"{category['name']}\n")
            id_map[category['id']] = i
    # print(id_map)

    anns = {}
    for ann in data['annotations']:
        imgid = ann['image_id']
        anns.setdefault(imgid, []).append(ann)

    print('got anns')
    # 将图片的名字写入txt文件中
    list_file = open(os.path.join(ana_txt_save_path, 'train2017.txt'), 'w')
    for img in tqdm(data['images']):
        filename = img["file_name"]
        img_width = img["width"]
        img_height = img["height"]
        img_id = img["id"]
        head, tail = os.path.splitext(filename)
        ana_txt_name = head + ".txt"  # 对应的txt名字，与jpg一致
        f_txt = open(os.path.join(ana_txt_save_path, ana_txt_name), 'w')

        ann_img = anns.get(img_id, [])
        for ann in ann_img:
            box = convert((img_width, img_height), ann["bbox"])
            f_txt.write("%s %s %s %s %s\n" % (id_map[ann["category_id"]], box[0], box[1], box[2], box[3]))
        f_txt.close()
        list_file.write('G:/coco/images/train2017/%s.jpg\n' % (head))
    list_file.close()