From 9d12a162f898b04eb64f0fb3366eeea7a6bbc8ed Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 10 Feb 2019 21:01:49 +0100 Subject: [PATCH] updates --- detect.py | 23 ++++------------------- utils/utils.py | 18 ++++++++++++++++-- 2 files changed, 20 insertions(+), 21 deletions(-) diff --git a/detect.py b/detect.py index 167dc587..a707b9e0 100755 --- a/detect.py +++ b/detect.py @@ -8,12 +8,8 @@ from utils.utils import * from utils import torch_utils -def unletterbox(img0_shape, letterbox_shape): - return None - - -def detect(cfg, weights, images, output='output', img_size=416, conf_thres=0.3, nms_thres=0.45, - save_txt=False, save_images=True): +def detect(cfg, weights, images, output='output', img_size=416, conf_thres=0.3, nms_thres=0.45, save_txt=False, + save_images=True): device = torch_utils.select_device() os.system('rm -rf ' + output) @@ -59,12 +55,8 @@ def detect(cfg, weights, images, output='output', img_size=416, conf_thres=0.3, save_img_path = os.path.join(output, path.split('/')[-1]) save_txt_path = save_img_path + '.txt' - # The amount of padding that was added - pad_x = max(im0.shape[0] - im0.shape[1], 0) * (img_size / max(im0.shape)) - pad_y = max(im0.shape[1] - im0.shape[0], 0) * (img_size / max(im0.shape)) - # Image height and width after padding is removed - unpad_h = img_size - pad_y - unpad_w = img_size - pad_x + # Rescale boxes from 416 to true image size + detections[:, :4] = scale_coords(img_size, detections[:, :4], im0.shape) unique_classes = detections[:, -1].cpu().unique() for i in unique_classes: @@ -72,13 +64,6 @@ def detect(cfg, weights, images, output='output', img_size=416, conf_thres=0.3, print('%g %ss' % (n, classes[int(i)]), end=', ') for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections: - # Rescale coordinates to original dimensions - y1 = (((y1 - pad_y // 2) / unpad_h) * im0.shape[0]).round() - x1 = (((x1 - pad_x // 2) / unpad_w) * im0.shape[1]).round() - y2 = (((y2 - pad_y // 2) / unpad_h) * im0.shape[0]).round() - x2 = (((x2 - pad_x // 2) / unpad_w) * im0.shape[1]).round() - x1, y1, x2, y2 = max(x1, 0), max(y1, 0), max(x2, 0), max(y2, 0) - if save_txt: # Write to file with open(save_txt_path, 'a') as file: file.write('%g %g %g %g %g %g\n' % (x1, y1, x2, y2, cls_pred, cls_conf * conf)) diff --git a/utils/utils.py b/utils/utils.py index 9a7853e3..7209c9bf 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -71,7 +71,8 @@ def weights_init_normal(m): torch.nn.init.constant_(m.bias.data, 0.0) -def xyxy2xywh(x): # Convert bounding box format from [x1, y1, x2, y2] to [x, y, w, h] +def xyxy2xywh(x): + # Convert bounding box format from [x1, y1, x2, y2] to [x, y, w, h] y = torch.zeros(x.shape) if x.dtype is torch.float32 else np.zeros(x.shape) y[:, 0] = (x[:, 0] + x[:, 2]) / 2 y[:, 1] = (x[:, 1] + x[:, 3]) / 2 @@ -80,7 +81,8 @@ def xyxy2xywh(x): # Convert bounding box format from [x1, y1, x2, y2] to [x, y, return y -def xywh2xyxy(x): # Convert bounding box format from [x, y, w, h] to [x1, y1, x2, y2] +def xywh2xyxy(x): + # Convert bounding box format from [x, y, w, h] to [x1, y1, x2, y2] y = torch.zeros(x.shape) if x.dtype is torch.float32 else np.zeros(x.shape) y[:, 0] = (x[:, 0] - x[:, 2] / 2) y[:, 1] = (x[:, 1] - x[:, 3] / 2) @@ -89,6 +91,18 @@ def xywh2xyxy(x): # Convert bounding box format from [x, y, w, h] to [x1, y1, x return y +def scale_coords(img_size, coords, img0_shape): + # Rescale x1, y1, x2, y2 from 416 to image size + gain = float(img_size) / max(img0_shape) # gain = old / new + pad_x = (img_size - img0_shape[1] * gain) / 2 # width padding + pad_y = (img_size - img0_shape[0] * gain) / 2 # height padding + coords[:, [0, 2]] -= pad_x + coords[:, [1, 3]] -= pad_y + coords[:, :4] /= gain + coords[:, :4] = torch.round(torch.clamp(coords[:, :4], min=0)) + return coords + + def ap_per_class(tp, conf, pred_cls, target_cls): """ Compute the average precision, given the recall and precision curves. Method originally from https://github.com/rafaelpadilla/Object-Detection-Metrics.