-
下载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