龙空技术网

基于yolo v4的自动化标注软件实现——附源码

自动驾驶前沿 1500

前言:

此时小伙伴们对“yolov4源码解读”都比较重视,姐妹们都想要分析一些“yolov4源码解读”的相关知识。那么小编也在网络上搜集了一些对于“yolov4源码解读””的相关内容,希望咱们能喜欢,朋友们一起来学习一下吧!

传统的标注,需要使用标注软件人工一个一个的标注,太耗费时间和精力,如果我们有个基础的、通用的模型,但是效果一般,然后想要重新标注,我们就可以使用现有的模型,生成VOC格式的标签文件。

原始图片

生成的VOC格式的标签

1、Automatic_labeling.py

"""  "********************************************************************************************文件名称 :Automatic_labeling.py*文件功能 :基于yolo v4的 自动标注软件版本:1.0内容:基于yolo v4的 自动标注软件,验证成功时间:2021.10.27作者:狄云********************************************************************************************"""import tensorflow as tfconfig = tf.compat.v1.ConfigProto(gpu_options=tf.compat.v1.GPUOptions(allow_growth=True))sess = tf.compat.v1.Session(config=config)# 1、导入相关包和库import cv2 as cvfrom PIL import Imageimport numpy as npimport osimport sysfrom yolo import YOLO #from yolo import YOLO, detect_videoyolo = YOLO()class OBJECT:  # Bounding box + annotations    def __init__(self, bb, classname):        self.xmin = int(round(bb[0]))#对浮点数进行近似,无第二个参数:取整        self.ymin = int(round(bb[1]))        self.xmax = int(round(bb[2]))        self.ymax = int(round(bb[3]))        self.classname = classname# VOC数据集格式xml_body_1 = """<annotation>        <folder>FOLDER</folder>        <filename>{FILENAME}</filename>        <path>{PATH}</path>        <source>                <database>Unknown</database>        </source>        <size>                <width>{WIDTH}</width>                <height>{HEIGHT}</height>                <depth>3</depth>        </size>        <segmented>0</segmented>"""xml_object = """ <object>                <name>{CLASS}</name>                <pose>Unspecified</pose>                <truncated>0</truncated>                <difficult>0</difficult>                <bndbox>                        <xmin>{XMIN}</xmin>                        <ymin>{YMIN}</ymin>                        <xmax>{XMAX}</xmax>                        <ymax>{YMAX}</ymax>                </bndbox>        </object>"""xml_body_2 = """</annotation>        """#创建VOC文件def create_voc_xml(xml_file, img_file,imgW, imgH, objects , display=False):    with open(xml_file, "w") as f:        f.write(xml_body_1.format(            **{'FILENAME': os.path.basename(img_file), 'PATH': img_file, 'WIDTH': imgW, 'HEIGHT': imgH}))        # for bba in listbba:        #     f.write(xml_object.format(        #         **{'CLASS': bba.classname, 'XMIN': bba.x1, 'YMIN': bba.y1, 'XMAX': bba.x2, 'YMAX': bba.y2}))        for object in objects:            f.write(xml_object.format(                **{'CLASS': object.classname, 'XMIN': object.xmin, 'YMIN': object.ymin, 'XMAX': object.xmax, 'YMAX': object.ymax}))        f.write(xml_body_2)        if display: print("New xml", xml_file)# def SAL_print_log(print_str):##     print('[ERROR]路径不正确: %s' % Picture_file_path)#     sys._getframe().f_linenoif __name__ == "__main__":    debug = 1    #图片地址    Picture_file_path = r'E:/1_Training_picture/11_car/car_train'    Xmls_file_path    = r'E:/1_Training_picture/11_car/car_xml'    # 确认上述地址是否存在    if not os.path.exists(Picture_file_path):        print('[ERROR]路径不正确: %s' % Picture_file_path)        sys.exit(1)    total_picture = os.listdir(Picture_file_path)  # 用于返回指定的文件夹包含的文件或文件夹的名字的列表。    num = len(total_picture)  # xml文件个数    print('共有 %d 张图片进行自动化标注......' % num)    for picture in total_picture:        # 确认上述地址是否存在        picture_name=Picture_file_path + '/'+ picture        print('正在读取: %s' % picture_name)        src = cv.imread(picture_name)        dst=src.copy()        height, width, bytesPerComponent = src.shape  # 取彩色图片的长、宽、通道        print( height, width, bytesPerComponent)        cv.namedWindow("3", 0)        cv.imshow("3", src)        cv.waitKey(10)        print('[DEBUG]运行到:%s 行'  % sys._getframe().f_lineno)        # RGBtoBGR满足opencv显示格式        frame = cv.cvtColor(src, cv.COLOR_BGR2RGB)        # 转变成Image        frame = Image.fromarray(np.uint8(frame))        img0, out_boxes, out_classes = yolo.detect_image(frame)        img0= np.array(img0)        # RGBtoBGR满足opencv显示格式        result = cv.cvtColor(img0, cv.COLOR_RGB2BGR)        print('[DEBUG]运行到:%s 行'  % sys._getframe().f_lineno)        objects = []        for i, c in list(enumerate(out_classes)):            print('[DEBUG]运行到:%s 行' %sys._getframe().f_lineno)            predicted_class=out_classes[i]            box = out_boxes[i]            top, left, bottom, right = box            top = top - 5            left = left - 5            bottom = bottom + 5            right = right + 5            #436 3762 948 4032            top = max(0, np.floor(top + 0.5).astype('int32'))            left = max(0, np.floor(left + 0.5).astype('int32'))            bottom = min(height, np.floor(bottom + 0.5).astype('int32'))            right = min(width, np.floor(right + 0.5).astype('int32'))            xmin = left            ymin = top            xmax=right            ymax=bottom            finall_boundingBoxes=[xmin,ymin,xmax,ymax]            objects.append(OBJECT(finall_boundingBoxes, predicted_class))            if debug:                print("finall_boundingBoxes_1=", finall_boundingBoxes)                cv.rectangle(dst, (int(xmin), (int)(ymin)), ((int)(xmax),(int)(ymax)), (0, 255, 0), 5)        if debug:            cv.namedWindow("result", 0)            cv.imshow('result', result)            cv.waitKey(1)            cv.namedWindow("dst", 0)            cv.imshow('dst', dst)            cv.waitKey(1)        if len(out_classes)>0:            print('[DEBUG]运行到:%s 行'  % sys._getframe().f_lineno)            print("picture=", picture)            xml=os.path.splitext(picture)            xmls_name=Xmls_file_path+"/" +  str(xml[0]) +".xml"            print("xmls_name=", xmls_name)            print('[DEBUG]运行到:%s 行'  % sys._getframe().f_lineno)            imgH=height            imgW=width            create_voc_xml(xmls_name, picture_name, imgW, imgH, objects, display=False)            print('[DEBUG]运行到:%s 行'  % sys._getframe().f_lineno)            #cv.waitKey(0)
2、yolo.py
import tensorflow as tfconfig = tf.compat.v1.ConfigProto(gpu_options=tf.compat.v1.GPUOptions(allow_growth=True))sess = tf.compat.v1.Session(config=config)import osimport numpy as npimport copyimport colorsysfrom timeit import default_timer as timerfrom keras import backend as Kfrom keras.models import load_modelfrom keras.layers import Inputfrom PIL import Image, ImageFont, ImageDrawfrom nets.yolo4 import yolo_body,yolo_evalfrom utils.utils import letterbox_image#--------------------------------------------##   使用自己训练好的模型预测需要修改2个参数#   model_path和classes_path都需要修改!#--------------------------------------------#class YOLO(object):    _defaults = {        "model_path": 'model_data/yolo_car_person_20201126.h5',#加载模型        "anchors_path": 'model_data/yolo_anchors.txt',        "classes_path": 'model_data/coco_classes.txt',        "Chinese_classes_path": 'model_data/my_Chinese_classes.txt',        "score" : 0.5,        "iou" : 0.3,        # 显存比较小可以使用416x416        # 显存比较大可以使用608x608        "model_image_size" : (416, 416)    }    target_type_list = []  ## 空列表,检测出的目标物体    target_type_Chinese_list = []  ## 空列表,检测出的目标物体    boxes_num=0#输出检测目标数量    @classmethod    def get_defaults(cls, n):        if n in cls._defaults:            return cls._defaults[n]        else:            return "Unrecognized attribute name '" + n + "'"    #---------------------------------------------------#    #   初始化yolo    #---------------------------------------------------#    def __init__(self, **kwargs):        self.__dict__.update(self._defaults)        self.class_names = self._get_class()        self.Chinse_class_names = self.get_class_Chinese_name()        self.anchors = self._get_anchors()        self.sess = K.get_session()        self.boxes, self.scores, self.classes = self.generate()    #---------------------------------------------------#    #   获得所有的分类    #---------------------------------------------------#    def _get_class(self):        classes_path = os.path.expanduser(self.classes_path)        with open(classes_path) as f:            class_names = f.readlines()        class_names = [c.strip() for c in class_names]        return class_names    # 读取对应中文标签    def get_class_Chinese_name(self):        classes_path = os.path.expanduser(self.Chinese_classes_path)        with open(classes_path, 'r', encoding='UTF-8') as f:            Chinese_class_names = f.readlines()        Chinese_class_names = [c.strip() for c in Chinese_class_names]        #print(Chinese_class_names)        return Chinese_class_names    #---------------------------------------------------#    #   获得所有的先验框    #---------------------------------------------------#    def _get_anchors(self):        anchors_path = os.path.expanduser(self.anchors_path)        with open(anchors_path) as f:            anchors = f.readline()        anchors = [float(x) for x in anchors.split(',')]        return np.array(anchors).reshape(-1, 2)    #---------------------------------------------------#    #   获得所有的分类    #---------------------------------------------------#    def generate(self):        model_path = os.path.expanduser(self.model_path)        assert model_path.endswith('.h5'), 'Keras model or weights must be a .h5 file.'                # 计算anchor数量        num_anchors = len(self.anchors)        num_classes = len(self.class_names)        # 载入模型,如果原来的模型里已经包括了模型结构则直接载入。        # 否则先构建模型再载入        try:            self.yolo_model = load_model(model_path, compile=False)        except:            self.yolo_model = yolo_body(Input(shape=(None,None,3)), num_anchors//3, num_classes)            self.yolo_model.load_weights(self.model_path)        else:            assert self.yolo_model.layers[-1].output_shape[-1] == \                num_anchors/len(self.yolo_model.output) * (num_classes + 5), \                'Mismatch between model and given anchor and class sizes'        print('{} model, anchors, and classes loaded.'.format(model_path))        # 画框设置不同的颜色        hsv_tuples = [(x / len(self.class_names), 1., 1.)                      for x in range(len(self.class_names))]        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))        self.colors = list(            map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),                self.colors))        # 打乱颜色        np.random.seed(10101)        np.random.shuffle(self.colors)        np.random.seed(None)        self.input_image_shape = K.placeholder(shape=(2, ))        boxes, scores, classes = yolo_eval(self.yolo_model.output, self.anchors,                num_classes, self.input_image_shape,                score_threshold=self.score, iou_threshold=self.iou)        return boxes, scores, classes    #---------------------------------------------------#    #   检测图片    #---------------------------------------------------#    def detect_image(self, image):        start = timer()        self.boxes_num=0        # 调整图片使其符合输入要求        new_image_size = self.model_image_size        boxed_image = letterbox_image(image, new_image_size)        image_data = np.array(boxed_image, dtype='float32')        image_data /= 255.        image_data = np.expand_dims(image_data, 0)  # Add batch dimension.        # 预测结果        out_boxes, out_scores, out_classes = self.sess.run(            [self.boxes, self.scores, self.classes],            feed_dict={                self.yolo_model.input: image_data,                self.input_image_shape: [image.size[1], image.size[0]],                K.learning_phase(): 0            })        print('Found {} boxes for {}'.format(len(out_boxes), 'img'))        #自己加的2020.3.4        self.boxes_num=len(out_boxes)#输出目标框数量        print(self.boxes_num)        #自己加的2020.3.4结束	    # 设置字体        font = ImageFont.truetype(font='font/simhei.ttf',                    size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))        thickness = (image.size[0] + image.size[1]) // 300        small_pic=[]        return_boxes=[]        return_classes=[]        for i, c in list(enumerate(out_classes)):            predicted_class = self.class_names[c]            #predicted_Chinse_class = self.Chinse_class_names[c]            box = out_boxes[i]            score = out_scores[i]            # 自己加的start            if predicted_class != 'car':#过滤非车辆类别               #print(predicted_class)               continue            return_boxes.append(box)            return_classes.append(predicted_class)            #self.target_type_list.append(predicted_class)  # 将标签赋值给列表            #self.target_type_Chinese_list.append(predicted_Chinse_class)  # 将标签赋值给列表            #print(self.target_type_list)            #print(self.target_type_Chinese_list)            #自己加的end            #print("box=", box)            top, left, bottom, right = box            top = top - 5            left = left - 5            bottom = bottom + 5            right = right + 5            top = max(0, np.floor(top + 0.5).astype('int32'))            left = max(0, np.floor(left + 0.5).astype('int32'))            bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32'))            right = min(image.size[0], np.floor(right + 0.5).astype('int32'))            print("top=", top,left,bottom,right,image.size[0],image.size[1])            # 画框框            label = '{} {:.2f}'.format(predicted_class, score)#显示英文            #label = '{} {:.2f}'.format(predicted_Chinse_class, score)#显示中文            draw = ImageDraw.Draw(image)            label_size = draw.textsize(label, font)            label = label.encode('utf-8')            print(label)                        if top - label_size[1] >= 0:                text_origin = np.array([left, top - label_size[1]])            else:                text_origin = np.array([left, top + 1])            for i in range(thickness):                draw.rectangle(                    [left + i, top + i, right - i, bottom - i],                    outline=self.colors[c])            draw.rectangle(                [tuple(text_origin), tuple(text_origin + label_size)],                fill=self.colors[c])            draw.text(text_origin, str(label,'UTF-8'), fill=(0, 0, 0), font=font)            del draw        end = timer()        print(end - start)        return image,return_boxes,return_classes        #return image    def close_session(self):        self.sess.close()

我使用的yolo v4的github公开源码:

该代码基于tensorflow+keras,将源码下载下来后,替我写的两个文件即可使用。

:效果

然后再手动检查一下,就可以快速得到数据集了。

标签: #yolov4源码解读