深度学习python之用Faster-rcnn 检测结果（txt文件） 在原图画出box

使用Faster-rcnn 的test_net.py 检测网络的mAP等精度会生成一个检测结果（txt文件），格式如下：

000004 0.972 302.8 94.5 512.0 150.0
000004 0.950 348.1 166.1 512.0 242.9
000004 0.875 1.0 25.7 292.6 126.3
000004 0.730 1.0 138.5 488.3 230.0
000004 0.699 1.0 120.9 145.5 139.9
000004 0.592 54.4 227.4 431.9 343.4
000004 0.588 1.0 159.8 18.8 231.6
000004 0.126 1.0 247.1 342.3 270.0
000004 0.120 1.0 225.4 185.7 309.3

每行分别为 名称 检测概率 xmin ymin xmax ymax

问题在于每一行只显示一个box数据，每幅图像可能包括多个box，需要判断提取的多行数据是不是属于同一图片

下面使用python提取这些数据，在原图上画出box并且保存起来

import os
import os.path
import numpy as np
import xml.etree.ElementTree as xmlET
from PIL import Image, ImageDraw
import cPickle as pickle  

txt_name = 'comp4_8a226fd7-753d-40fc-8013-f68d2a465579_det_test_ship.txt'
file_path_img = '/home/JPEGImages'
save_file_path = '/home/detect_results'


source_file = open(txt_name)

img_names = []
for line in source_file:
    staff = line.split()
    img_name = staff[0]
    img_names.append(img_name)

name_dict = {}
for i in img_names:
    if img_names.count(i)>0:
        name_dict[i] = img_names.count(i) 

source_file.close()

source_file = open(txt_name)
for idx in name_dict:
    img = Image.open(os.path.join(file_path_img, idx + '.jpg')) 
    draw = ImageDraw.Draw(img)
    for i in xrange(name_dict[idx]):
        line = source_file.readline()
        staff = line.split()
        score = staff[1]
        box = staff[2:6]
        draw.rectangle([int(np.round(float(box[0]))), int(np.round(float(box[1]))), 
                    int(np.round(float(box[2]))), int(np.round(float(box[3])))], outline=(255, 0, 0))
    img.save(os.path.join(save_file_path, idx + '.jpg'))  

source_file.close()

运行完即可在保存文件夹中得到效果图。