-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathget_dr_txt.py
110 lines (90 loc) · 4.72 KB
/
get_dr_txt.py
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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
#----------------------------------------------------#
# 获取测试集的detection-result和images-optional
# 具体视频教程可查看
# https://www.bilibili.com/video/BV1zE411u7Vw
#----------------------------------------------------#
import copy
import os
import numpy as np
import torch
from PIL import Image
from tqdm import tqdm
from frcnn import FRCNN
from utils.utils import get_new_img_size
'''
这里设置的门限值较低是因为计算map需要用到不同门限条件下的Recall和Precision值。
所以只有保留的框足够多，计算的map才会更精确，详情可以了解map的原理。
计算map时输出的Recall和Precision值指的是门限为0.5时的Recall和Precision值。
此处获得的./input/detection-results/里面的txt的框的数量会比直接predict多一些，这是因为这里的门限低，
目的是为了计算不同门限条件下的Recall和Precision值，从而实现map的计算。
这里的self.iou指的是非极大抑制所用到的iou，具体的可以了解非极大抑制的原理，
如果低分框与高分框的iou大于这里设定的self.iou，那么该低分框将会被剔除。
可能有些同学知道有0.5和0.5:0.95的mAP，这里的self.iou=0.5不代表mAP0.5。
如果想要设定mAP0.x，比如设定mAP0.75，可以去get_map.py设定MINOVERLAP。
'''
class mAP_FRCNN(FRCNN):
#---------------------------------------------------#
# 检测图片
#---------------------------------------------------#
def detect_image(self,image_id,image):
self.confidence = 0.01
self.iou = 0.45
f = open("./input/detection-results/"+image_id+".txt","w")
#-------------------------------------#
# 转换成RGB图片，可以用于灰度图预测。
#-------------------------------------#
image = image.convert("RGB")
image_shape = np.array(np.shape(image)[0:2])
old_width, old_height = image_shape[1], image_shape[0]
old_image = copy.deepcopy(image)
#---------------------------------------------------------#
# 给原图像进行resize，resize到短边为600的大小上
#---------------------------------------------------------#
width,height = get_new_img_size(old_width, old_height)
image = image.resize([width,height], Image.BICUBIC)
#-----------------------------------------------------------#
# 图片预处理，归一化。
#-----------------------------------------------------------#
photo = np.transpose(np.array(image,dtype = np.float32)/255, (2, 0, 1))
with torch.no_grad():
images = torch.from_numpy(np.asarray([photo]))
if self.cuda:
images = images.cuda()
roi_cls_locs, roi_scores, rois, _ = self.model(images)
#-------------------------------------------------------------#
# 利用classifier的预测结果对建议框进行解码，获得预测框
#-------------------------------------------------------------#
outputs = self.decodebox.forward(roi_cls_locs[0], roi_scores[0], rois, height = height, width = width, nms_iou = self.iou, score_thresh = self.confidence)
#---------------------------------------------------------#
# 如果没有检测出物体，返回原图
#---------------------------------------------------------#
if len(outputs)==0:
return old_image
outputs = np.array(outputs)
bbox = outputs[:,:4]
label = outputs[:, 4]
conf = outputs[:, 5]
bbox[:, 0::2] = (bbox[:, 0::2]) / width * old_width
bbox[:, 1::2] = (bbox[:, 1::2]) / height * old_height
for i, c in enumerate(label):
predicted_class = self.class_names[int(c)]
score = str(conf[i])
left, top, right, bottom = bbox[i]
f.write("%s %s %s %s %s %s\n" % (predicted_class, score[:6], str(int(left)), str(int(top)), str(int(right)),str(int(bottom))))
f.close()
return
frcnn = mAP_FRCNN()
image_ids = open('VOCdevkit/VOC2007/ImageSets/Main/test.txt').read().strip().split()
if not os.path.exists("./input"):
os.makedirs("./input")
if not os.path.exists("./input/detection-results"):
os.makedirs("./input/detection-results")
if not os.path.exists("./input/images-optional"):
os.makedirs("./input/images-optional")
for image_id in tqdm(image_ids):
image_path = "./VOCdevkit/VOC2007/JPEGImages/"+image_id+".jpg"
image = Image.open(image_path)
image = image.convert("RGB")
# image.save("./input/images-optional/"+image_id+".jpg")
frcnn.detect_image(image_id,image)
print("Conversion completed!")