【转换格式脚本】yolo\visdrone\detection转换COCO格式_yolo转coco

1. yolo2coco.py

import imp
import json
import cv2
import os

bnd_id_start = 0

times = 0

json_dict = {
    "images"     : [],
    "type"       : "instances",
    "annotations": [],
    "categories" : []
}


# 这里是你的txt文件的读取
# with open('train.txt','r') as f:
#     data = f.readlines()
raw_images_path = 'images'
raw_labels_path = 'labels'
data = os.listdir(raw_images_path)


bnd_id = bnd_id_start

#类别的名字(cid,cate)对应
classes = ["your dataset categories"]

classes_count_obj = {0:0, 1:0, 2:0, 3:0, 4:0, 5:0, 6:0, 7:0, 8:0, 9:0, 10:0, 11:0, 12:0, 13:0, 14:0, 15:0, 16:0, 17:0, 18:0, 19:0, 20:0, 21:0}

classes_count_imgs = {0:0, 1:0, 2:0, 3:0, 4:0, 5:0, 6:0, 7:0, 8:0, 9:0, 10:0, 11:0, 12:0, 13:0, 14:0, 15:0, 16:0, 17:0, 18:0, 19:0, 20:0, 21:0}

for d in data:
    filename = d     #这里可能修改，txt文件每一行第一个属性是图片路径，通过split()函数把图像名分离出来就行
    img = cv2.imread(os.path.join(raw_images_path, filename))
    txtFile = filename.replace('jpeg', 'txt').replace('jpg', 'txt')
    # import pdb;pdb.set_trace()
    try:
        height,width = img.shape[0],img.shape[1]
        image_id = filename.split(".")[0]
    except:
        times += 1
        print('file is error')

# type 已经填充

#定义image 填充到images里面
    image = {
        'file_name' : filename,  #文件名
        'height'    : height,    #图片的高
        'width'     : width,     #图片的宽
        'id'        : image_id   #图片的id，和图片名对应的
    }
    json_dict['images'].append(image)
  
    with open(os.path.join(raw_labels_path, txtFile), 'r') as fr:
        labelList = fr.readlines()
        cls_list = set()  # 用于统计在一张图上每个类的类别id
        for c in labelList:
            
            label, xmin, ymin, w, h = c.strip().split(" ")
            label = int(label)
            classes_count_obj[label] += 1 
            cls_list.add(label) # 如果有这个类，将其加入到cls_list中
            xmin = float(xmin)
            ymin = float(ymin)
            w = float(w)
            h = float(h)
            x1 = width * xmin - 0.5 * width * w
            y1 = height * ymin - 0.5 * height * h
            x2 = width * xmin + 0.5 * width * w
            y2 = height * ymin + 0.5 * height * h
            o_width = abs(x2 - x1)
            o_height = abs(y2 - y1)

            area = o_width * o_height

            # #定义annotation
            annotation = {
                'area'          : area,  #
                'iscrowd'       : 0,
                'image_id'      : image_id,  #图片的id
                'bbox'          :[x1, y1, o_width,o_height],
                'category_id'   : label, #类别的id 通过这个id去查找category里面的name
                'id'            : bnd_id,  #唯一id ,可以理解为一个框一个Id
                'ignore'        : 0,
                'segmentation'  : [[x1, y1, x1 + o_width, y1, x1 + o_width, y1 + o_height, x1, y1 + o_height]]
            }

            json_dict['annotations'].append(annotation)

            bnd_id += 1

        for i in cls_list:
            if i in classes_count_imgs:
                classes_count_imgs[i] += 1


#定义categories

for i in range(len(classes)):

    cate = classes[i]
    cid = i
    category = {
        'supercategory' : 'none',
        'id'            : cid,  #类别的id ,一个索引，主键作用，和别的字段之间的桥梁
        'name'          : cate  #类别的名字
    }

    json_dict['categories'].append(category)



json_fp = open("your json path",'w')
json_str = json.dumps(json_dict, indent=4)
json_fp.write(json_str)
json_fp.close()

print("cls obj nums:", classes_count_obj)
print()
print("cls img nums:", classes_count_imgs)
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

