yolov3目标检测数据增强和训练的详细过程

yolov3的编译,数据增强、训练、测试

yoloV3的官网 https://pjreddie.com/darknet/yolo/

一、编译和训练

1、修改Makefile文件,修改GPU、cudnn、opencv三个参数,调整为1。

2、然后这里如果需要修改保存模型保存的次数,比如想要1000次保存一次模型,修改为100次保存一下模型,这个时候需要修改sample文件夹中的detector.c文件,将其中的第138行修改一下即可。(这里看一下就知道怎么修改了了)

3、如果报了opencv的错误,把src中的image_opencv.cpp文件中mat_to_image函数修改为下面的部分:

image mat_to_image(Mat m)
{
    //IplImage ipl = m;
    IplImage ipl = cvIplImage(m);
    image im = ipl_to_image(&ipl);
    rgbgr_image(im);
    return im;
}

4、接下来在darknet窗口下打开一个命令窗口,进行make,编译一下即可。

5、接着可以测试一下yolov3是不是能用,下载预训练权重文件 https://pjreddie.com/media/files/yolov3.weights,然后测试命令如下:

./darknet detect cfg/yolov3.cfg yolov3.weights data/dog.jpg

6、测试可用以后,准备数据集,数据集的结构如下:
VOCdevkit/VOC2007/Annotations JPEGImages

其中Annotations放标好的标签为xml格式,JPEGImages放准备好的图片,另外那两个文件夹是运行脚本gen_files之后生成的,先不用管。

准备好数据集之后,在darknet文件夹下存放一个gen_files的脚本,脚本里面的classes列表里面需要修改成自己的数据集的标签。之后在darknet下打开终端运行

python gen_files

其中gen_files如下:

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import random

classes=["1","2","3","4","5","6","7","8","9","10","12","13"]


def clear_hidden_files(path):
    dir_list = os.listdir(path)
    for i in dir_list:
        abspath = os.path.join(os.path.abspath(path), i)
        if os.path.isfile(abspath):
            if i.startswith("._"):
                os.remove(abspath)
        else:
            clear_hidden_files(abspath)

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

