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,21 +89,21 @@ def detect(
# Rescale coordinates to original dimensions # Rescale coordinates to original dimensions
box_h = ((y2 - y1) / unpad_h) * img.shape[0] box_h = ((y2 - y1) / unpad_h) * img.shape[0]
box_w = ((x2 - x1) / unpad_w) * img.shape[1] box_w = ((x2 - x1) / unpad_w) * img.shape[1]
y1 = (((y1 - pad_y // 2) / unpad_h) * img.shape[0]).round().item() y1 = (((y1 - pad_y // 2) / unpad_h) * img.shape[0]).round()
x1 = (((x1 - pad_x // 2) / unpad_w) * img.shape[1]).round().item() x1 = (((x1 - pad_x // 2) / unpad_w) * img.shape[1]).round()
x2 = (x1 + box_w).round().item() x2 = (x1 + box_w).round()
y2 = (y1 + box_h).round().item() y2 = (y1 + box_h).round()
x1, y1, x2, y2 = max(x1, 0), max(y1, 0), max(x2, 0), max(y2, 0) x1, y1, x2, y2 = max(x1, 0), max(y1, 0), max(x2, 0), max(y2, 0)
# write to file # write to file
if save_txt: if save_txt:
with open(results_txt_path, 'a') as file: with open(results_txt_path, 'a') as file:
file.write(('%g %g %g %g %g %g \n') % (x1, y1, x2, y2, cls_pred, cls_conf * conf)) file.write(('%g %g %g %g %g %g\n') % (x1, y1, x2, y2, cls_pred, cls_conf * conf))
if save_images: if save_images:
# Add the bbox to the plot # Add the bbox to the plot
label = '%s %.2f' % (classes[int(cls_pred)], conf) 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) plot_one_box([x1, y1, x2, y2], img, label=label, color=color)
if save_images: if save_images: