yolov7裁剪出识别结果

开头先附上两篇有关yolov7的使用教程
①yolov7保姆级教程：链接
②yolov7数据增强与数据划分：链接

本篇主要内容为：在识别出来目标照片后，对目标进行裁剪

1.准备照片与标签

标签的生成在执行代码中加入–save-txt，完整命令如下：

python detect.py --weights .\runs\train\yolov72\weights\best.pt --source .\datasets\findcontact\testImages\ --save-txt

上面的文件路径根据实际情况修改即可
执行后，会生成一个labels文件，里面是yolo格式的标签，将labels与图片放到同一个文件夹

2.代码

代码部分是转载。
根据实际情况，修改文件路径即可。w、h的值，根据需要可自行调整。

# -*- coding: utf-8 -*-
# @Author  : 大叔azhe
# @Time    : 2022/9/12 20:00
# @Function:图片裁剪
import os
import cv2
import shutil
def caijian():
    # 仅支持JPEG,PNG,JPG格式图片
    path = "D:/jmcode/2/yolov7-main/runs/detect/exp5/labels"  # jpg图片和对应的生成结果的txt标注文件，放在一起
    path3 = "D:/jmcode/2/yolov7-main/runs/detect/exp5/cut/roi"  # 裁剪出来的小图保存的根目录
    path6 = "D:/jmcode/2/yolov7-main/runs/detect/exp5/cut"  # 裁剪出来的小图保存的根目录
    w = 4000  # 原始图片resize
    h = 4000
    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
if __name__ == '__main__':  
    caijian()