import sys import cv2 import random import os from utils.util import * from imutils import paths

路径

img_dir = r'C:\Users\86157\Desktop\mtcnn\CCPD2020\ccpd_green\train' pos_save_dir = r'C:\Users\86157\Desktop\mtcnn\data\train\positive' part_save_dir = r'C:\Users\86157\Desktop\mtcnn\data\train\part' neg_save_dir = r'C:\Users\86157\Desktop\mtcnn\data\train\negative'

创建文件夹

for folder in [pos_save_dir, part_save_dir, neg_save_dir]: if not os.path.exists(folder): os.mkdir(folder)

存储正样本、负样本、部分样本的标签

f1 = open(os.path.join('data', 'pos_12_train.txt'), 'w') f2 = open(os.path.join('data', 'neg_12_train.txt'), 'w') f3 = open(os.path.join('data', 'part_12_train.txt'), 'w')

将img_dir目录下的所有图像文件打乱顺序，并输出打乱后的总数

img_paths = list(paths.list_images(img_dir)) random.shuffle(img_paths) num = len(img_paths) print('%d pics in total' % num)

用于记录正样本、负样本、部分样本的数量

p_idx, n_idx, d_idx = 0, 0, 0

遍历所有图像

for idx, im_path in enumerate(img_paths, start=1): print(im_path)

# 获取文件名和后缀名
basename = os.path.basename(im_path)
imgname, suffix = os.path.splitext(basename)
imgname_split = imgname.split('-')
rec_x1y1 = imgname_split[2].split('_')[0].split('&')
rec_x2y2 = imgname_split[2].split('_')[1].split('&')
x1, y1, x2, y2 = int(rec_x1y1[0]), int(rec_x1y1[1]), int(rec_x2y2[0]), int(rec_x2y2[1])

# 以box的形式存储bbox信息
boxes = np.zeros((1, 4), dtype=np.int32)
boxes[0, 0], boxes[0, 1], boxes[0, 2], boxes[0, 3] = x1, y1, x2, y2

# 读取图像
img = cv2.imread(im_path)

# 生成负样本
neg_num = 0
while neg_num < 35:
    size_x = np.random.randint(47, min(img.shape[:2]) / 2)
    size_y = np.random.randint(12, min(img.shape[:2]) / 2)
    nx = np.random.randint(0, img.shape[1] - size_x)
    ny = np.random.randint(0, img.shape[0] - size_y)
    crop_box = np.array([nx, ny, nx + size_x, ny + size_y])

    Iou = IoU(crop_box, boxes)

    cropped_im = img[ny: ny + size_y, nx: nx + size_x, :]
    resized_im = cv2.resize(cropped_im, (47, 12), interpolation=cv2.INTER_LINEAR)

    if np.max(Iou) < 0.3:
        # Iou with all gts must below 0.3
        save_file = os.path.join(neg_save_dir, "%s.jpg" % n_idx)
        f2.write(save_file + ' 0\n')
        cv2.imwrite(save_file, resized_im)
        n_idx += 1
        neg_num += 1

# 生成正样本和部分样本
for box in boxes:
    # box (x_left, y_top, w, h)
    x1, y1, x2, y2 = box
    w = x2 - x1 + 1
    h = y2 - y1 + 1

    # 生成负样本，确保和所有bbox的IoU都小于0.3
    for i in range(5):
        size_x = np.random.randint(47, min(img.shape[:2]) / 2)
        size_y = np.random.randint(12, min(img.shape[:2]) / 2)
        # delta_x and delta_y are offsets of (x1, y1)
        delta_x = np.random.randint(max(-size_x, -x1), w)
        delta_y = np.random.randint(max(-size_y, -y1), h)
        nx1 = max(0, x1 + delta_x)
        ny1 = max(0, y1 + delta_y)

        if nx1 + size_x > img.shape[1] or ny1 + size_y > img.shape[0]:
            continue
        crop_box = np.array([nx1, ny1, nx1 + size_x, ny1 + size_y])
        Iou = IoU(crop_box, boxes)

        cropped_im = img[ny1: ny1 + size_y, nx1: nx1 + size_x, :]
        resized_im = cv2.resize(cropped_im, (47, 12), interpolation=cv2.INTER_LINEAR)

        if np.max(Iou) < 0.3:
            # Iou with all gts must below 0.3
            save_file = os.path.join(neg_save_dir, "%s.jpg" % n_idx)
            f2.write(save_file + ' 0\n')
            cv2.imwrite(save_file, resized_im)
            n_idx += 1
    # 生成正样本和部分样本
    for i in range(20):
        size_x = np.random.randint(int(min(w, h) * 0.8), np.ceil(1.25 * max(w, h)))
        size_y = np.random.randint(int(min(w, h) * 0.8), np.ceil(1.25 * max(w, h)))

        # delta here is the offset of box center
        delta_x = np.random.randint(-w * 0.2, w * 0.2)
        delta_y = np.random.randint(-h * 0.2, h * 0.2)

        nx1 = max(x1 + w / 2 + delta_x - size_x / 2, 0)
        ny1 = max(y1 + h / 2 + delta_y - size_y / 2, 0)
        nx2 = nx1 + size_x
        ny2 = ny1 + size_y

        if nx2 > img.shape[1] or ny2 > img.shape[0]:
            continue
        crop_box = np.array([nx1, ny1, nx2, ny2])

        offset_x1 = (x1 - nx1) / float(size_x)
        offset_y1 = (y1 - ny1) / float(size_y)
        offset_x2 = (x2 - nx2) / float(size_x)
        offset_y2 = (y2 - ny2) / float(size_y)

        cropped_im = img[int(ny1): int(ny2), int(nx1): int(nx2), :]
        resized_im = cv2.resize(cropped_im, (47, 12), interpolation=cv2.INTER_LINEAR)

        box_ = box.reshape(1, -1)
        if IoU(crop_box, box_) >= 0.65:
            save_file = os.path.join(pos_save_dir, "%s.jpg" % p_idx)
            f1.write(save_file + ' 1 %.2f %.2f %.2f %.2f\n' % (offset_x1, offset_y1, offset_x2, offset_y2))
            cv2.imwrite(save_file, resized_im)
            p_idx += 1
        elif IoU(crop_box, box_) >= 0.4 and d_idx < 1.2 * p_idx + 1:
            save_file = os.path.join(part_save_dir, "%s.jpg" % d_idx)
            f3.write(save_file + ' -1 %.2f %.2f %.2f %.2f\n' % (offset_x1, offset_y1, offset_x2, offset_y2))
            cv2.imwrite(save_file, resized_im)
            d_idx += 1

# 输出当前处理进度
print("%s images done, pos: %s part: %s neg: %s" % (idx, p_idx, d_idx, n_idx))

关闭文件

for f in [f1, f2, f3]: f.close(