def convert_annotation(image_id):
    in_file = open('VOCdevkit/VOC2007/Annotations/%s.xml' %image_id)
    out_file = open('VOCdevkit/VOC2007/labels/%s.txt' %image_id, 'w')
    tree=ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
        bb = convert((w,h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
    in_file.close()
    out_file.close()

wd = os.getcwd()
wd = os.getcwd()
work_sapce_dir = os.path.join(wd, "VOCdevkit/")
if not os.path.isdir(work_sapce_dir):
    os.mkdir(work_sapce_dir)
work_sapce_dir = os.path.join(work_sapce_dir, "VOC2007/")
if not os.path.isdir(work_sapce_dir):
    os.mkdir(work_sapce_dir)
annotation_dir = os.path.join(work_sapce_dir, "Annotations/")
if not os.path.isdir(annotation_dir):
        os.mkdir(annotation_dir)
clear_hidden_files(annotation_dir)
image_dir = os.path.join(work_sapce_dir, "JPEGImages/")
if not os.path.isdir(image_dir):
        os.mkdir(image_dir)
clear_hidden_files(image_dir)
VOC_file_dir = os.path.join(work_sapce_dir, "ImageSets/")
if not os.path.isdir(VOC_file_dir):
        os.mkdir(VOC_file_dir)
VOC_file_dir = os.path.join(VOC_file_dir, "Main/")
if not os.path.isdir(VOC_file_dir):
        os.mkdir(VOC_file_dir)

train_file = open(os.path.join(wd, "2007_train.txt"), 'w')
test_file = open(os.path.join(wd, "2007_test.txt"), 'w')
train_file.close()
test_file.close()
VOC_train_file = open(os.path.join(work_sapce_dir, "ImageSets/Main/train.txt"), 'w')
VOC_test_file = open(os.path.join(work_sapce_dir, "ImageSets/Main/test.txt"), 'w')
VOC_train_file.close()
VOC_test_file.close()
if not os.path.exists('VOCdevkit/VOC2007/labels'):
    os.makedirs('VOCdevkit/VOC2007/labels')
train_file = open(os.path.join(wd, "2007_train.txt"), 'a')
test_file = open(os.path.join(wd, "2007_test.txt"), 'a')
VOC_train_file = open(os.path.join(work_sapce_dir, "ImageSets/Main/train.txt"), 'a')
VOC_test_file = open(os.path.join(work_sapce_dir, "ImageSets/Main/test.txt"), 'a')
list = os.listdir(image_dir) # list image files
probo = random.randint(1, 100)
print("Probobility: %d" % probo)
for i in range(0,len(list)):
    path = os.path.join(image_dir,list[i])
    if os.path.isfile(path):
        image_path = image_dir + list[i]
        voc_path = list[i]
        (nameWithoutExtention, extention) = os.path.splitext(os.path.basename(image_path))
        (voc_nameWithoutExtention, voc_extention) = os.path.splitext(os.path.basename(voc_path))
        annotation_name = nameWithoutExtention + '.xml'
        annotation_path = os.path.join(annotation_dir, annotation_name)
    probo = random.randint(1, 100)
    print("Probobility: %d" % probo)
    if(probo < 75):
        if os.path.exists(annotation_path):
            train_file.write(image_path + '\n')
            VOC_train_file.write(voc_nameWithoutExtention + '\n')
            convert_annotation(nameWithoutExtention)
    else:
        if os.path.exists(annotation_path):
            test_file.write(image_path + '\n')
            VOC_test_file.write(voc_nameWithoutExtention + '\n')
            convert_annotation(nameWithoutExtention)
train_file.close()
test_file.close()
VOC_train_file.close()
VOC_test_file.close()

在VOCdevkit / VOC2007目录下可以看到生成了文件夹labels ,同时在darknet下生成了两个文件2007_train.txt和2007_test.txt。2007_train.txt和2007_test.txt分

别给出了训练图片文件和测试图片文件的列表,含有每个图片的路径和文件名。另外,在VOCdevkit / VOC2007/ImageSets/Main目录下生成了两个文件test.txt和

train.txt,分别给出了训练图片文件和测试图片文件的列表,但只含有每个图片的文件名(不含路径和扩展名)。labels下的文件是JPEGImages文件夹下每一个图像的

yolo格式的标注文件,这是由Annotations的xml标注文件转换来的。最终训练只需要:2007_train.txt,2007_test.txt,labels下的标注文件和 VOCdevkit

/VOC2007/JPEGImages下的图像文件。

7、修改配置文件

1)新建data/voc.names文件

可以复制data/voc.names再根据自己情况的修改;可以重新命名如:data/voc-ball.names

2)新建 cfg/voc.data文件

可以复制cfg/voc.data再根据自己情况的修改;可以重新命名如:cfg/voc-ball.data

3)新建cfg/yolov3-voc.cfg

可以复制cfg/yolov3-voc.cfg再根据自己情况的修改;可以重新命名cfg/yolov3-voc-ball.cfg

在cfg/yolov3-voc.cfg文件中,三个yolo层和各自前面的conv层的参数需要修改:

三个yolo层都要改:yolo层中的class为类别数,每一个yolo层前的conv层中的filters =(类别+5) 3*

例如:

yolo层 classes=1, conv层 filters=18

yolo层 classes=2, conv层 filters=21

yolo层 classes=4,conv层 filters=27

这里的改变如下图所示:

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-oYlfWvk0-1608043513202)(/home/wl/.config/Typora/typora-user-images/image-20201214193720130.png)]

这里的classes类别为12类的话,filters就需要修改为(class_num+5)*3

在darknet目录下面放入预训练模型:https://pjreddie.com/media/files/darknet53.conv.74

8、进行训练

修改好之后就可以使用下面的命令进行训练了:

./darknet detector train cfg/voc-ball.data cfg/yolov3-voc-ball.cfg darknet53.conv.74

如需要存储训练日志,执行

./darknet detector train cfg/voc-ball.data cfg/yolov3-voc-ball.cfg darknet53.conv.74 2>1 | tee visualization/train_yolov3_ball.log

执行前需要建立visualizaton/train_logs即可

