labelme转COCO数据集（物体检测）

本次采用的数据集是Labelme标注的数据集，地址：链接：https://pan.baidu.com/s/1nxo9-NpNWKK4PwDZqwKxGQ 提取码：kp4e，需要将其转为COCO格式的数据集。转换代码如下：

新建labelme2coco.py

import argparse
import json
import matplotlib.pyplot as plt
import skimage.io as io
import cv2
from labelme import utils
import numpy as np
import glob
import PIL.Image

REQUIRE_MASK = False
labels = {'aircraft': 1, 'oiltank': 2}


class labelme2coco(object):
    def __init__(self, labelme_json=[], save_json_path='./new.json'):
        '''
        :param labelme_json: the list of all labelme json file paths
        :param save_json_path: the path to save new json
        '''
        self.labelme_json = labelme_json
        self.save_json_path = save_json_path
        self.images = []
        self.categories = []
        self.annotations = []
        # self.data_coco = {}
        self.label = []
        self.annID = 1
        self.height = 0
        self.width = 0
        self.require_mask = REQUIRE_MASK
        self.save_json()

    def data_transfer(self):
        for num, json_file in enumerate(self.labelme_json):
            if not json_file == self.save_json_path:
                with open(json_file, 'r') as fp:
                    data = json.load(fp)
                    self.images.append(self.image(data, num))
                    for shapes in data['shapes']:
                        print("label is ")
                        print(shapes['label'])
                        label = shapes['label']
                        #                    if label[1] not in self.label:
                        if label not in self.label:
                            print("find new category: ")
                            self.categories.append(self.categorie(label))
                            print(self.categories)
                            # self.label.append(label[1])
                            self.label.append(label)
                        points = shapes['points']
                        self.annotations.append(self.annotation(points, label, num))
                        self.annID += 1

    def image(self, data, num):
        image = {}
        img = utils.img_b64_to_arr(data['imageData'])
        height, width = img.shape[:2]
        img = None
        image['height'] = height
        image['width'] = width
        image['id'] = num + 1
        image['file_name'] = data['imagePath'].split('/')[-1]
        self.height = height
        self.width = width
        return image

    def categorie(self, label):
        categorie = {}
        categorie['supercategory'] = label
        #        categorie['supercategory'] = label
        categorie['id'] = labels[label]  # 0 默认为背景
        categorie['name'] = label
        return categorie

    def annotation(self, points, label, num):
        annotation = {}
        print(points)
        x1 = points[0][0]
        y1 = points[0][1]
        x2 = points[1][0]
        y2 = points[1][1]
        contour = np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]])  # points = [[x1, y1], [x2, y2]] for rectangle
        contour = contour.astype(int)
        area = cv2.contourArea(contour)
        print("contour is ", contour, " area = ", area)
        annotation['segmentation'] = [list(np.asarray([[x1, y1], [x2, y1], [x2, y2], [x1, y2]]).flatten())]
        # [list(np.asarray(contour).flatten())]
        annotation['iscrowd'] = 0
        annotation['area'] = area
        annotation['image_id'] = num + 1

        if self.require_mask:
            annotation['bbox'] = list(map(float, self.getbbox(points)))
        else:
            x1 = points[0][0]
            y1 = points[0][1]
            width = points[1][0] - x1
            height = points[1][1] - y1
            annotation['bbox'] = list(np.asarray([x1, y1, width, height]).flatten())

        annotation['category_id'] = self.getcatid(label)
        annotation['id'] = self.annID
        return annotation

    def getcatid(self, label):
        for categorie in self.categories:
            #            if label[1]==categorie['name']:
            if label == categorie['name']:
                return categorie['id']
        return -1

    def getbbox(self, points):
        polygons = points
        mask = self.polygons_to_mask([self.height, self.width], polygons)
        return self.mask2box(mask)

    def mask2box(self, mask):

        # np.where(mask==1)
        index = np.argwhere(mask == 1)
        rows = index[:, 0]
        clos = index[:, 1]

        left_top_r = np.min(rows)  # y
        left_top_c = np.min(clos)  # x
        right_bottom_r = np.max(rows)
        right_bottom_c = np.max(clos)
        return [left_top_c, left_top_r, right_bottom_c - left_top_c, right_bottom_r - left_top_r]

    def polygons_to_mask(self, img_shape, polygons):
        mask = np.zeros(img_shape, dtype=np.uint8)
        mask = PIL.Image.fromarray(mask)
        xy = list(map(tuple, polygons))
        PIL.ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1)
        mask = np.array(mask, dtype=bool)
        return mask

    def data2coco(self):
        data_coco = {}
        data_coco['images'] = self.images
        data_coco['categories'] = self.categories
        data_coco['annotations'] = self.annotations
        return data_coco

    def save_json(self):
        print("in save_json")
        self.data_transfer()
        self.data_coco = self.data2coco()

        print(self.save_json_path)
        json.dump(self.data_coco, open(self.save_json_path, 'w'), indent=4)


labelme_json = glob.glob('LabelmeData/*.json')
from sklearn.model_selection import train_test_split

trainval_files, test_files = train_test_split(labelme_json, test_size=0.2, random_state=55)
import os

if not os.path.exists("datasets/coco/annotations"):
    os.makedirs("datasets/coco/annotations/")
if not os.path.exists("/datasets/coco/train2017"):
    os.makedirs("datasets/coco/train2017")
if not os.path.exists("datasets/coco/val2017"):
    os.makedirs("datasets/coco/val2017")
labelme2coco(trainval_files, 'datasets/coco/annotations/instances_train2017.json')
labelme2coco(test_files, 'datasets/coco/annotations/instances_val2017.json')
import shutil

for file in trainval_files:
    shutil.copy(os.path.splitext(file)[0] + ".jpg", "datasets/coco/train2017/")
for file in test_files:
    shutil.copy(os.path.splitext(file)[0] + ".jpg", "datasets/coco/val2017/")


  

 

  
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
  
157
  
158
  
159
  
160
  
161
  
162
  
163
  
164
  
165
  
166
  
167
  
168
  
169
  
170
  
171
  
172
  
173
  
174
  
175
 

文章来源: wanghao.blog.csdn.net，作者：AI浩，版权归原作者所有，如需转载，请联系作者。

原文链接：wanghao.blog.csdn.net/article/details/120664330