-
Notifications
You must be signed in to change notification settings - Fork 29
/
detect.py
136 lines (107 loc) · 5.42 KB
/
detect.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
import argparse
import time
from models import *
from utils.datasets import *
from utils.utils import *
cuda = torch.cuda.is_available()
device = torch.device('cuda:0' if cuda else 'cpu')
parser = argparse.ArgumentParser()
parser.add_argument('-image_folder', type=str, default='data/samples', help='path to images')
parser.add_argument('-output_folder', type=str, default='output', help='path to outputs')
parser.add_argument('-plot_flag', type=bool, default=True)
parser.add_argument('-txt_out', type=bool, default=False)
parser.add_argument('-cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path')
parser.add_argument('-weights_path', type=str, default='weights/latest.pt', help='weight file path')
parser.add_argument('-class_path', type=str, default='data/icdar.names', help='path to class label file')
parser.add_argument('-conf_thres', type=float, default=0.2, help='object confidence threshold')
parser.add_argument('-nms_thres', type=float, default=0.1, help='iou threshold for non-maximum suppression')
parser.add_argument('-batch_size', type=int, default=1, help='size of the batches')
parser.add_argument('-img_size', type=int, default=32 * 19, help='size of each image dimension')
opt = parser.parse_args()
print(opt)
def detect(opt):
os.system('rm -rf ' + opt.output_folder)
os.makedirs(opt.output_folder, exist_ok=True)
# Load model
model = Darknet(opt.cfg, opt.img_size)
weights_path = opt.weights_path
if weights_path.endswith('.weights'): # saved in darknet format
load_weights(model, weights_path)
else: # endswith('.pt'), saved in pytorch format
checkpoint = torch.load(weights_path, map_location='cpu')
model.load_state_dict(checkpoint['model'])
del checkpoint
model.to(device).eval()
# Set Dataloader
classes = load_classes(opt.class_path) # Extracts class labels from file
dataloader = load_images(opt.image_folder, batch_size=opt.batch_size, img_size=opt.img_size)
imgs = [] # Stores image paths
img_detections = [] # Stores detections for each image index
prev_time = time.time()
for batch_i, (img_paths, img) in enumerate(dataloader):
print(batch_i, img.shape, end=' ')
# Get detections
with torch.no_grad():
chip = torch.from_numpy(img).unsqueeze(0).to(device)
pred = model(chip)
pred = pred[pred[:, :, 8] > opt.conf_thres]
if len(pred) > 0:
detections = non_max_suppression(pred.unsqueeze(0), 0.1, opt.nms_thres)
img_detections.extend(detections)
imgs.extend(img_paths)
print('Batch %d... (Done %.3f s)' % (batch_i, time.time() - prev_time))
prev_time = time.time()
# Bounding-box colors
color_list = [[random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)] for _ in range(len(classes))]
if len(img_detections) == 0:
return
# Iterate through images and save plot of detections
for img_i, (path, detections) in enumerate(zip(imgs, img_detections)):
print("image %g: '%s'" % (img_i, path))
if opt.plot_flag:
img = cv2.imread(path)
# The amount of padding that was added
pad_x = max(img.shape[0] - img.shape[1], 0) * (opt.img_size / max(img.shape))
pad_y = max(img.shape[1] - img.shape[0], 0) * (opt.img_size / max(img.shape))
# Image height and width after padding is removed
unpad_h = opt.img_size - pad_y
unpad_w = opt.img_size - pad_x
# Draw bounding boxes and labels of detections
if detections is not None:
unique_classes = detections[:, -1].cpu().unique()
bbox_colors = random.sample(color_list, len(unique_classes))
# write results to .txt file
results_img_path = os.path.join(opt.output_folder, path.split('/')[-1])
results_txt_path = results_img_path + '.txt'
if os.path.isfile(results_txt_path):
os.remove(results_txt_path)
for i in unique_classes:
n = (detections[:, -1].cpu() == i).sum()
print('%g %ss' % (n, classes[int(i)]))
for P1_x, P1_y, P2_x, P2_y, P3_x, P3_y, P4_x, P4_y, conf, cls_conf, cls_pred in detections:
P1_y = max((((P1_y - pad_y // 2) / unpad_h) * img.shape[0]).round().item(), 0)
P1_x = max((((P1_x - pad_x // 2) / unpad_w) * img.shape[1]).round().item(), 0)
P2_y = max((((P2_y - pad_y // 2) / unpad_h) * img.shape[0]).round().item(), 0)
P2_x = max((((P2_x - pad_x // 2) / unpad_w) * img.shape[1]).round().item(), 0)
P3_y = max((((P3_y - pad_y // 2) / unpad_h) * img.shape[0]).round().item(), 0)
P3_x = max((((P3_x - pad_x // 2) / unpad_w) * img.shape[1]).round().item(), 0)
P4_y = max((((P4_y - pad_y // 2) / unpad_h) * img.shape[0]).round().item(), 0)
P4_x = max((((P4_x - pad_x // 2) / unpad_w) * img.shape[1]).round().item(), 0)
# write to file
if opt.txt_out:
with open(results_txt_path, 'a') as file:
file.write(('%g %g %g %g %g %g %g %g %g %g \n') % (P1_x, P1_y, P2_x, P2_y, P3_x, P3_y, P4_x, P4_y, cls_pred, cls_conf * conf))
if opt.plot_flag:
# Add the bbox to the plot
label = '%s %.2f' % (classes[int(cls_pred)], conf)
color = bbox_colors[int(np.where(unique_classes == int(cls_pred))[0])]
plot_one_box([P1_x, P1_y, P2_x, P2_y, P3_x, P3_y, P4_x, P4_y], img, label=None, color=color)
cv2.imshow(path.split('/')[-1], img)
cv2.waitKey(0)
cv2.destroyAllWindows()
if opt.plot_flag:
# Save generated image with detections
cv2.imwrite(results_img_path.replace('.bmp', '.jpg').replace('.tif', '.jpg'), img)
if __name__ == '__main__':
torch.cuda.empty_cache()
detect(opt)