YOLOv8 target detection and image cropping

Preface

We have previously trained model files for identifying whether people are wearing masks, which can effectively identify whether people are wearing masks. This article will explain how to crop the identified targets.

Target Recognition

Need to specify save_txt=True, save the data annotation file txt

yolo predict model=runs/detect/train26/weights/best.pt source=ultralytics/assets/mask save_txt=True

Check the annotation file directory and file format

Coordinate transformation

The format in the data annotation file (all values ​​except categories are normalized)

w: picture pixel width

h: picture pixel height

0 0.301826 0.367765 0.123616 0.229143
label x_center y_center _width _height

x1=(x_center-width/2) * w
x2=(x_center+width/2) * w
y1=(y_center-height/2) * h
y2=(y_center+height/2) * h

Trim script

import os
from PIL import Image
import shutil

def findSingleFile(path):
        # 创建 cutpictures 文件夹（先判断）
    cutp = os.path.join(path, "cutpictures") 
        # 判断文件夹是否存在
    if os.path.exists(cutp):
        # 如果文件夹存在，先删除再创建
        # 递归删除文件夹
        shutil.rmtree(cutp)
        os.makedirs(cutp)
    else:
        # 如果文件夹不存在，直接创建
        os.makedirs(cutp)

    for filename in os.listdir(path):
        if not os.path.isdir(os.path.join(path,filename)):
            # 无后缀文件名
            filename_nosuffix = filename.split(".")[0]
            # 文件后缀
            file_suffix = filename.split(".")[1]
            # print(filename_nosuffix)

            img_path = os.path.join(path,filename)
            label_path = os.path.join(path,'labels',filename_nosuffix+".txt")

            # print(img_path)
            # print(label_path)
            # 生成裁剪图片（遍历 txt 每一行）eg: mask_0_1.jpg
            # 0 裁剪的图片序号 1 类别序号
            img = Image.open(img_path)
            w, h = img.size
            with open(label_path, 'r+', encoding='utf-8') as f:
                # 读取txt文件中的第一行，数据类型str
                lines = f.readlines()
                # 根据空格切割字符串，最后得到的是一个list
                for index, line in enumerate(lines):
                    msg = line.split(" ")
                    category = int(msg[0])
                    x_center = float(msg[1])
                    y_center = float(msg[2])
                    width = float(msg[3])
                    height = float(msg[4])
                    x1 = int((x_center - width / 2) * w)  # x_center - width/2
                    y1 = int((y_center - height / 2) * h)  # y_center - height/2
                    x2 = int((x_center + width / 2) * w)  # x_center + width/2
                    y2 = int((y_center + height / 2) * h)  # y_center + height/2
                    # print(x1, ",", y1, ",", x2, ",", y2, "," ,category)
                    # 保存图片
                    img_roi = img.crop((x1, y1, x2, y2))
                    save_path = os.path.join(cutp, "{}_{}_{}.{}".format(filename_nosuffix, index, category, file_suffix))
                    img_roi.save(save_path)

    print("裁剪图片存放目录：", cutp)


def main():
    import argparse

    # 创建 ArgumentParser 对象
    parser = argparse.ArgumentParser(description='输入目标检测裁剪目录')

    # 添加参数
    parser.add_argument('--dir', help='目录名', required=True)

    # 解析命令行参数
    args = parser.parse_args()

    dir = args.dir
    # print('目录参数:', dir)

    findSingleFile(dir)
    return

if __name__ == '__main__':
    main()

Execute the script (need to specify the --dir directory location parameter)

python cutpictures.py --dir /home/hualiujie/baoxinshagnchuan/ultralytics-main-cgh/runs/detect/predict6

Generated successfully, file name meaning:{无后缀文件名}_{裁剪图片序号}_{标签类别序号}

Reference article

Use yolov5 for target detection and crop the detected targets_Target Detection_Cerebellar Ax AI Eating Meat-Huawei Cloud Developer Alliance (csdn.net)

This article is published by OpenWrite, a blog that publishes multiple articles !