目前yolo仅支持检测图片或单个文件夹,但在很多时候需要对成百上千个文件夹中图片进行检测,再根据得到的位置信息txt文件来截取图片,如何一步完成呢,详情见下文。
在detect.py中将save_txt设为store_false,则表示默认保存检测到的txt锚框位置信息
在最前面加入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到一起,为了接下来从图片中截取出锚框图片。
在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根据自己要放的地方设置,是截取小图要保存的根目录。
重复运行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中加入的代码,可以实现分别检测后再截取保存到各自文件夹下。