【深度学习】修改源码以实现预测整个文件夹和保存txt文件（自用笔记

努力敲代码的小盆友2024-09-16 18:19

此笔记是为记录一下，为解决某模型只能预测一张图，并且不能生成相应的txt文件的问题

python 复制代码

def main():
    if os.path.isdir(directory_path):
        files = os.listdir(directory_path)
        for file in files:
            file_path = os.path.join(directory_path, file)
            if os.path.isfile(file_path):
                try:
                    origin_img = np.asarray(Image.open(file_path).convert('RGB'))
                    bboxes = [[1, 2, 3, 4], [5, 6, 7, 8]]
                    scores = [100, 99]
                    cls_inds = [0, 1]
                    vis_res = visualize(origin_img, bboxes, scores, cls_inds, conf=0.6, save_name=os.path.basename(file_path), save_result=True)
                    print(os.path.basename(file_path))
                except Exception as e:
                    print(f"读取文件 {file} 时发生错误：{e}")
    else:
        print(f"{directory_path} 不是一个目录。")

python 复制代码

def visualize(image, bboxes, scores, cls_inds, conf, save_name='vis.jpg', save_result=True):
    # 写入txt文件 一行一行的写入
    vis_img, labels = vis(image, bboxes, scores, cls_inds, conf, class_names)
    img_name = save_name.split('.')[0]
    txt_name = img_name + '.txt'
    if save_result:  # 把已经框好的照片写进去
        save_path = os.path.join(output_dir, save_name)
        print(f"save visualization results at {save_path}")
        save_txt_path = os.path.join(output_dir, txt_name)
        for i in range(len(labels)):
            label = labels[i]
            with open(save_txt_path, 'a') as fp:
                fp.write(label)

        # cv2.imwrite(save_path, vis_img[:, :, ::-1])
    return vis_img

python 复制代码

def vis(img, boxes, scores, cls_ids, conf=0.5, class_names=None):
    # 返回labels数组
    labels = []
    for i in range(len(boxes)):
        box = boxes[i]
        cls_id = int(cls_ids[i])
        score = scores[i]
        if score < conf:
            continue
        x0 = int(box[0])
        y0 = int(box[1])
        x1 = int(box[2])
        y1 = int(box[3])
        x2 = float(box[0])
        y2 = float(box[1])
        x3 = float(box[2])
        y3 = float(box[3])
        score_txt = float(score)

        label = str(x2)+' '+str(y2)+' '+str(x3)+' '+str(y3)+' '+str(score)+' '+str(cls_id)+'\n'
        labels.append(label)
    return img, labels