YOLO-V5-s 实现人脸框检测

• 下载yolov-v5代码，并安装依赖库（推荐adaconda创建虚拟环境后安装）

  github地址：https://github.com/ultralytics/yolov5

  yolo-v5官方文档：https://docs.ultralytics.com/yolov5/tutorials/train_custom_data/#train-on-custom-data
  

  git clone https://github.com/ultralytics/yolov5  # clone
  cd yolov5
  pip install -r requirements.txt  # install

 

 

• 下载人脸检测训练集
  WIDER FACE训练集：http://shuoyang1213.me/WIDERFACE/index.html
  

  

   

  
  
  
  
• 将WINDER FACE人脸框信息转为YOLO格式
  由于训练速度慢，验证集数据量最少，本次将使用验证集作为训练集使用
  将训练集图片和对应人脸框信息放到YOLO-v5项目下
  

  
  
  
  

•  

  使用脚本将WINDER FACE数据转化为YOLO格式

  import os
  import shutil
  
  
  class WiderFace(object):
      def __init__(self, label_path, img_dir_path, save_path, class_num):
          self.class_num = class_num
          self.label_path = label_path
          self.img_dir_path = img_dir_path
          self.save_path = save_path
          # self.train_percent = train_percent
          # self.val_percent = val_percent
          # self.test_percent = test_percent
          self.yolo_img_path = os.path.join(save_path, "images")
          self.yolo_label_path = os.path.join(save_path, "labels")
          if not os.path.exists(self.yolo_img_path ):
              os.makedirs(self.yolo_img_path)
          if not os.path.exists(self.yolo_label_path ):
              os.makedirs(self.yolo_label_path)
  
      def analysis_line(self, line):
          """判断当前行数据类型，并返回"""
          if "." in line:
              img_path = os.path.join(self.img_dir_path, line)
              img_path = img_path.replace("\n", "")
              return "img", img_path
          if " " in line:
              # 左、上、宽、高
              landmark = line.split(" ")[:4]
              return "bbox", landmark
          return "box_sum", line.replace("\n", "")
  
      def save_img(self, img_path):
          """复制图片到新目录并按照图片数量重新命名"""
          file_count = len(os.listdir(self.yolo_img_path))
          new_file_name = "%05d" % file_count + os.path.splitext(img_path)[-1]
          new_file_path = os.path.join(self.yolo_img_path, new_file_name)
          shutil.copy(img_path, new_file_path)
          return new_file_name
  
      def save_yolo(self):
          with open(self.label_path, "r") as f:
              bbox_list = []
              cur_img_name = ""
              for line in f:
                  line_type, data = self.analysis_line(line)
                  if line_type == 'img':
                      # 显示图片和人脸框
                      # if bbox_list:
                      #     from PIL import Image, ImageDraw
                      #     new_file_path = os.path.join(self.yolo_img_path, cur_img_name)
                      #     img = Image.open(new_file_path)
                      #     draw = ImageDraw.Draw(img)
                      #     for item in bbox_list:
                      #         draw.rectangle(xy=item, outline="red")
                      #     img.show()
                      #     img.close()
                      #     bbox_list.clear()
  
                      # 拷贝图片并重新命名
                      img_path = os.path.join(self.yolo_img_path, data)
                      cur_img_name = self.save_img(img_path)
                  elif line_type == 'bbox':
                      # 存储所属分类和人脸框信息
                      from PIL import Image
                      new_file_path = os.path.join(self.yolo_img_path, cur_img_name)
                      img = Image.open(new_file_path)
                      w, h = img.size
                      x_center, y_center = int(data[0])+int(data[2])/2, int(data[1])+int(data[3])/2
                      yolo_label = "1 %s %s %s %s\n" % (x_center/w, y_center/h, int(data[2])/w, int(data[3])/h)
                      label_path = os.path.join(self.yolo_label_path, cur_img_name.split(".")[0]) + ".txt"
                      with open(label_path, "a+") as f:
                          f.writelines(yolo_label)
  
                      # 存储人脸框信息
                      xy = (int(data[0]), int(data[1]), int(data[0])+int(data[2]), int(data[1])+int(data[3]))
                      bbox_list.append(xy)
  
  
  if __name__ == "__main__":
      label_path = "/Users/wb-fcj414969/project/yolov5-master/samples/wider_face/wider_face_val_bbx_gt.txt"
      img_dir_path = "/Users/wb-fcj414969/project/yolov5-master/samples/wider_face/WIDER_val/images"
      save_dir_path = "/Users/wb-fcj414969/project/yolov5-master/samples/wider_face/yolo/"
      # train_percent = 0.6
      # val_percent = 0.2
      # test_percent = 0.2
      wider_face = WiderFace(label_path, img_dir_path, save_dir_path, 1)
      wider_face.save_yolo()
  

  　　

• 新建数据集配置文件face.yaml放入data文件中

  # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
  # COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
  # Example usage: python train.py --data coco128.yaml
  # parent
  # ├── yolov5
  # └── datasets
  #     └── coco128  ← downloads here (7 MB)
  
  
  # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
  path:*****/samples/wider_face/yolo  # dataset root dir
  train: images  # train images (relative to 'path') 128 images
  val: images  # val images (relative to 'path') 128 images
  test:  # test images (optional)
  
  # Classes
  names:
    0: person
    1: face
  
  
  # Download script/URL (optional)
  #download: https://ultralytics.com/assets/coco128.zip
  

  　　

   

• shell命令启动训练脚本
  python train.py --img 640 --epochs 3 --data face.yaml --weights yolov5s.pt