机器视觉 AI 数据集制作_一键识别物体数据集制作

工业中，机器视觉物体分拣时，需要制作，数据集，那么，一般情况下，可以选择几个物体的几张图片，或者视频，将待识别的物体的掩模扣取出来，随机的贴在 传送带背景中，并批量自动的写成 VOC 数据集

使用图像处理的技术手段，将上述的目标的掩模扣取出来，或者使用 ps 的技术扣取掩模均可。

# -*- coding :  utf-8 -*-
# @Data      :  2019-08-16
# @Author    :  xm
# @Email     :
# @File      :  image_process.py
# Desctiption:  求取图像中物体的边界矩形

import numpy as np
import cv2
import os


def calculatBoundImage(src_Image):
    """
    求取图像中物体的边界矩形框
    :param src_Image: 输出的源图像
    :return: 返回图像中的物体边界矩形
    """

    tmp_image = src_Image.copy()
    #print(tmp_image)
    if (len(tmp_image.shape) == 3):
        tmp_image = cv2.cvtColor(tmp_image, cv2.COLOR_BGR2GRAY)
    # 自适应阈值进行二值化
    thresh_image = cv2.adaptiveThreshold(tmp_image, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 71, 10)
    thresh_image = cv2.morphologyEx(thresh_image, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (25, 25)))
    # 寻找最外层轮廓
    contours_ls, hierarchy = cv2.findContours(thresh_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_KCOS)
    pnt_cnt_ls = np.array([tmp_contour.shape[0] for tmp_contour in contours_ls])

    contour_image = src_Image.copy()
    contours_idx = np.argmax(pnt_cnt_ls)
    contour_image = cv2.drawContours(contour_image, contours_ls, contours_idx, (0, 0, 255))
    longest_contour = contours_ls[contours_idx]

    countour_image_gray = np.zeros(src_Image.shape, dtype=np.uint8)
    countour_image_gray = cv2.drawContours(countour_image_gray, contours_ls, contours_idx, (1, 1, 1), cv2.FILLED)
    obj_image = src_Image * countour_image_gray
    bound_box = cv2.boundingRect(longest_contour)
    return bound_box, contour_image, obj_image


def rotateImage(src_Image, angle_deg, rotate_center=None):
    """
    对目标图像进行旋转
    :param src_Image: 输入的源图像
    :param angle_deg: 旋转的角度
    :param rotate_center: 旋转的中心
    :return: 旋转后的图片
    """
    (h, w) = src_Image.shape[:2]
    if rotate_center is None:
        rotate_center = ((w -1) / 2, (h - 1) / 2)
    rot_mat = cv2.getRotationMatrix2D(rotate_center, angle_deg, 1.0)
    rot_iamge = cv2.warpAffine(src_Image, rot_mat, (w, h))
    return rot_iamge


def VideotoImage(video_file, folder_path):
    """
    数据的视频保存为提取之后的物体图
    :param video_file: 视频文件
    :param folder_path: 保存图片的路径
    :return: 保存的图片
    """
    video_cap = cv2.VideoCapture(video_file)
    image_idx = 2000
    while True:
        ret, frame = video_cap.read()
        if (frame is None):
            continue
        bound_box, contour_image, obj_image = calculatBoundImage(frame)
        bound_thres = 4500

        if (bound_box[2] > bound_thres or bound_box[3] > bound_thres):
            continue
        contour_image = cv2.rectangle(contour_image, (bound_box[0], bound_box[1]),(bound_box[0] + bound_box[2],bound_box[1] + bound_box[3]), (225, 0, 0), thickness=2)
        #cv2.imshow('frame', contour_image)

        image_name = str(image_idx).zfill(6) + '.jpg'
        image_idx += 1
        if image_idx % 2 == 0:
            cv2.imwrite(folder_path + image_name, obj_image)
        cv2.waitKey(25)
        if 0xFF & cv2.waitKey(5) == 27:
            break
    video_cap.release()


def BatchImageProcess(image_path, folder_path):
    """
    批量图片物体提取，背景为黑色
    :param Image_path: 图片的路径
    :param folder_path: 图像处理之后的保存路径
    :return: 保存的图片
    """
    image_file_list = os.listdir(image_path)
    # 获取物体图像的文件名
    image_idx = 0
    for image_name in range(len(image_file_list)):
        obj_image_path = image_path + image_file_list[image_idx]
        src_Image = cv2.imread(obj_image_path)

        bound_box, contour_image, obj_image = calculatBoundImage(src_Image)
        bound_thres = 4500

        if (bound_box[2] > bound_thres or bound_box[3] > bound_thres):
            continue
        contour_image = cv2.rectangle(contour_image, (bound_box[0], bound_box[1]), (bound_box[0] + bound_box[2], bound_box[1] + bound_box[3]), (225, 0, 0), thickness=2)
        #cv2.imshow('frame', contour_image)
        image_name = str(image_idx).zfill(6) + '.jpg'
        cv2.imwrite(folder_path + image_name, obj_image)
        image_idx += 1


