[YOLO专题-17]：YOLO V5 - 如何把YOLO训练数据集批量转换成带矩形框的图片

Posted 2022-02-01 文火冰糖的硅基工坊

作者主页(文火冰糖的硅基工坊)：文火冰糖（王文兵）的博客_文火冰糖的硅基工坊_CSDN博客

本文网址：https://blog.csdn.net/HiWangWenBing/article/details/122344955

---

目录

前言：

第1章 主要的流程

第2章 工具的使用

第3章 工具源代码实现

---

前言：

yolo的训练数据集，原始图片与yolo的标签文件是分开存储的，有时候，我们期望把标签的box与原始图片整合在一起，看一看标签后的图片的效果。这就需要专门的转换工具来完成。

本文就是提供了这样的一种python工具，来实现此功能。

第1章 主要的流程

（1）创建输出目录：os.makedirs（）

（2）从输入路径中依次读取图片文件：cv2.imread（）

（3）从输入路径中读取对应的同名的标签文件：open(label_file) as f:

（4）从标签文件中依次读取多个Box标签信息：f.read()

（5）把yolo标签矩形框的坐标，还原成正常的坐标值：

                    x1 = int(width  * (box[1] - box[3]/2))
                    y1 = int(height * (box[2] - box[4]/2))
                    x2 = int(width  * (box[1] + box[3]/2))
                    y2 = int(height * (box[2] + box[4]/2))

（6）在指定的图片上画矩形框：cv2.rectangle（）

（7）存储带有矩形框的生成图片：cv2.imwrite（）

第2章 工具的使用

python convertYoloLabelToBox.py --input train_data\\train\\ --output tmp_box

convertYoloLabelToBox.py ： 工具名称

（1）--input train_data\\train\\： 包含yolo训练数据集的根目录，该目录中必须包括images和lables子目录。

• images：原始的训练图片
• lables：图片对应的yolo标签文件

（2）--output tmp_box：输出图片的目录

第3章 工具源代码实现

import os
import numpy as np
import cv2
import argparse

def convertYoloLabelToBox(input, output):
    # the root path of images
    images_path = os.path.join(input,'images')
    
    # the root path of yolo labels
    labes_path = os.path.join(input,'labels')
    
    if not os.path.exists(output):
        os.makedirs(output)
    
    # 获取图片文件列表
    file_list = os.listdir(images_path)
    
    # 依次进行转换
    for filename in file_list:
        if 'jpg' in filename:
            # 构建image文件名称
            img_file = os.path.join(images_path, filename)
            
            # 读取Image文件
            image = cv2.imread(img_file)
            
            # 获取图片的尺寸
            height = image.shape[0]
            width  = image.shape[1]
            
            # 构建yolo Label文件名称
            label_file = os.path.join(labes_path, filename.split('.')[0]+'.txt')
            
            # 打开yolo label文件
            with open(label_file) as f:
                # 读取label文件内容: 由N行数据组成，每一行代表一个物体的Box，每个Box有1个分类参数 + 4个坐标参数
                box_list = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32)
                
                for box in box_list:
                    # 重构box的坐标，恢复成常规坐标值
                    x1 = int(width  * (box[1] - box[3]/2))
                    y1 = int(height * (box[2] - box[4]/2))
                    x2 = int(width  * (box[1] + box[3]/2))
                    y2 = int(height * (box[2] + box[4]/2))
                    
                    # 在图片上画矩形框
                    start_point = (x1, y1)
                    end_point = (x2, y2)
                    cv2.rectangle(img = image, pt1 = start_point, pt2 = end_point, color = (0, 255, 0), thickness = 1)
            # 存储具有矩形框的图片 
            cv2.imwrite(os.path.join(output, filename), image)
    print('Finished, output path =', output)
    return

def parse_opt():
    # define argparse object
    parser = argparse.ArgumentParser()
    
    # add argument for command line
    parser.add_argument('--input',      type=str, help='The input  directory with YOLO Labels and images')
    parser.add_argument('--output',     type=str, help='The output directory to save images with boxes')
    
    # parse arges from command line
    opt = parser.parse_args()
    print("input  =", opt.input)
    print("output =", opt.output)
    
    # return opt
    return opt

def main(opt):
    convertYoloLabelToBox(**vars(opt))
if __name__ == '__main__':
    opt = parse_opt()
    main(opt)

---

作者主页(文火冰糖的硅基工坊)：文火冰糖（王文兵）的博客_文火冰糖的硅基工坊_CSDN博客

本文网址：https://blog.csdn.net/HiWangWenBing/article/details/122344955