深度学习目标检测常用工具型代码：检测结果按 类别 绘制于原图上（绘制图片保存到另一文件夹下）

# -*- coding: utf-8 -*-
"""
    作用是将检测出的类结果（如VOC是20类，也就是20个txt）画在原图上，并另存到一个路径下
    srcpath表示txt的路径 
    dstpath表示原始图片的路径
    
    txt格式如（文末有截图）：  
    Task2_baseball-diamond.txt
    P0034 0.840593 1593.0 1443.0 1778.0 1600.0  
    
    
"""
import os
import cv2

#像这种常用的函数可以作为参考
def custombasename(fullname):
    return os.path.basename(os.path.splitext(fullname)[0])

def GetFileFromThisRootDir(dir,ext = None):
  allfiles = []
  needExtFilter = (ext != None)
  for root,dirs,files in os.walk(dir):
    for filespath in files:
      filepath = os.path.join(root, filespath)
      extension = os.path.splitext(filepath)[1][1:]
      if needExtFilter and extension in ext:
        allfiles.append(filepath)
      elif not needExtFilter:
        allfiles.append(filepath)
  return allfiles
  
def draw_bbox_by_det_result(srcpath, dstpath):
    filelist = GetFileFromThisRootDir(srcpath)
    for fullname in filelist:
        classname = custombasename(fullname)
        classname = classname[6:]#这个地方需要自己改，懒得再改了。
        with open(fullname, 'r') as f_in:
            nameboxdict = {}
            lines = f_in.readlines()
            splitlines = [x.strip().split(' ') for x in lines]
            for splitline in splitlines:
                 oriname = splitline[0]
                 confidence = splitline[1]
                 poly = list(map(float, splitline[2:]))
                 det = poly
                 det.append(confidence)
                 det = list(map(float, det))
                 if (oriname not in nameboxdict):
                     nameboxdict[oriname] = []
                 nameboxdict[oriname].append(det)
                #字典的for in是根据key做的
            for imgname in nameboxdict:
                im_file = os.path.join(dstpath+'/'+imgname+'.png')
                im = cv2.imread(im_file)
                for det in nameboxdict[imgname]:
                        #print('det:', det)
                    confidence = det[-1]
                    bbox = det[0:-1]
                    cv2.rectangle(im,(int(bbox[0]),int(bbox[1])),(int(bbox[2]),int(bbox[3])),(255,255,0),2)
                    cv2.putText(im, '{:s} {:.3f}'.format(classname,confidence), (int(bbox[0]), int(bbox[1] - 5)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,255), 2)
                #这里的路径表示保存图片的路径 下面有效果展示
                cv2.imwrite(os.path.join('/media/ygx/Elements/VOCdevkit2007/testsplit1.5/1','%s_%s.txt.png' % (imgname,classname)),im)

if __name__ == '__main__':
    draw_bbox_by_det_result(r'/media/ygx/参赛/最后提交的zip文件/small-vehicle最终所用/nms/nms0.3', r'/media/ygx/Elements/DOTA比赛原始下载数据/test')