def main():
    image_path = "/home/xm/workspace/ImageProcess/tmp/circle/"
    folder_path = "/home/xm/workspace/ImageProcess/tmp/"
    BatchImageProcess(image_path, folder_path)


# def main():
#     src_Image = cv2.imread("./Images/00001.png")
#     bound_box, contour_image, obj_image = calculatBoundImage(src_Image)
#     print("bound_box", bound_box)
#
#     cv2.namedWindow("input image", cv2.WINDOW_AUTOSIZE)
#     cv2.imshow("input image", contour_image)
#
#
#     # 一般源图像进行旋转再提取轮廓
#     rot_image = rotateImage(src_Image, 20, rotate_center=None)
#     cv2.imshow("obj image", obj_image)
#     cv2.imshow("rot image", rot_image)
#     cv2.waitKey(0)
#
#     # vide_file = "./Images/blue_1_82.mp4"
#     # folder_path = "./results/"
#     #
#     # VideotoImage(vide_file, folder_path)


if __name__ == "__main__":
    main()

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

# -*- coding :  utf-8 -*-
# @Data      :  2019-08-17
# @Author    :  xm
# @Email     :
# @File      :  ImageDataSetGeneration.py
# Desctiption:  生成物体分类的图像数据集

import numpy as np
import cv2
import os
from lxml import etree, objectify


def rotateImage(src_image, rotate_deg):
    """
    对图像进行旋转
    :param src_image: 输入源图像
    :param rotate_dog: 旋转角度
    :return: 旋转后的图像
    """
    img_h, img_w = src_image.shape[0:2]
    rotate_mat = cv2.getRotationMatrix2D((img_w / 2.0, img_h / 2.0), rotate_deg, 1.0)
    dst_image = cv2.warpAffine(src_image, rotate_mat, (img_w, img_h))
    return dst_image