注:如果是yolov4,上述的步骤是一样的,只不过这里的测试命令变为

./darknet detector train cfg/voc-ball.data cfg/yolov4-voc-ball.cfg yolov4.conv.137 -map

9、训练建议

batch=16

subdivisions=8

把max_batches设置为 (classes*2000);但最小为4000。例如如果训练3个目标类别,

把max_batches设置为 (classes*2000);但最小为4000。例如如果训练3个目标类别,

把steps改为max_batches的80% and 90%;例如steps=4800, 5400。

为增加网络分辨率可增大height和width的值,但必须是32的倍数 (height=608, width=608 or 32的整数倍) 。这有助于提高检测精度。

10、测试图片

./darknet detector test cfg/voc-ball.data cfg/yolov4-voc-ball-test.cfg backup/yolov4-voc-ball_final.weights testfiles/img1.jpg

测试视频

./darknet detector demo cfg/voc-ball.data cfg/yolov4-voc-ball-test.cfg backup/yolov4-voc-ball_final.weights testfiles/messi.mp4

11、性能统计

  1. 计算mAP方法1
    统计 mAP@IoU=0.50:
./darknet detector map cfg/voc-ball.data cfg/yolov4-voc-ball-test.cfg backup/yolov4-voc-ball_final.weights  

统计 mAP@IoU=0.75:

./darknet detector map cfg/voc-ball.data cfg/yolov4-voc-ball-test.cfg backup/yolov4-voc-ball_final.weights -iou_thresh 0.75
  1. 计算mAP方法2
    首先执行
./darknet detector valid cfg/voc-ball.data cfg/yolov4-voc-ball-test.cfg backup/yolov4-voc-ball_final.weights  

生成results/comp4_det_test_ball.txt文件
然后执行

python reval_voc.py --voc_dir VOCdevkit --year 2007 --image_set test --classes data/voc-ball.names testball

生成testball/ball_pr.pkl文件

画出PR曲线

然后可画出PR曲线,

修改文件draw_pr.py

fr = open(‘testball/ball_pr.pkl’,‘rb’)
执行

python draw_pr.py

12、Anchor Box先验框聚类分析与修改

  1. 使用k-means聚类获得自己数据集的先验框大小
./darknet detector calc_anchors cfg/voc-ball.data -num_of_clusters 9 -width 608 -height 608
  1. 修改cfg文件中的先验框大小

  2. 重新训练和测试

二、数据增强

这里直接上华科大佬许炳伟的代码

# -*- coding=utf-8 -*-

# 包括:
#     1. 裁剪(需改变bbox)
#     2. 平移(需改变bbox)
#     3. 改变亮度
#     4. 加噪声
#     5. 旋转角度(需要改变bbox)
#     6. 镜像(需要改变bbox)
#     7. cutout
# 注意:
#     random.seed(),相同的seed,产生的随机数是一样的!!
import sys

ros_path = '/opt/ros/kinetic/lib/python2.7/dist-packages'

if ros_path in sys.path:
    sys.path.remove(ros_path)

import cv2
import time
import random
import os
import math
import numpy as np
from skimage.util import random_noise
from skimage import exposure