2. visdrone2coco.py

# coding:utf-8
# 该代码只提取类别1和2
import os
import cv2
from tqdm import tqdm
import json
import uuid
from os.path import join as opj
import shutil

mapping = json.load(open(os.path.join("uuid_mapping.json")))
# 如果你有{image_name:uuid}这个映射，可以添加到上面的路径

def convert_to_cocodetection(dir, output_dir):
    train_dir = opj(dir, "VisDrone2019-DET-train")
    val_dir = opj(dir, "VisDrone2019-DET-val")
    test_dir = opj(dir, "VisDrone2019-DET-test-dev")
    train_annotations = opj(train_dir, "annotations")
    val_annotations = opj(val_dir, "annotations")
    test_annotations = opj(test_dir, "annotations")
    train_images = opj(train_dir, "images")
    val_images = opj(val_dir, "images")
    test_images = opj(test_dir, "images")
    id_num = 0
    num = 0
 
    categories = [
        {"supercategory": "ignored regions", "id": 0, "name": "ignored regions"},
        {"supercategory": "pedestrian", "id": 1, "name": "pedestrian"},
        {"supercategory": "people", "id": 2, "name": "people"},
        {"supercategory": "bicycle", "id": 3, "name": "bicycle"},
        {"supercategory": "car", "id": 4, "name": "car"},
        {"supercategory": "van", "id": 5, "name": "van"},
        {"supercategory": "truck", "id": 6, "name": "truck"},
        {"supercategory": "tricycle", "id": 7, "name": "tricycle"},
        {"supercategory": "awning-tricycle", "id": 8, "name": "awning-tricycle"},
        {"supercategory": "bus", "id": 9, "name": "bus"},
        {"supercategory": "motor", "id": 10, "name": "motor"},
        {"supercategory": "others", "id": 11, "name": "others"}
    ]
    # categories = [{"id": 0, "name": "person"}]
    for mode in ["train", "val", "test"]:
        images = []
        annotations = []
        print(f"start loading {mode} data...")
        if mode == "train":
            set = os.listdir(train_annotations)
            annotations_path = train_annotations
            images_path = train_images
        elif mode == "test":
            set = os.listdir(test_annotations)
            annotations_path = test_annotations
            images_path = test_images
        else:
            set = os.listdir(val_annotations)
            annotations_path = val_annotations
            images_path = val_images
        for i in tqdm(set):
            f = open(opj(annotations_path, i), "r")
            name = i.replace(".txt", "")
            image = {}
            height, width = cv2.imread(opj(images_path, name + ".jpg")).shape[:2]
            file_name = opj("images", "visdrone", "raw", mapping[name] + ".jpg") # 若没有mapping，把这个行改成你需要的file_name即可
            image["file_name"] = file_name
            image["coco_url"] = ""
            image["height"] = height
            image["width"] = width
            image["id"] = mapping[name]
            images.append(image)         
            for line in f.readlines():
                annotation = {}
                line = line.replace("\n", "")
                if line.endswith(","):  # filter data
                    line = line.rstrip(",")
                line_list = [int(i) for i in line.split(",")]
                category_id = line_list[5]
                # if category_id == 0 or category_id == 11:
                #     continue
                # else:
                bbox_xywh = [line_list[0], line_list[1], line_list[2], line_list[3]]
                annotation["image_id"] = mapping[name]
                annotation["bbox"] = bbox_xywh
                annotation["category_id"] = category_id
                annotation["id"] = id_num
                annotation["iscrowd"] = 0
                annotation["segmentation"] = [[bbox_xywh[0],bbox_xywh[1],bbox_xywh[0]+bbox_xywh[2],bbox_xywh[1],bbox_xywh[0]+bbox_xywh[2],bbox_xywh[1]+bbox_xywh[3],bbox_xywh[0],bbox_xywh[1]+bbox_xywh[3]]]
                annotation["area"] = bbox_xywh[2] * bbox_xywh[3]
                id_num += 1
                annotations.append(annotation)
            num += 1
            # shutil.copy2(opj(images_path, name + ".jpg"), opj("/data0/public/data/visdrone/images/Visdrone_OD/raw", uid + ".jpg"))

        dataset_dict = {}
        dataset_dict["images"] = images
        dataset_dict["annotations"] = annotations
        dataset_dict["categories"] = categories
        dataset_dict["info"] = dict()
        dataset_dict["licenses"] = dict()
        json_str = json.dumps(dataset_dict, indent=4)
        with open(f'{output_dir}/split_{mode}.json', 'w') as json_file:
            json_file.write(json_str)
    print("json file write done...")
 
 
