目标检测任务中的Bounding box绘制

发表于 |

背景

计算机视觉目标检测任务中，输出的result是一系列图像及对应的txt文档，文档里面记录了检测到类别的bounding box信息：其中每行是一个类别，第1列代表哪一类，第2列是该类别置信度，后面4列是该bounding box的左上、右下的坐标信息(以图的左上方为坐标0点)，如下所示：

7 0.9482899 157.50272 192.24951 216.85303 248.50638
12 0.7890044 247.94312 211.89938 304.7307 280.8666
18 0.5444266 101.857285 215.64401 181.53456 261.96036

解决方案

为了直观显示每张图像包含了哪些类别，我们用python对bounding box进行了重绘，并将图像按类别放置于对应的文件夹中，文件夹名称为1～21，以字典形式对应于21个类别，代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
# -*- coding:utf-8 -*-
#!/usr/bin/env Python
import os
import shutil
from PIL import Image, ImageDraw

#==========================建立索引文件夹==========================

Path1_A = "/home/tucoder/Desktop/demo1_result/S1/" # 图像按类型分配到各个文件夹
Path2_A = "/home/tucoder/Desktop/demo1_result/S2/" # 存储类型文本

dict1 = {'1':'飞机','2':'自行车','3':'鸟','4':'船','5':'瓶子','6':'公共汽车','7':'车','8':'猫','9':'椅子','10':'牛','11':'餐桌','12':'狗', \
    '13':'马','14':'摩托车','15':'人','16':'植物','17':'羊','18':'沙发','19':'火车','20':'电视'} # 21类
    
os.mkdir(Path1_A)
  
os.mkdir(Path2_A)

for i in range(1, 21):
        os.mkdir(Path1_A+str(i))

def draw(Path1, x1, y1, x2, y2, path2):
    '''
    绘制矩形框函数
    原pic位置path1
    处理后的pic被存储到path2中
    '''
    im = Image.open(Path1)
    draw = ImageDraw.Draw(im)
    line = 5
    x, y = x1, y1
    x_1, y_1 = x2, y2

    for i in range(1, line + 1):
        draw.rectangle((x + (line - i), y + (line - i), x_1 + i, y_1 + i), outline='red')
    im.save(path2)

#==========================main==========================
if __name__ == "__main__":
    result = []
    alllist=os.listdir("/home/tucoder/Desktop/demo1_result")
    for i in alllist:
        if "." in i:
            aa, bb=i.split(".")
            if bb == "txt":
                result.append(aa+"."+bb)
    A = set()
    for filename in result:
        set1 = set()
        string = '' # 每幅图片包含其class属性
        str1 = filename[:filename.rindex(".")]
        with open(filename, 'r') as f:
            for line in f.readlines():
                list1 = line.split()
                if float(list1[1]) >= 0.80: # 类别置信度>=0.8，判定其为真
                    A.add(list1[0])
                    Path_1 = str1+".jpg"
                    Path_2 = Path1_A+'/'+list1[0]+"/"+str1+".jpg"
                    if list1[0] not in set1:
                        set1.add(list1[0]) 
                        draw(Path_1, float(list1[2]), float(list1[3]), float(list1[4]), float(list1[5]), Path_2)
                        string = string +list1[0]+"-"
                    else:
                        draw(Path_2, float(list1[2]), float(list1[3]), float(list1[4]), float(list1[5]), Path_2)
        set1.clear()
        List = []
    for str_i in list(A):
        List.append(dict1[str_i])
    print(List)
    with open(Path2_A+'text', 'w') as f:
        f.write(str(List))

最终效果如图：

图片名称