def calculateBoundImage(src_image):
    """
    求图像中物体的边界矩形
    :param src_image: 源图像
    :return: 图像中物体的边界矩形、轮廓图、目标图像
    """

    tmp_image = src_image.copy()
    if len(tmp_image.shape) == 3:
        tmp_image = cv2.cvtColor(tmp_image, cv2.COLOR_BGR2GRAY)
    ret, thresh_images = cv2.threshold(tmp_image, 0, 255,cv2.THRESH_BINARY)
    contours_ls, _ = cv2.findContours(thresh_images, cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
    all_points = np.concatenate(contours_ls, axis=0)
    bound_box = cv2.boundingRect(all_points)
    return bound_box


def randomMoveObjectInImage(src_image, src_bound_box):
    """
    将物体在图像中随机摆放
    :param src_image: 背景图 COCO/VOC
    :param src_bound_box: 原始边界框
    :return: 相机旋转后的边界框
    """
    x, y, w, h = src_bound_box
    img_h, img_w = src_image.shape[0:2]
    img_h -= h
    img_w -= w
    random_array = np.random.uniform(0.0, 1.0, 2)
    bbox_x = np.floor(img_w * random_array[0])
    bbox_y = np.floor(img_h * random_array[1])
    return np.array([bbox_x, bbox_y, w, h])


def calculateIOU(bound_box_1, bound_box_2):
    """
    计算两个 bound_box 之间的 IOU
    :param bound_box_1: 边界框 1, shape [x, y, w, h]
    :param bound_box_2: 边界框 2，shape [x, y, w, h]
    :return: 两个 bound box 之间的 IOU 值
    """

    min_xy = np.maximum(bound_box_1[0:2], bound_box_2[0:2])
    max_xy = np.minimum(bound_box_1[0:2] + bound_box_2[2:4],
                        bound_box_2[0:2] + bound_box_2[2:4])

    delta_xy = max_xy - min_xy
    intersection_area = delta_xy[0] * delta_xy[1]
    if (intersection_area < 0):
        return
    box_area_1 = bound_box_1[2] * bound_box_1[3]
    box_area_2 = bound_box_2[2] * bound_box_2[3]

    union_area = box_area_1 + box_area_2 - intersection_area
    return intersection_area / union_area


def resizeObjectImage(src_image, max_min_box_size):
    """
    对物体图像进行随机缩放
    :param src_image: 原始图像
    :param max_min_box_size: 缩放后图像中的物体的 bound box 的最大边的范围
    :return: 缩放后的图像
    """
    src_bbox = calculateBoundImage(src_image)
    src_bbox_max = np.max(src_bbox[2:4])
    cur_bbox_max = np.random.uniform(max_min_box_size[1], max_min_box_size[0], 1)[-1]
    cur_ratio = cur_bbox_max / src_bbox_max

    src_h, src_w = src_image.shape[0:2]
    dst_h, dst_w = np.floor(src_h * cur_ratio), np.floor(src_w * cur_ratio)
    dst_image = cv2.resize(src_image, (np.int(dst_w), np.int(dst_h)))
    return dst_image


def addObjectToImage(backgroup_image, obj_image, bound_box):
    """
    将目标物体添加到背景图中
    :param backgroup_image: 背景图
    :param obj_image: 目标物体图
    :param bound_box: 边界矩形框
    :return: 添加了目标物体的背景图
    """

    tmp_image = obj_image.copy()
    if len(tmp_image.shape) == 3:
        tmp_image = cv2.cvtColor(tmp_image, cv2.COLOR_BGR2GRAY)
    mask = tmp_image > 5

    min_x, min_y, max_x, max_y = bound_box[0], bound_box[1], bound_box[0] + bound_box[2], bound_box[1] + bound_box[3]
    backgroup_image[np.int(min_y):np.int(max_y), np.int(min_x):np.int(max_x)][mask] = obj_image[mask]
    return backgroup_image


def formImageAndlabel(background_image, obj_ls, max_min_size_ration, iou_thres):
    """
    形成训练图像，并生成对应的 label 列表
    :param background_image: 输入背景图
    :param obj_ls: 目标 list
    :param max_min_size_ration: 最大最小旋转角度
    :param iou_thres: IOU 阈值
    :return: 返训练的图像，对应的 label
    """

    max_ratio, min_ratio = max_min_size_ration
    image_size = np.min(background_image.shape[0:2])
    dst_image = background_image.copy()
    max_min_box_size = [np.floor(max_ratio * image_size), np.floor(min_ratio * image_size)]
    label_ls = []
    for obj_image, obj_name in obj_ls:
        # 对图像进行随机缩放
        resize_obj_image = resizeObjectImage(obj_image, max_min_box_size)
        # 对图像进行随机旋转
        rotate_image = rotateImage(resize_obj_image, np.random.uniform(0, 360, 1)[-1])
        # 多次迭代, 直到将图像平移到适当位置为止
        src_bbox = calculateBoundImage(rotate_image)
        sub_obj_image = rotate_image[src_bbox[1]:src_bbox[1] + src_bbox[3], src_bbox[0]:src_bbox[0] + src_bbox[2]]
        iter_cnt = 100
        if len(label_ls) == 0:
            iter_cnt = 1
        for iter_idx in range(iter_cnt):
            dst_bbox = randomMoveObjectInImage(dst_image, src_bbox)
            if len(label_ls) != 0:
                is_fit = True
                for tmp_box, tmp_obj_name in label_ls:
                    #print("....", tmp_box)
                    #print("+++++", dst_bbox)
                    IOU = calculateIOU(tmp_box, dst_bbox)
                    if (IOU is not None) and (IOU > iou_thres):
                        is_fit = False
                        break
                if is_fit == False:
                    continue
                else:
                    break
        dst_image = addObjectToImage(dst_image, sub_obj_image, dst_bbox)
        label_ls.append([dst_bbox, obj_name])
    return dst_image, label_ls


def formImageLableXML(src_image, image_file_name, label_info, label_path):
    """
    生成图片的 label XML
    :param src_image: 原始图像
    :param image_file_name: 图像的文件名
    :param label_infor: 标签信息
    :param label_path: 标签的路径
    :return: XML
    """

    ele = objectify.ElementMaker(annotate=False)
    anno_tree = ele.annotation(
        ele.folder('VOC2019_xm'),
        ele.filename(image_file_name),
        ele.source(
            ele.database('The VOC2019 Database'),
            ele.annotation('PASCAL VOC2019'),
            ele.image('flickr'),
            ele.flickrid('264265361')
            ),
        ele.owner(
            ele.flickrid('xm'),
            ele.name('xm')
        ),
        ele.size(
            ele.width(str(src_image.shape[0])),
            ele.height(str(src_image.shape[1])),
            ele.depth(str(src_image.shape[2]))
        ),
        ele.segmented('0')
    )
    for cur_box, cur_obj_name in label_info:
        cur_ele = objectify.ElementMaker(annotate=False)
        cur_tree = cur_ele.object(
            ele.name(cur_obj_name),
            ele.pose('Frontal'),
            ele.truncated('0'),
            ele.difficult('0'),
            ele.bndbox(
                ele.xmin(str(cur_box[0])),
                ele.ymin(str(cur_box[1])),
                ele.xmax(str(cur_box[0] + cur_box[2])),
                ele.ymax(str(cur_box[1] + cur_box[3]))
            )
        )
        anno_tree.append(cur_tree)
    etree.ElementTree(anno_tree).write(label_path, pretty_print=True)


def main():
    obj_name_ls = ['circle', 'square']
    # 各种物体对应的图像的路径
    base_obj_file_name = '/home/xm/workspace/ImageProcess/DataSet/'
    obj_file_name = [base_obj_file_name + cur_obj for cur_obj in obj_name_ls]
    print(obj_file_name)
    # 每个种类的样本数量
    obj_count = 600

    # 图像中物体最大的数量
    image_max_obj_cnt = 2

    # 图像中物体的 bound box 的最大尺寸点，整个图像最小尺寸比例，
    max_size_radio = 0.45
    min_size_radio = 0.20

    # 图像的总数
    image_count = len(obj_name_ls) * 600

    # 数据集的保存路径
    dataset_basic_path = '/home/xm/workspace/ImageProcess/COCO/VOCdevkit/VOC2019/'
    image_folder = dataset_basic_path + 'JPEGImages/'
    #print(image_folder)
    label_folder = dataset_basic_path + 'Annotations/'
    #print(label_folder)
    image_set_folder = dataset_basic_path + 'ImageSets/Main/'
    #print(image_set_folder)

    for data_idx in range(image_count):
        # 获取 VOC 数据集中图像文件夹中所有文件的名称
        voc_folder_dir = '/home/xm/workspace/ImageProcess/VOC'
        voc_image_file_list = os.listdir(voc_folder_dir)
        #获取物体图像的文件名列表
        obj_image_ls_ls = []
        for obj_image_dir in obj_name_ls:
            cur_image_dir = base_obj_file_name + obj_image_dir
            obj_image_ls_ls.append(os.listdir(cur_image_dir))

        # 随机取一张 VOC 图做背景
        background_image_file = voc_image_file_list[np.random.randint(0, len(voc_image_file_list), 1)[-1]]
        background_image_file = voc_folder_dir + '/' + background_image_file
        background_image = cv2.imread(background_image_file)

        # 随机取若干物体
        obj_image_name_ls = []
        obj_cnt = np.random.randint(1, image_max_obj_cnt, 1)[-1]
        for obj_idx in range(obj_cnt):
            cur_obj_idx = np.random.randint(0, len(obj_image_ls_ls), 1)[-1]
            cur_obj_image_ls = obj_image_ls_ls[cur_obj_idx]
            cur_obj_file = cur_obj_image_ls[np.random.randint(0, len(cur_obj_image_ls), 1)[-1]]
            cur_obj_image = cv2.imread(base_obj_file_name + obj_name_ls[cur_obj_idx] + '/' + cur_obj_file)
            obj_image_name_ls.append([cur_obj_image, obj_name_ls[cur_obj_idx]])


        # 随机生成图像
        get_image, label_ls = formImageAndlabel(background_image, obj_image_name_ls, [max_size_radio, min_size_radio], iou_thres=0.05)
        #
        # # 保存图像与标签
        cur_image_name = str(data_idx).zfill(6) + '.jpg'
        #print(cur_image_name)
        cur_label_name = str(data_idx).zfill(6) + '.xml'
        #print(cur_label_name)

        cv2.imwrite(image_folder + cur_image_name, get_image)
        formImageLableXML(get_image, cur_image_name, label_ls, label_folder + cur_label_name)

        for obj_bbox, obj_name in label_ls:
            pnt_1 = tuple(map(int, obj_bbox[0:2]))
            pnt_2 = tuple(map(int, obj_bbox[0:2]))
            cv2.rectangle(get_image, pnt_1, pnt_2, (0, 0, 255))
        print(cur_image_name)

        cv2.imshow("get image", get_image)
        cv2.waitKey(10)
    train_set_name = 'train.txt'
    train_val_name = 'val.txt'
    test_set_name = 'test.txt'
    idx_thre = np.floor(0.6 * image_count)
    idx_thre_ = np.floor(0.8 * image_count)

    train_file = open(image_set_folder + train_set_name, 'w')
    for line_idx in range(int(idx_thre)):
        line_str = str(line_idx).zfill(6) + '\n'
        train_file.write(line_str)
    train_file.close()

    train_val_file = open(image_set_folder + train_val_name, 'w')
    for line_idx in range(int(idx_thre), int(idx_thre_)):
        line_str = str(line_idx).zfill(6) + '\n'
        train_val_file.write(line_str)
    train_val_file.close()

    test_file = open(image_set_folder + test_set_name, 'w')
    for line_idx in range(int(idx_thre_), image_count):
        line_str = str(line_idx).zfill(6) + '\n'
        test_file.write(line_str)
    test_file.close()


if __name__ == '__main__':
    main()
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
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316