此笔记是为记录一下，为解决某模型只能预测一张图，并且不能生成相应的txt文件的问题
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} 不是一个目录。")
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
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:continuex0 = 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