Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

updates

This commit is contained in:
Glenn Jocher 2019-02-10 21:01:49 +01:00
parent e057f52780
commit 9d12a162f8
2 changed files with 20 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
detect.py
View File

@ -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))

18
utils/utils.py
View File

@ -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.