yolov5/v7/v8自动检测多个文件夹及截取锚框

目前yolo仅支持检测图片或单个文件夹,但在很多时候需要对成百上千个文件夹中图片进行检测,再根据得到的位置信息txt文件来截取图片,如何一步完成呢,详情见下文。

  1. 在detect.py中将save_txt设为store_false,则表示默认保存检测到的txt锚框位置信息

  1. 在最前面加入import shutil,接着ctrl+f找到 if save_txt or save_img:所在位置,在其后加上

    if save_txt or save_img:
        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir/'labels'}" if save_txt else ''
        #print(f"Results saved to {save_dir}{s}")
    global old_path
    global new_path
    old_path = str(save_dir)+'/labels/'
    new_path = str(save_dir)
    for ipath in os.listdir(old_path):
        fulldir = os.path.join(old_path,ipath)
        print(fulldir)
        if os.path.isfile(fulldir):
            shutil.copy(fulldir,new_path)

这段代码是因为yolo会将测后图片与位置信息的txt分开文件夹放置,这里把他们move到一起,为了接下来从图片中截取出锚框图片。

  1. 在detect.py中拉到最后,找到detect(),在其后加入path = ...

        else:
            detect()
            #shutil.move()
            path = new_path         # jpg图片和对应的生成结果的txt标注文件,放在一起
            path3 = "data/person_reid_zhibo_allpic/"    # 裁剪出来的小图保存的根目录

            ss1=time.localtime().tm_min
            ss2=time.localtime().tm_sec
            #ss = 0
            
            path2 = os.path.join(path3,f"{ss1}{ss2}")
            os.mkdir(path2)

            w = 128                         # 原始图片resize
            h = 256
            img_total = []
            txt_total = []

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

            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)
                    img = cv2.imread(path1)
                    img = cv2.resize(img,(w,h),interpolation = cv2.INTER_CUBIC)        # resize 图像大小,否则roi区域可能会报错
                    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)                        # 如果不resize图片统一大小,可能会得到有的roi为[]导致报错         
                            filename_last = img_+"_"+str(n)+".jpg"    # 裁剪出来的小图文件名
                            # print(filename_last)
                            #path2 = os.path.join(path3,"1")           # 在path3路径下创建一个roi文件夹
                            #os.mkdir(path2)
                            print('path2:', path2)                    # 裁剪小图的保存位置
                            cv2.imwrite(os.path.join(path2,filename_last),roi)
                            n = n+1
                else:
                    continue

这段根据当前的时分建立文件夹名称,将每次detect后截取的图片放置在不同文件夹,这里的w和h可以根据自己的要求设置,注意w和h是截取之前原图的尺寸,不是截取后的尺寸。path3根据自己要放的地方设置,是截取小图要保存的根目录。

  1. 重复运行detect.py ,只需执行以下几行代码

import os

#global dir_count

for dir in os.listdir('./your_pic_dir/'):
    #dir_count += len(dir)
    s1 = os.path.basename(dir)
    s2 = os.path.join('./your_pic_dir/',s1)

    run = f'python detect_label.py --weights runs/train/yolov7-tiny_320_desk/weights/best.pt --img-size 320 --source {s2}'
    os.system(run)

print('done')

注意这里的your_pic_dir是你要检测的所有文件夹的根目录,s2是根目录下遍历每个文件夹,这样就可以重复运行detect.py对根目录下每个文件夹进行检测了。再结合在detect.py中加入的代码,可以实现分别检测后再截取保存到各自文件夹下。

猜你喜欢

转载自blog.csdn.net/yjcccccc/article/details/129305210
今日推荐