def get_test_namelist(dir, out_dir):
    full_path = out_dir + "/" + "test.txt"
    file = open(full_path, 'w')
    for name in tqdm(os.listdir(dir)):
        name = name.replace(".txt", "")
        file.write(name + "\n")
    file.close()
    return None
 
 
def centerxywh_to_xyxy(boxes):
    """
    args:
        boxes:list of center_x,center_y,width,height,
    return:
        boxes:list of x,y,x,y,cooresponding to top left and bottom right
    """
    x_top_left = boxes[0] - boxes[2] / 2
    y_top_left = boxes[1] - boxes[3] / 2
    x_bottom_right = boxes[0] + boxes[2] / 2
    y_bottom_right = boxes[1] + boxes[3] / 2
    return [x_top_left, y_top_left, x_bottom_right, y_bottom_right]
 
 
def centerxywh_to_topleftxywh(boxes):
    """
    args:
        boxes:list of center_x,center_y,width,height,
    return:
        boxes:list of x,y,x,y,cooresponding to top left and bottom right
    """
    x_top_left = boxes[0] - boxes[2] / 2
    y_top_left = boxes[1] - boxes[3] / 2
    width = boxes[2]
    height = boxes[3]
    return [x_top_left, y_top_left, width, height]
 
 
def clamp(coord, width, height):
    if coord[0] < 0:
        coord[0] = 0
    if coord[1] < 0:
        coord[1] = 0
    if coord[2] > width:
        coord[2] = width
    if coord[3] > height:
        coord[3] = height
    return coord
 
 
if __name__ == '__main__':
    convert_to_cocodetection(r"your src path",r"your output path")
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
143
144
145
146
147
148
149
150
151
152
153
154
155
156

3. 把yolov5 detect的输出转成coco

import json
import cv2
import os

times = 0

pred_json = []

json_dict = {
    "image_id": int,
    "bbox": list,
    "category_id" : int,
    "score": float
}

raw_images_path = 'images'
raw_labels_path = 'labels'
data = os.listdir(raw_images_path)

for d in data:
    filename = d     #这里可能修改，txt文件每一行第一个属性是图片路径，通过split()函数把图像名分离出来就行
    img = cv2.imread(os.path.join(raw_images_path, filename))
    txtFile = filename.replace('jpeg', 'txt').replace('jpg', 'txt')
    # import pdb;pdb.set_trace()  
    try:
        height,width = img.shape[0],img.shape[1]
        image_id = filename.split(".")[0]
    except:
        times += 1
        print('file is error')
    
    if not os.path.exists(os.path.join(raw_labels_path, txtFile)):
        continue
    with open(os.path.join(raw_labels_path, txtFile), 'r') as fr:
        labelList = fr.readlines()
        for c in labelList:
            
            label, xmin, ymin, w, h, score = c.strip().split(" ")
            label = int(label)
            xmin = float(xmin)
            ymin = float(ymin)
            w = float(w)
            h = float(h)
            score = float(score)
            x1 = width * xmin - 0.5 * width * w
            y1 = height * ymin - 0.5 * height * h
            x2 = width * xmin + 0.5 * width * w
            y2 = height * ymin + 0.5 * height * h
            
            json_dict['category_id'] = label
            json_dict['image_id'] = image_id
            json_dict['score'] = score
            box = [x1, y1, x2, y2]
            json_dict['bbox'] = box
        
            pred_json.append(json_dict)

print(len(pred_json))

json_fp = open("best_predictions.json",'w')
json_str = json.dumps(pred_json, indent=4)
json_fp.write(json_str)
json_fp.close()

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

4. visdrone-MOT转coco

# coding:utf-8
"""
YOLO 格式的数据集转化为 COCO 格式的数据集
--root_dir 输入根路径
--save_path 保存文件的名字
"""

import os
import cv2
import json
from tqdm import tqdm
import argparse
import progressbar
from PIL import Image

parser = argparse.ArgumentParser()
parser.add_argument('--root_dir', default='VisDrone2019/Test',type=str, help="root path of images and labels, include ./images and ./labels and classes.txt")
parser.add_argument('--save_path', type=str,default='annotations/test.json', help="if not split the dataset, give a path to a json file")

arg = parser.parse_args()

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

def yolo2coco(arg):
    root_path = arg.root_dir
    print("Loading data from ",root_path)

    assert os.path.exists(root_path)
    originLabelsDir = os.path.join(root_path, 'labels')                                        
    originImagesDir = os.path.join(root_path, 'JPEGImages')
    classes = '0'
    # images dir name
    indexes = os.listdir(originImagesDir)

    dataset = {'categories': [], 'annotations': [], 'images': []}
    for i, cls in enumerate(classes, 0):     
        dataset['categories'].append({'id': i, 'name': cls, 'supercategory': 'mark'})
    
    # 标注的id
    ann_id_cnt = 0
    for k, index in enumerate(tqdm(indexes)):
        # 支持 png jpg 格式的图片。
        txtFile = index.replace('images','txt').replace('.jpg','.txt').replace('.png','.txt')
        # 读取图像的宽和高
        im = cv2.imread(os.path.join(root_path, 'JPEGImages/') + index)
        height, width, _ = im.shape
        # 添加图像的信息
        dataset['images'].append({'file_name': index,
                                    'id': k,
                                    'width': width,
                                    'height': height})
        # if not os.path.exists(os.path.join(originLabelsDir, txtFile)):
        #     # 如没标签，跳过，只保留图片信息。
        #     continue
        with open(os.path.join(originLabelsDir, txtFile), 'r') as fr:
            labelList = fr.readlines()
            if len(labelList) == 0:
                # import pdb;pdb.set_trace()

                dataset['annotations'].append({
                    'image_name': index[:-4],
                    'area': 0,
                    'bbox': [],
                    'category_id': 0,
                    'id': ann_id_cnt,  # 目标总数
                    'image_id': k,
                    'iscrowd': 0,
                    # mask, 矩形是从左上角点按顺时针的四个顶点
                    'segmentation': []
                })
                ann_id_cnt += 1
            else:
                for label in labelList:
                    label = label.strip().split()
                    x = float(label[1])
                    y = float(label[2])
                    w = float(label[3])
                    h = float(label[4])

                    # convert x,y,w,h to x1,y1,x2,y2 (left-top + w + h)
                    H, W, _ = im.shape
                    x1 = (x - w / 2) 
                    y1 = (y - h / 2)  # top-left
                    x2 = (x + w / 2) 
                    y2 = (y + h / 2)  # right-bottom
                    # 标签序号从0开始计算
                    cls_id = int(label[0])   
                    width = max(0, x2 - x1)
                    height = max(0, y2 - y1)
                    dataset['annotations'].append({
                        'image_name': index[:-4],
                        'area': width * height,
                        'bbox': [x1, y1, width, height],
                        'category_id': cls_id,
                        'id': ann_id_cnt,
                        'image_id': k,
                        'iscrowd': 0,
                        # mask, 矩形是从左上角点按顺时针的四个顶点
                        'segmentation': [[x1, y1, x2, y1, x2, y2, x1, y2]]
                    })
                    ann_id_cnt += 1

    # 保存结果
    folder = os.path.join("VisDrone2019/", 'annotations')
    if not os.path.exists(folder):
        os.makedirs(folder)

    json_name = os.path.join("VisDrone2019/", arg.save_path)
    # import pdb;pdb.set_trace()
    with open(json_name, 'w') as f:
        json.dump(dataset, f, indent=4)
        print('Save annotation to {}'.format(json_name))

if __name__ == "__main__":

    yolo2coco(arg)


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