python根据yolov5输出的label txt文件,批量截取目标框图片并保存

1.根据label标签和图片image,从图片中截取相应的位置,并根据类别进行存储。
txt文件中每一行包含分类类别,中心点xy坐标,以及wh(全部都是归一化的信息)
python根据yolov5输出的label txt文件,批量截取目标框图片并保存_第1张图片
2.根据label信息,截取相应的位置,能够判断图片和label信息相对应,并且批量截图.
代码如下:


from cProfile import label
import os
from pathlib import Path
from tkinter import Image
import cv2
from PIL import Image
import numpy as np

# gain and pad 扩大框的尺寸
def split_img(img, local_infor, n, name, width=640, height=640, gain=1, pad=0):
    cls = str(local_infor[0])
    x = float(local_infor[1]) * width
    y = float(local_infor[2]) * height
    w = int((float(local_infor[3]) * width) * gain + pad)
    h = int((float(local_infor[4]) * height) * gain + pad)

    x1 = int(x - w / 2.0)
    y1 = int(y - h / 2.0)
    x2 = int(x + w / 2.0)
    y2 = int(y + h / 2.0)

    split_img = img[y1: y2, x1: x2]
    save_split_path = save_split_file + '/' + cls
    if not os.path.exists(save_split_path):
        os.makedirs(save_split_path)
    save_path = save_split_path + os.path.sep + '{0}.{1}'.format(name + '_' + str(n), 'jpg')
    cv2.imwrite(save_path, split_img)


if __name__ == '__main__':
    label_file = "D:/testpy/split/label"# 标签存储的位置
    img_file = "D:/testpy/split/img"# 图像存储的位置
    save_split_file = "D:/testpy/split/save"# 截图存储的位置
    img_lists = os.listdir(img_file)
    label_lists = os.listdir(label_file)

    for img_list in img_lists:
        img_path = img_file + '/' + img_list
        if os.path.exists(img_path):
            img = cv2.imread(img_path)
            img_h, img_w, img_c = img.shape
            label_stem = Path(img_list)
            label_name = label_file + '/' + label_stem.stem + ".txt"
            print("Do Image:", img_list)
            if os.path.exists(label_name):
                with open(label_name, 'r', encoding='utf-8') as f:
                    label_infors = f.readlines()
                    f.close()
            		n = 1
            		for infor in label_infors:
                		local_infors = infor.split(" ")
                		split_img(img, local_infors, n, name, img_h, img_w)
                		n = n + 1
			else:
				print("Image:" + Path(img_path).name + " no label")

输出结果如下:
python根据yolov5输出的label txt文件,批量截取目标框图片并保存_第2张图片
python根据yolov5输出的label txt文件,批量截取目标框图片并保存_第3张图片

python根据yolov5输出的label txt文件,批量截取目标框图片并保存_第4张图片

你可能感兴趣的:(深度学习,python,开发语言,深度学习)