# 显示带标签显示的图片
def show_pic(img, bboxes=None, labels=None):
    '''
    输入:
        img:图像array
        bboxes:图像的所有boudning box list, 格式为[[x_min, y_min, x_max, y_max]....]
        names:每个box对应的名称
    '''
    #     cv2.imwrite('./1.jpg', img)
    #     img = cv2.imread('./1.jpg')
    img = img / 255
    for i in range(len(bboxes)):
        bbox = bboxes[i]
        x_min = bbox[0]
        y_min = bbox[1]
        x_max = bbox[2]
        y_max = bbox[3]
        cv2.rectangle(img, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 3)
        cv2.putText(img, labels[i], (int(x_min), int(y_min)), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
    cv2.namedWindow('pic', 0)  # 1表示原图
    cv2.moveWindow('pic', 0, 0)
    cv2.resizeWindow('pic', 1200, 800)  # 可视化的图片大小
    cv2.imshow('pic', img)
    if cv2.waitKey(1) == ord('q'):
        cv2.destroyAllWindows()
        sys.exit()


#     cv2.destroyAllWindows()
#     os.remove('./1.jpg')

# 图像均为cv2读取
class DataAugmentForObjectDetection():
    def __init__(self, rotation_rate=0.5, max_rotation_angle=30,
                 crop_rate=0.5, shift_rate=0.5, change_light_rate=0.5,
                 add_noise_rate=0.5, flip_rate=0.5,
                 cutout_rate=0.5, cut_out_length=50, cut_out_holes=1, cut_out_threshold=0.5):
        self.rotation_rate = rotation_rate
        self.max_rotation_angle = max_rotation_angle
        self.crop_rate = crop_rate
        self.shift_rate = shift_rate
        self.change_light_rate = change_light_rate
        self.add_noise_rate = add_noise_rate
        self.flip_rate = flip_rate
        self.cutout_rate = cutout_rate

        self.cut_out_length = cut_out_length
        self.cut_out_holes = cut_out_holes
        self.cut_out_threshold = cut_out_threshold

    # 加噪声
    def _addNoise(self, img):
        '''
        输入:
            img:图像array
        输出:
            加噪声后的图像array,由于输出的像素是在[0,1]之间,所以得乘以255
        '''
        # random.seed(int(time.time()))
        # return random_noise(img, mode='gaussian', seed=int(time.time()), clip=True)*255
        return random_noise(img, mode='gaussian', clip=True) * 255

    # 调整亮度
    def _changeLight(self, img):
        # random.seed(int(time.time()))
        flag = random.uniform(0.5, 1.5)  # flag>1为调暗,小于1为调亮
        return exposure.adjust_gamma(img, flag)

    # cutout
    def _cutout(self, img, bboxes, length=100, n_holes=1, threshold=0.5):
        '''
        原版本:https://github.com/uoguelph-mlrg/Cutout/blob/master/util/cutout.py
        Randomly mask out one or more patches from an image.
        Args:
            img : a 3D numpy array,(h,w,c)
            bboxes : 框的坐标
            n_holes (int): Number of patches to cut out of each image.
            length (int): The length (in pixels) of each square patch.
        '''

        def cal_iou(boxA, boxB):
            '''
            boxA, boxB为两个框,返回iou
            boxB为bouding box
            '''

            # determine the (x, y)-coordinates of the intersection rectangle
            xA = max(boxA[0], boxB[0])
            yA = max(boxA[1], boxB[1])
            xB = min(boxA[2], boxB[2])
            yB = min(boxA[3], boxB[3])

            if xB <= xA or yB <= yA:
                return 0.0

            # compute the area of intersection rectangle
            interArea = (xB - xA + 1) * (yB - yA + 1)

            # compute the area of both the prediction and ground-truth
            # rectangles
            boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
            boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)

            # compute the intersection over union by taking the intersection
            # area and dividing it by the sum of prediction + ground-truth
            # areas - the interesection area
            # iou = interArea / float(boxAArea + boxBArea - interArea)
            iou = interArea / float(boxBArea)

            # return the intersection over union value
            return iou

        # 得到h和w
        if img.ndim == 3:
            h, w, c = img.shape
        else:
            _, h, w, c = img.shape

        mask = np.ones((h, w, c), np.float32)

        for n in range(n_holes):

            chongdie = True  # 看切割的区域是否与box重叠太多

            while chongdie:
                y = np.random.randint(h)
                x = np.random.randint(w)

                y1 = np.clip(y - length // 2, 0,
                             h)  # numpy.clip(a, a_min, a_max, out=None), clip这个函数将将数组中的元素限制在a_min, a_max之间,大于a_max的就使得它等于 a_max,小于a_min,的就使得它等于a_min
                y2 = np.clip(y + length // 2, 0, h)
                x1 = np.clip(x - length // 2, 0, w)
                x2 = np.clip(x + length // 2, 0, w)

                chongdie = False
                for box in bboxes:
                    if cal_iou([x1, y1, x2, y2], box) > threshold:
                        chongdie = True
                        break

            mask[y1: y2, x1: x2, :] = 0.

        # mask = np.expand_dims(mask, axis=0)
        img = img * mask

        return img

    # 旋转
    def _rotate_img_bbox(self, img, bboxes, angle=5, scale=1.):
        '''
        参考:https://blog.csdn.net/u014540717/article/details/53301195crop_rate
        输入:
            img:图像array,(h,w,c)
            bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
            angle:旋转角度
            scale:默认1
        输出:
            rot_img:旋转后的图像array
            rot_bboxes:旋转后的boundingbox坐标list
        '''
        # ---------------------- 旋转图像 ----------------------
        w = img.shape[1]
        h = img.shape[0]
        # 角度变弧度
        rangle = np.deg2rad(angle)  # angle in radians
        # now calculate new image width and height
        nw = (abs(np.sin(rangle) * h) + abs(np.cos(rangle) * w)) * scale
        nh = (abs(np.cos(rangle) * h) + abs(np.sin(rangle) * w)) * scale
        # ask OpenCV for the rotation matrix
        rot_mat = cv2.getRotationMatrix2D((nw * 0.5, nh * 0.5), angle, scale)
        # calculate the move from the old center to the new center combined
        # with the rotation
        rot_move = np.dot(rot_mat, np.array([(nw - w) * 0.5, (nh - h) * 0.5, 0]))
        # the move only affects the translation, so update the translation
        # part of the transform
        rot_mat[0, 2] += rot_move[0]
        rot_mat[1, 2] += rot_move[1]
        # 仿射变换
        rot_img = cv2.warpAffine(img, rot_mat, (int(math.ceil(nw)), int(math.ceil(nh))), flags=cv2.INTER_LANCZOS4)

        # ---------------------- 矫正bbox坐标 ----------------------
        # rot_mat是最终的旋转矩阵
        # 获取原始bbox的四个中点,然后将这四个点转换到旋转后的坐标系下
        rot_bboxes = list()
        for bbox in bboxes:
            xmin = bbox[0]
            ymin = bbox[1]
            xmax = bbox[2]
            ymax = bbox[3]
            point1 = np.dot(rot_mat, np.array([(xmin + xmax) / 2, ymin, 1]))
            point2 = np.dot(rot_mat, np.array([xmax, (ymin + ymax) / 2, 1]))
            point3 = np.dot(rot_mat, np.array([(xmin + xmax) / 2, ymax, 1]))
            point4 = np.dot(rot_mat, np.array([xmin, (ymin + ymax) / 2, 1]))
            # 合并np.array
            concat = np.vstack((point1, point2, point3, point4))
            # 改变array类型
            concat = concat.astype(np.int32)
            # 得到旋转后的坐标
            rx, ry, rw, rh = cv2.boundingRect(concat)
            rx_min = rx
            ry_min = ry
            rx_max = rx + rw
            ry_max = ry + rh
            # 加入list中
            rot_bboxes.append([rx_min, ry_min, rx_max, ry_max])

        return rot_img, rot_bboxes

    # 裁剪
    def _crop_img_bboxes(self, img, bboxes):
        '''
        裁剪后的图片要包含所有的框
        输入:
            img:图像array
            bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
        输出:
            crop_img:裁剪后的图像array
            crop_bboxes:裁剪后的bounding box的坐标list
        '''
        # ---------------------- 裁剪图像 ----------------------
        w = img.shape[1]
        h = img.shape[0]
        x_min = w  # 裁剪后的包含所有目标框的最小的框
        x_max = 0
        y_min = h
        y_max = 0
        for bbox in bboxes:
            x_min = min(x_min, bbox[0])
            y_min = min(y_min, bbox[1])
            x_max = max(x_max, bbox[2])
            y_max = max(y_max, bbox[3])

        d_to_left = x_min  # 包含所有目标框的最小框到左边的距离
        d_to_right = w - x_max  # 包含所有目标框的最小框到右边的距离
        d_to_top = y_min  # 包含所有目标框的最小框到顶端的距离
        d_to_bottom = h - y_max  # 包含所有目标框的最小框到底部的距离

        # 随机扩展这个最小框
        crop_x_min = int(x_min - random.uniform(0, d_to_left))
        crop_y_min = int(y_min - random.uniform(0, d_to_top))
        crop_x_max = int(x_max + random.uniform(0, d_to_right))
        crop_y_max = int(y_max + random.uniform(0, d_to_bottom))

        # 随机扩展这个最小框 , 防止别裁的太小
        # crop_x_min = int(x_min - random.uniform(d_to_left//2, d_to_left))
        # crop_y_min = int(y_min - random.uniform(d_to_top//2, d_to_top))
        # crop_x_max = int(x_max + random.uniform(d_to_right//2, d_to_right))
        # crop_y_max = int(y_max + random.uniform(d_to_bottom//2, d_to_bottom))

        # 确保不要越界
        crop_x_min = max(0, crop_x_min)
        crop_y_min = max(0, crop_y_min)
        crop_x_max = min(w, crop_x_max)
        crop_y_max = min(h, crop_y_max)

        crop_img = img[crop_y_min:crop_y_max, crop_x_min:crop_x_max]

        # ---------------------- 裁剪boundingbox ----------------------
        # 裁剪后的boundingbox坐标计算
        crop_bboxes = list()
        for bbox in bboxes:
            crop_bboxes.append([bbox[0] - crop_x_min, bbox[1] - crop_y_min, bbox[2] - crop_x_min, bbox[3] - crop_y_min])

        return crop_img, crop_bboxes

    # 平移
    def _shift_pic_bboxes(self, img, bboxes):
        '''
        参考:https://blog.csdn.net/sty945/article/details/79387054
        平移后的图片要包含所有的框
        输入:
            img:图像array
            bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
        输出:
            shift_img:平移后的图像array
            shift_bboxes:平移后的bounding box的坐标list
        '''
        # ---------------------- 平移图像 ----------------------
        w = img.shape[1]
        h = img.shape[0]
        x_min = w  # 裁剪后的包含所有目标框的最小的框
        x_max = 0
        y_min = h
        y_max = 0
        for bbox in bboxes:
            x_min = min(x_min, bbox[0])
            y_min = min(y_min, bbox[1])
            x_max = max(x_max, bbox[2])
            y_max = max(y_max, bbox[3])

        d_to_left = x_min  # 包含所有目标框的最大左移动距离
        d_to_right = w - x_max  # 包含所有目标框的最大右移动距离
        d_to_top = y_min  # 包含所有目标框的最大上移动距离
        d_to_bottom = h - y_max  # 包含所有目标框的最大下移动距离

        x = random.uniform(-(d_to_left - 1) / 3, (d_to_right - 1) / 3)
        y = random.uniform(-(d_to_top - 1) / 3, (d_to_bottom - 1) / 3)

        M = np.float32([[1, 0, x], [0, 1, y]])
        # x为向左或右移动的像素值,正为向右负为向左; y为向上或者向下移动的像素值,正为向下负为向上
        try:
            shift_img = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]))
        except Exception as e:
            print("error")

        # ---------------------- 平移boundingbox ----------------------
        shift_bboxes = list()
        for bbox in bboxes:
            shift_bboxes.append([bbox[0] + x, bbox[1] + y, bbox[2] + x, bbox[3] + y])

        return shift_img, shift_bboxes

    # 镜像
    def _filp_pic_bboxes(self, img, bboxes):
        '''
            参考:https://blog.csdn.net/jningwei/article/details/78753607
            平移后的图片要包含所有的框
            输入:
                img:图像array
                bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
            输出:
                flip_img:平移后的图像array
                flip_bboxes:平移后的bounding box的坐标list
        '''
        # ---------------------- 翻转图像 ----------------------
        import copy
        flip_img = copy.deepcopy(img)
        #         if random.random() < 0.5:    #0.5的概率水平翻转,0.5的概率垂直翻转
        horizon = True
        #         else:
        #             horizon = False
        h, w, _ = img.shape
        if horizon:  # 水平翻转
            flip_img = cv2.flip(flip_img, 1)  # 1是水平,-1是水平垂直
        else:
            flip_img = cv2.flip(flip_img, 0)

        # ---------------------- 调整boundingbox ----------------------
        flip_bboxes = list()
        for box in bboxes:
            x_min = box[0]
            y_min = box[1]
            x_max = box[2]
            y_max = box[3]
            if horizon:
                flip_bboxes.append([w - x_max, y_min, w - x_min, y_max])
            else:
                flip_bboxes.append([x_min, h - y_max, x_max, h - y_min])

        return flip_img, flip_bboxes

    def dataAugment(self, img, bboxes):
        '''
        图像增强
        输入:
            img:图像array
            bboxes:该图像的所有框坐标
        输出:
            img:增强后的图像
            bboxes:增强后图片对应的box
        '''
        change_num = 0  # 改变的次数
        print('------')
        while change_num < 1:  # 默认至少有一种数据增强生效
            if random.random() < self.crop_rate:  # 裁剪
                print('裁剪')
                change_num += 1
                img, bboxes = self._crop_img_bboxes(img, bboxes)

            if random.random() > self.rotation_rate:  # 旋转
                print('旋转')
                change_num += 1
                angle = random.uniform(-self.max_rotation_angle, self.max_rotation_angle)
                #                 angle = random.sample([90, 180, 270],1)[0]
                scale = random.uniform(0.7, 0.8)
                img, bboxes = self._rotate_img_bbox(img, bboxes, angle, scale)

            if random.random() < self.shift_rate:  # 平移
                print('平移')
                change_num += 1
                img, bboxes = self._shift_pic_bboxes(img, bboxes)

            if random.random() > self.change_light_rate:  # 改变亮度
                print('亮度')
                change_num += 1
                img = self._changeLight(img)

            if random.random() < self.add_noise_rate:  # 加噪声
                print('加噪声')
                change_num += 1
                img = self._addNoise(img)

            #             if random.random() < self.cutout_rate:  #cutout
            #                 print('cutout')
            #                 change_num += 1
            #                 img = self._cutout(img, bboxes, length=self.cut_out_length, n_holes=self.cut_out_holes, threshold=self.cut_out_threshold)

            #             if random.random() < self.flip_rate:    #翻转
            #                 print('翻转')
            #                 change_num += 1
            #                 img, bboxes = self._filp_pic_bboxes(img, bboxes)
            print('\n')
        # print('------')
        return img, bboxes



