This commit is contained in:
Glenn Jocher 2019-02-08 23:03:27 +01:00
parent 334660d58f
commit e77de1c3c7
1 changed files with 6 additions and 6 deletions

View File

@ -89,10 +89,10 @@ def detect(
# Rescale coordinates to original dimensions
box_h = ((y2 - y1) / unpad_h) * img.shape[0]
box_w = ((x2 - x1) / unpad_w) * img.shape[1]
y1 = (((y1 - pad_y // 2) / unpad_h) * img.shape[0]).round().item()
x1 = (((x1 - pad_x // 2) / unpad_w) * img.shape[1]).round().item()
x2 = (x1 + box_w).round().item()
y2 = (y1 + box_h).round().item()
y1 = (((y1 - pad_y // 2) / unpad_h) * img.shape[0]).round()
x1 = (((x1 - pad_x // 2) / unpad_w) * img.shape[1]).round()
x2 = (x1 + box_w).round()
y2 = (y1 + box_h).round()
x1, y1, x2, y2 = max(x1, 0), max(y1, 0), max(x2, 0), max(y2, 0)
# write to file
@ -103,7 +103,7 @@ def detect(
if save_images:
# 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])]
color = bbox_colors[list(unique_classes).index(cls_pred)]
plot_one_box([x1, y1, x2, y2], img, label=label, color=color)
if save_images: