python实现:目标检测裁剪图片样本,根据类标签文件进行裁剪保存

我在进行目标检测时候,比如红绿灯检测,目标区域很小,样本杂乱。

 

想要筛选错误样本的话,很困难。可以把目标区域裁剪出来。人大脑处理对于这样的异己比较敏感。样本量较少的话可以自己筛一筛。样本量较大的话,可以训练一个分类模型帮你筛一下。

它就可以实现一个目标检测的数据集可以转化为一个分类的数据集。

这个是根据txt标签文件进行筛选的,大同小异,别的文件类型也就是登录读取信息不一样。

它最终会读取你目标检测的每一个类,然后每个类会创建一个就叫这个类名的文件夹,然后裁剪出来的这一类的图片。最后图片的文件名就是你的”原文件名+num“。

#根据预测出来的txt文件裁剪图片
import os
import cv2
from tqdm import tqdm

image_input = '/home/xys/CloundShiProjects/traffic_light/trafficlight_dect/data/JPEGImages/'
txt_input = '/home/xys/CloundShiProjects/traffic_light/trafficlight_dect/data/labels/'
path_output = "/home/xys/CloundShiProjects/traffic_light/trafficlight_dect/crop/"    # 裁剪出来的小图保存的根目录
class_names_path = '/home/xys/CloundShiProjects/traffic_light/trafficlight_dect/data/classes.txt'

img_total = []
txt_total = []

def read_class_name(path):        #读取path下的类别民
    f = open(path,'r')
    classes_name = []
    for i in f.readlines():
        classes_name.append(i.strip())
    return classes_name
classes_name = read_class_name(class_names_path)

file_image = os.listdir(image_input)
for filename in file_image:#在做jpg文件名列表
     first,last = os.path.splitext(filename)
     img_total.append(first)

file_txt = os.listdir(txt_input)
for filename in file_txt:#在做txt文件名列表
     first,last = os.path.splitext(filename)
     txt_total.append(first)

for img_ in tqdm(img_total):
    if img_ in txt_total:
        filename_img = img_+".jpg"
        path1 = os.path.join(image_input,filename_img)
        img = cv2.imread(path1)
        filename_txt = img_+'.txt'     #预测出来的txt文件没有后缀名,有则加 {+".txt"}
        h = img.shape[0]
        w = img.shape[1]
        n = 1
        with open(os.path.join(txt_input,filename_txt),"r+",encoding="utf-8",errors="ignore") as f:
            for line in f:
                aa = line.split(" ")
                # if not int(aa[0]) == 0: continue     #判断需要裁剪的类别:0--vehicle
                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)需要调参,否则裁剪出来的小图可能不太好
                if roi.size == 0: continue
                filename_last = img_+"_"+str(n)+".jpg"      # 裁剪出来的小图文件名
                x = int(aa[0])
                path2 = os.path.join(path_output,classes_name[x])           # 需要在path_output路径下创建一个cut_txt文件夹
                if not os.path.exists(path2):
                    os.mkdir(path2)
                # print('path2:', path2)                    # 裁剪小图的保存位置
                cv2.imwrite(os.path.join(path2,filename_last),roi)
                n = n+1
    else:
        continue

希望可以帮到你们!

你可能感兴趣的:(yolov4,python,目标检测,开发语言,计算机视觉,opencv)