YOLOV5实现目标检测后先验框裁剪

# -*- coding: utf-8 -*-
# @Author  : 大叔azhe
# @Time    : 2022/9/12 20:00
# @Function:图片裁剪


def caijian():
    # 仅支持JPEG,PNG,JPG格式图片
    path = "/root/cxc/python/yolov5-3.1/inference/output"  # jpg图片和对应的生成结果的txt标注文件，放在一起
    path3 = "/root/cxc/python/yolov5-3.1/inference/cut/roi"  # 裁剪出来的小图保存的根目录
    path6 = "/root/cxc/python/yolov5-3.1/inference/cut"  # 裁剪出来的小图保存的根目录
    w = 400  # 原始图片resize
    h = 400
    img_total = []
    txt_total = []

    file = os.listdir(path)
    for filename in file:
        first, last = os.path.splitext(filename)
        # if (last in [".jpg",".jpeg","png"]  ):  # 图片的后缀名
        #     img_total.append(first)
        # # print(img_total)
        # else:
        #     txt_total.append(first)

        if (last in [".txt"]):  # 图片的后缀名
            txt_total.append(first)
            # print(img_total)
        else:
            img_total.append(first)
    if os.path.exists(path3):
        shutil.rmtree(path3)
        os.mkdir(path3)
    else:
        os.mkdir(path3)

    for img_ in img_total:
        if img_ in txt_total:
            filename_img = img_ + ".jpg"  # 图片的后缀名
            # print('filename_img:', filename_img)
            path1 = os.path.join(path, filename_img)
            a = os.path.exists(path1)
            if (a == False):
                filename_img = img_ + ".jpeg"  # 图片的后缀名
                # print('filename_img:', filename_img)
                path1 = os.path.join(path, filename_img)
                a = os.path.exists(path1)
            if (a == False):
                filename_img = img_ + ".png"  # 图片的后缀名
                # print('filename_img:', filename_img)
                path1 = os.path.join(path, filename_img)
                a = os.path.exists(path1)
            print("文件是否存在{}".format(a))
            img = cv2.imread(path1)

            img = cv2.resize(img, (w, h), interpolation=cv2.INTER_CUBIC)  # resize 图像大小，否则提取先验框时因原图差异区域可能会报错
            filename_txt = img_ + ".txt"
            # print('filename_txt:', filename_txt)
            n = 1
            with open(os.path.join(path, filename_txt), "r+", encoding="utf-8", errors="ignore") as f:
                for line in f:
                    aa = line.split(" ")
                    x_center = w * float(aa[1])  # aa[1]左上点的x坐标
                    y_center = h * float(aa[2])  # aa[2]左上点的y坐标
                    width = int(w * float(aa[3]))  # aa[3]图片width
                    height = int(h * float(aa[4]))  # aa[4]图片height
                    lefttopx = int(x_center - width / 2.0)
                    lefttopy = int(y_center - height / 2.0)
                    roi = img[lefttopy + 1:lefttopy + height + 3,
                          lefttopx + 1:lefttopx + width + 1]  # [左上y:右下y,左上x:右下x] (y1:y2,x1:x2)需要调参，否则裁剪出来的小图可能不太好
                    print('roi:', roi)
                    filename_last = img_ + "_" + str(n) + ".jpg"  # 裁剪出来的小图文件名
                    # print(filename_last)
                    path2 = os.path.join(path6, "roi")  # 需要在path3路径下创建一个roi文件夹
                    print('path2:', path2)  # 裁剪小图的保存位置
                    cv2.imwrite(os.path.join(path2, filename_last), roi)
                    n = n + 1
        else:
            continue