updates
This commit is contained in:
parent
eeb546ed6f
commit
8ed89d8c88
12
detect.py
12
detect.py
|
@ -61,25 +61,19 @@ def detect(opt):
|
|||
imgs = [] # Stores image paths
|
||||
img_detections = [] # Stores detections for each image index
|
||||
prev_time = time.time()
|
||||
detections = None
|
||||
for batch_i, (img_paths, img) in enumerate(dataloader):
|
||||
print(batch_i, img.shape, end=' ')
|
||||
preds = []
|
||||
|
||||
# Get detections
|
||||
with torch.no_grad():
|
||||
# Normal orientation
|
||||
chip = torch.from_numpy(img).unsqueeze(0).to(device)
|
||||
pred = model(chip)
|
||||
pred = pred[pred[:, :, 4] > opt.conf_thres]
|
||||
|
||||
if len(pred) > 0:
|
||||
preds.append(pred.unsqueeze(0))
|
||||
|
||||
if len(preds) > 0:
|
||||
detections = non_max_suppression(torch.cat(preds, 1), opt.conf_thres, opt.nms_thres)
|
||||
img_detections.extend(detections)
|
||||
imgs.extend(img_paths)
|
||||
detections = non_max_suppression(pred.unsqueeze(0), opt.conf_thres, opt.nms_thres)
|
||||
img_detections.extend(detections)
|
||||
imgs.extend(img_paths)
|
||||
|
||||
print('Batch %d... (Done %.3fs)' % (batch_i, time.time() - prev_time))
|
||||
prev_time = time.time()
|
||||
|
|
|
@ -5,7 +5,7 @@ import numpy as np
|
|||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
# set printoptions
|
||||
# Set printoptions
|
||||
torch.set_printoptions(linewidth=1320, precision=5, profile='long')
|
||||
np.set_printoptions(linewidth=320, formatter={'float_kind': '{11.5g}'.format}) # format short g, %precision=5
|
||||
|
||||
|
@ -19,7 +19,7 @@ def load_classes(path):
|
|||
return names
|
||||
|
||||
|
||||
def modelinfo(model):
|
||||
def modelinfo(model): # Plots a line-by-line description of a PyTorch model
|
||||
nparams = sum(x.numel() for x in model.parameters())
|
||||
ngradients = sum(x.numel() for x in model.parameters() if x.requires_grad)
|
||||
print('\n%4s %70s %9s %12s %20s %12s %12s' % ('', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma'))
|
||||
|
@ -39,17 +39,17 @@ def xview_class_weights(indices): # weights of each class in the training set,
|
|||
return weights[indices]
|
||||
|
||||
|
||||
def plot_one_box(x, im, color=None, label=None, line_thickness=None):
|
||||
tl = line_thickness or round(0.003 * max(im.shape[0:2])) # line thickness
|
||||
def plot_one_box(x, img, color=None, label=None, line_thickness=None): # Plots one bounding box on image img
|
||||
tl = line_thickness or round(0.003 * max(img.shape[0:2])) # line thickness
|
||||
color = color or [random.randint(0, 255) for _ in range(3)]
|
||||
c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
|
||||
cv2.rectangle(im, c1, c2, color, thickness=tl)
|
||||
cv2.rectangle(img, c1, c2, color, thickness=tl)
|
||||
if label:
|
||||
tf = max(tl - 1, 1) # font thickness
|
||||
t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
|
||||
c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
|
||||
cv2.rectangle(im, c1, c2, color, -1) # filled
|
||||
cv2.putText(im, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
|
||||
cv2.rectangle(img, c1, c2, color, -1) # filled
|
||||
cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
|
||||
|
||||
|
||||
def weights_init_normal(m):
|
||||
|
@ -61,13 +61,22 @@ def weights_init_normal(m):
|
|||
torch.nn.init.constant_(m.bias.data, 0.0)
|
||||
|
||||
|
||||
def xyxy2xywh(box):
|
||||
xywh = np.zeros(box.shape)
|
||||
xywh[:, 0] = (box[:, 0] + box[:, 2]) / 2
|
||||
xywh[:, 1] = (box[:, 1] + box[:, 3]) / 2
|
||||
xywh[:, 2] = box[:, 2] - box[:, 0]
|
||||
xywh[:, 3] = box[:, 3] - box[:, 1]
|
||||
return xywh
|
||||
def xyxy2xywh(x): # Convert bounding box format from [x1, y1, x2, y2] to [x, y, w, h]
|
||||
y = np.zeros(x.shape)
|
||||
y[:, 0] = (x[:, 0] + x[:, 2]) / 2
|
||||
y[:, 1] = (x[:, 1] + x[:, 3]) / 2
|
||||
y[:, 2] = x[:, 2] - x[:, 0]
|
||||
y[:, 3] = x[:, 3] - x[:, 1]
|
||||
return y
|
||||
|
||||
|
||||
def xywh2xyxy(x): # Convert bounding box format from [x, y, w, h] to [x1, y1, x2, y2]
|
||||
y = np.zeros(x.shape)
|
||||
y[:, 0] = (x[:, 1] - x[:, 3] / 2)
|
||||
y[:, 1] = (x[:, 2] - x[:, 4] / 2)
|
||||
y[:, 2] = (x[:, 1] + x[:, 3] / 2)
|
||||
y[:, 3] = (x[:, 2] + x[:, 4] / 2)
|
||||
return y
|
||||
|
||||
|
||||
def compute_ap(recall, precision):
|
||||
|
@ -98,9 +107,6 @@ def compute_ap(recall, precision):
|
|||
|
||||
|
||||
def bbox_iou(box1, box2, x1y1x2y2=True):
|
||||
# if len(box1.shape) == 1:
|
||||
# box1 = box1.reshape(1, 4)
|
||||
|
||||
"""
|
||||
Returns the IoU of two bounding boxes
|
||||
"""
|
||||
|
|
Loading…
Reference in New Issue