# -*- coding=utf-8 -*-
import xml.etree.ElementTree as ET
import xml.dom.minidom as DOC

# 从xml文件中提取bounding box信息, 格式为[[x_min, y_min, x_max, y_max, name]]
def parse_xml(xml_path):
    '''
    输入:
        xml_path: xml的文件路径
    输出:
        从xml文件中提取bounding box信息, 格式为[[x_min, y_min, x_max, y_max, name]]
    '''
    tree = ET.parse(xml_path)
    root = tree.getroot()
    objs = root.findall('object')
    coords = list()
    for ix, obj in enumerate(objs):
        name = obj.find('name').text
        box = obj.find('bndbox')
        x_min = int(float(box[0].text))
        y_min = int(float(box[1].text))
        x_max = int(float(box[2].text))
        y_max = int(float(box[3].text))
        coords.append([x_min, y_min, x_max, y_max, name])
    return coords


import os
from lxml.etree import Element, SubElement, tostring
from xml.dom.minidom import parseString
from PIL import Image


# 保存xml文件函数的核心实现,输入为图片名称image_name,分类category(一个列表,元素与bbox对应),bbox(一个列表,与分类对应),保存路径save_dir ,通道数channel
def save_xml(image_name, category, bbox, file_dir='/home/xbw/wurenting/dataset_3/',
             save_dir='/home/xxx/voc_dataset/Annotations/', channel=3):
    file_path = file_dir
    img = Image.open(file_path + image_name)
    width = img.size[0]
    height = img.size[1]

    node_root = Element('annotation')

    node_folder = SubElement(node_root, 'folder')
    node_folder.text = 'VOC2007'

    node_filename = SubElement(node_root, 'filename')
    node_filename.text = image_name

    node_size = SubElement(node_root, 'size')
    node_width = SubElement(node_size, 'width')
    node_width.text = '%s' % width

    node_height = SubElement(node_size, 'height')
    node_height.text = '%s' % height

    node_depth = SubElement(node_size, 'depth')
    node_depth.text = '%s' % channel

    for i in range(len(bbox)):
        left, top, right, bottom = bbox[i][0], bbox[i][1], bbox[i][2], bbox[i][3]
        node_object = SubElement(node_root, 'object')
        node_name = SubElement(node_object, 'name')
        node_name.text = category[i]
        node_difficult = SubElement(node_object, 'difficult')
        node_difficult.text = '0'
        node_bndbox = SubElement(node_object, 'bndbox')
        node_xmin = SubElement(node_bndbox, 'xmin')
        node_xmin.text = '%s' % left
        node_ymin = SubElement(node_bndbox, 'ymin')
        node_ymin.text = '%s' % top
        node_xmax = SubElement(node_bndbox, 'xmax')
        node_xmax.text = '%s' % right
        node_ymax = SubElement(node_bndbox, 'ymax')
        node_ymax.text = '%s' % bottom

    xml = tostring(node_root, pretty_print=True)
    dom = parseString(xml)

    save_xml = os.path.join(save_dir, image_name.replace('jpg', 'xml'))
    with open(save_xml, 'wb') as f:
        f.write(xml)

    return


import shutil

need_aug_num = 1

dataAug = DataAugmentForObjectDetection()

source_pic_root_path = '/home/wl/import/last_data/VOCdevkit/VOC2007/JPEGImages/'
source_xml_root_path = '/home/wl/import/last_data/VOCdevkit/VOC2007/Annotations/'
img_save_path = '/home/wl/import/last_data/VOCdevkit/VOC2007/aug_img/'
save_dir = '/home/wl/import/last_data/VOCdevkit/VOC2007/aug_label/'

for parent, _, files in os.walk(source_pic_root_path):
    for file in files:
        cnt = 0
        while cnt < need_aug_num:
            pic_path = os.path.join(parent, file)
            xml_path = os.path.join(source_xml_root_path, file[:-4]+'.xml')
            coords = parse_xml(xml_path)        #解析得到box信息,格式为[[x_min,y_min,x_max,y_max,name]]
            coordss = [coord[:4] for coord in coords]
            labels = [coord[4] for coord in coords]
            img = cv2.imread(pic_path)
            show_pic(img, coordss,labels)    # 原图

            auged_img, auged_bboxes = dataAug.dataAugment(img, coordss)
            cnt += 1
            cv2.imwrite(img_save_path+file[:-4]+'_1.jpg',auged_img)
            save_xml(file[:-4]+'_1.jpg',labels,auged_bboxes,file_dir = img_save_path,save_dir=save_dir)
            show_pic(auged_img, auged_bboxes,labels)  # 强化后的图
cv2.destroyAllWindows()






#测试label是否正确
import shutil

# need_aug_num = 1
#
# dataAug = DataAugmentForObjectDetection()
#
# source_pic_root_path = '/home/xbw/darknet_boat/darknet/scripts/VOCdevkit/VOC2007/add_990/990_add/'
# source_xml_root_path = '/home/xbw/darknet_boat/darknet/scripts/VOCdevkit/VOC2007/add_990/990_xml/'
#
# for parent, _, files in os.walk(source_pic_root_path):
#     for file in files:
#         cnt = 0
#         while cnt < need_aug_num:
#             pic_path = os.path.join(parent, file)
#             xml_path = os.path.join(source_xml_root_path, file[:-4]+'.xml')
#             coords = parse_xml(xml_path)        #解析得到box信息,格式为[[x_min,y_min,x_max,y_max,name]]
#             coordss = [coord[:4] for coord in coords]
#             labels = [coord[4] for coord in coords]
#             img = cv2.imread(pic_path)
#             show_pic(img, coordss,labels)    # 原图
#             cnt += 1
# cv2.destroyAllWindows()

最后yolo也是可以断点接着进行训练的,只需要在命令里面把那个预训练模型的命令修改为自己上次训练的模型即可。

猜你喜欢

转载自blog.csdn.net/qq_32651847/article/details/111242484