Python 深度学习目标检测结果可视化

深度学习目标检测的标签数据是以一系列点的形式存储在文件中，输出结果也是以点的形式表示，难以进行目视判读。本代码可以根据检测结果在原始影像上绘制边界框，实现检测结果的可视化。

在OBB的目标检测中，DOTA数据的标签以[x1 y1 x2 y2 x3 y3 x4 y4 class difficult]格式记录在txt文件中，如下所示：

 标签文件与影像文件的名称一一对应：

将标签数据绘制在影像上的可视化代码如下：

# -*- coding: utf-8 -*-
import os, cv2
import numpy as np
import tifffile


# 读取标签文件，返回一系列边界框的坐标点及类别信息
# 文件中一行代表一个边界框，每行由"x1 y1 x2 y2 x3 y3 x4 y4 class difficult"组成
def read_label(label_file):
    with open(label_file, 'r') as f:
        box_data = []  # 存储坐标点信息
        label_data = []  # 存储标类别信息
        for line in f.readlines():
            curLine = line.strip().split(" ")
            x1 = float(curLine[0])
            y1 = float(curLine[1])
            x2 = float(curLine[2])
            y2 = float(curLine[3])
            x3 = float(curLine[4])
            y3 = float(curLine[5])
            x4 = float(curLine[6])
            y4 = float(curLine[7])
            box = np.array([[x1, y1], [x2, y2], [x3, y3], [x4, y4]], np.int32)
            label = curLine[8]
            box_data.append(box)
            label_data.append(label)
    return box_data, label_data


label_dir = './labelTxt/'  # 存储标签的文件夹
image_dir = './images/'  # 存储影像的文件夹
out_dir = './showlabel/'  # 保存输出文件的文件夹
label_files = os.listdir(label_dir)

n = 0
for label_file in label_files:  # 遍历标签文件
    if '.txt' in label_file:
        label_path = label_dir + label_file  # 标签文件路径
        image_path = image_dir + label_file.replace('.txt', '.tif')  # 影像文件路径（影像与标签文件名对应，替换后缀即可）
        im = cv2.imread(image_path)
        box_data, label_data = read_label(label_path)  # 读取标签文件中的坐标点信息
        for i in range(len(box_data)):
            x1 = box_data[i][0][0]  # 第一个坐标点，为了固定类别标签的位置
            y1 = box_data[i][0][1]
            bbox = np.array(box_data[i], np.int32)
            bbox = bbox.reshape((-1, 1, 2))
            cv2.polylines(im, [bbox], True, (4, 7, 250), 1)  # 绘制边界框
            font = cv2.FONT_HERSHEY_SIMPLEX
            # cv2.putText(im, label_data[i], (x1, y1 - 7), font, 0.5, (6, 230, 230), 1)    #可以在边界框上标注类别
        label_show_file = out_dir + label_file.replace('.txt', '.tif')
        cv2.imwrite(label_show_file, im)  # 保存可视化后的影像