forked from eriklindernoren/PyTorch-YOLOv3
-
Notifications
You must be signed in to change notification settings - Fork 0
/
detect_single.py
147 lines (118 loc) · 4.28 KB
/
detect_single.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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
from __future__ import division
import argparse
import cv2
from models import *
from utils.datasets import *
from utils.utils import *
parser = argparse.ArgumentParser()
parser.add_argument(
"--model_config_path", type=str, required=True, help="path to model config file"
)
parser.add_argument(
"--data_config_path", type=str, required=True, help="path to data config file"
)
parser.add_argument(
"--weights_path", type=str, required=True, help="path to weights file"
)
parser.add_argument(
"--conf_thres", type=float, default=0.8, help="object confidence threshold"
)
parser.add_argument(
"--nms_thres",
type=float,
default=0.4,
help="iou thresshold for non-maximum suppression",
)
parser.add_argument(
"--area_thres", type=float, default=0, help="ignore objs with area < threshold"
)
parser.add_argument(
"--use_cuda", action="store_true", help="whether to use cuda if available"
)
opt = parser.parse_args()
for x in opt.__dict__:
print("%25s: %s" % (x, opt.__dict__[x]))
print("-" * 80)
FONT = cv2.FONT_HERSHEY_TRIPLEX
COLORS = [
tuple(255 * np.array(plt.get_cmap("tab20")(i)[:-1])) for i in np.linspace(0, 1, 20)
]
cuda = torch.cuda.is_available() and opt.use_cuda
Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
# Get data configuration
data_config = parse_data_config(opt.data_config_path)
num_classes = int(data_config["classes"])
classes = load_names(data_config["names"])
for x, y in data_config.items():
print("%25s: %s" % (x, y))
print("-" * 80)
# Initiate model
model = Darknet(opt.model_config_path)
model.load_weights(opt.weights_path)
model.cuda()
model.eval()
img_size = int(model.hyperparams["height"])
for x, y in model.hyperparams.items():
print("%25s: %s" % (x, y))
print("-" * 80)
print("Model loading done")
# Get dataloader
while True:
test_path = input("Image path: ")
path, img = SingleImage(test_path, img_size=img_size)[0]
input_img = Variable(img.type(Tensor).unsqueeze(0))
# Get detections
with torch.no_grad():
detections = model(input_img)
detections = non_max_suppression(
detections, num_classes, opt.conf_thres, opt.nms_thres
)
detections = detections[0]
filtered_detections = []
for d in detections:
x1, y1, x2, y2 = Tensor(d[:4]).view(1, -1)[0]
area = (x2 - x1) * (y2 - y1) / (img_size * img_size)
if area >= opt.area_thres:
filtered_detections.append(d)
detections = torch.stack(filtered_detections)
img = cv2.imread(path)
# The amount of padding that was added
pad_x = max(img.shape[0] - img.shape[1], 0) * (img_size / max(img.shape))
pad_y = max(img.shape[1] - img.shape[0], 0) * (img_size / max(img.shape))
# Image height and width after padding is removed
unpad_h = img_size - pad_y
unpad_w = img_size - pad_x
# Draw bounding boxes and labels of detections
if detections is not None:
unique_labels = detections[:, -1].cpu().unique()
bbox_colors = random.sample(COLORS, len(unique_labels))
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
area = (x2 - x1) * (y2 - y1) / (img_size * img_size)
# rescale coordinates to original dimensions
x1 = ((x1 - pad_x // 2) / unpad_w) * img.shape[1]
y1 = ((y1 - pad_y // 2) / unpad_h) * img.shape[0]
x2 = ((x2 - pad_x // 2) / unpad_w) * img.shape[1]
y2 = ((y2 - pad_y // 2) / unpad_h) * img.shape[0]
color = bbox_colors[int(np.where(unique_labels == int(cls_pred))[0])]
# draw bbox over image
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
# add label
cv2.putText(
img, str(int(cls_pred)), (x1, y1 - 3), FONT, 1, (255, 255, 255), 1
)
print(
"\t+ Coords: [%4d, %4d, %4d, %4d], Class: %s, ObjConf: %.5f, ClassProb: %.5f, Area: %.5f"
% (
x1,
y1,
x2,
y2,
classes[int(cls_pred)],
conf.item(),
cls_conf.item(),
area,
)
)
# Save generated image with detections
save_path = "detections/single.png"
cv2.imwrite(save_path, img)