updates
This commit is contained in:
parent
7283f26f6f
commit
b181c61f4b
|
@ -24,7 +24,7 @@ def detect(
|
|||
device = torch_utils.select_device()
|
||||
print("Using device: \"{}\"".format(device))
|
||||
|
||||
# os.system('rm -rf ' + output)
|
||||
os.system('rm -rf ' + output)
|
||||
os.makedirs(output, exist_ok=True)
|
||||
|
||||
data_config = parse_data_config(data_config_path)
|
||||
|
@ -66,6 +66,7 @@ def detect(
|
|||
|
||||
# Get detections
|
||||
with torch.no_grad():
|
||||
# cv2.imwrite('zidane_416.jpg', 255 * img.transpose((1, 2, 0))[:, :, ::-1]) # letterboxed
|
||||
img = torch.from_numpy(img).unsqueeze(0).to(device)
|
||||
# pred = torch.onnx._export(model, img, 'weights/model.onnx', verbose=True); return # ONNX export
|
||||
pred = model(img)
|
||||
|
|
15
models.py
15
models.py
|
@ -89,7 +89,7 @@ class Upsample(torch.nn.Module):
|
|||
self.mode = mode
|
||||
|
||||
def forward(self, x):
|
||||
return nn.functional.interpolate(x, scale_factor=self.scale_factor, mode=self.mode)
|
||||
return F.interpolate(x, scale_factor=self.scale_factor, mode=self.mode)
|
||||
|
||||
|
||||
class YOLOLayer(nn.Module):
|
||||
|
@ -120,9 +120,10 @@ class YOLOLayer(nn.Module):
|
|||
nG = int(self.img_dim / stride) # number grid points
|
||||
self.grid_x = torch.arange(nG).repeat(nG, 1).view([1, 1, nG, nG]).float()
|
||||
self.grid_y = torch.arange(nG).repeat(nG, 1).t().view([1, 1, nG, nG]).float()
|
||||
self.scaled_anchors = torch.FloatTensor([(a_w / stride, a_h / stride) for a_w, a_h in anchors])
|
||||
self.anchor_w = self.scaled_anchors[:, 0:1].view((1, nA, 1, 1))
|
||||
self.anchor_h = self.scaled_anchors[:, 1:2].view((1, nA, 1, 1))
|
||||
self.grid_y = torch.arange(nG).repeat(nG, 1).t().view([1, 1, nG, nG]).float()
|
||||
self.anchor_wh = torch.FloatTensor([(a_w / stride, a_h / stride) for a_w, a_h in anchors]) # scale anchors
|
||||
self.anchor_w = self.anchor_wh[:, 0:1].view((1, nA, 1, 1))
|
||||
self.anchor_h = self.anchor_wh[:, 1:2].view((1, nA, 1, 1))
|
||||
self.weights = class_weights()
|
||||
|
||||
self.loss_means = torch.ones(6)
|
||||
|
@ -177,7 +178,7 @@ class YOLOLayer(nn.Module):
|
|||
gy + height / 2), 4) # x1y1x2y2
|
||||
|
||||
tx, ty, tw, th, mask, tcls, TP, FP, FN, TC = \
|
||||
build_targets(p_boxes, p_conf, p_cls, targets, self.scaled_anchors, self.nA, self.nC, nG, batch_report)
|
||||
build_targets(p_boxes, p_conf, p_cls, targets, self.anchor_wh, self.nA, self.nC, nG, batch_report)
|
||||
|
||||
tcls = tcls[mask]
|
||||
if x.is_cuda:
|
||||
|
@ -319,8 +320,8 @@ class Darknet(nn.Module):
|
|||
if ONNX_export:
|
||||
# Produce a single-layer *.onnx model (upsample ops not working in PyTorch 1.0 export yet)
|
||||
output = output[0].squeeze().transpose(0, 1) # first layer reshaped to 85 x 507
|
||||
output[5:] = torch.nn.functional.softmax(torch.sigmoid(output[5:]) * output[4:5], dim=0) # SSD-like conf
|
||||
return output[5:], output[:4] # ONNX scores, boxes
|
||||
output[5:85] = F.softmax(output[5:85], dim=0) * output[4:5] # SSD-like conf
|
||||
return output[5:85], output[:4] # ONNX scores, boxes
|
||||
|
||||
return sum(output) if is_training else torch.cat(output, 1)
|
||||
|
||||
|
|
|
@ -309,8 +309,6 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4):
|
|||
# cross-class NMS (experimental)
|
||||
cross_class_nms = False
|
||||
if cross_class_nms:
|
||||
# thresh = 0.85
|
||||
thresh = nms_thres
|
||||
a = pred.clone()
|
||||
_, indices = torch.sort(-a[:, 4], 0) # sort best to worst
|
||||
a = a[indices]
|
||||
|
@ -325,7 +323,7 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4):
|
|||
if len(close) > 0:
|
||||
close = close + i + 1
|
||||
iou = bbox_iou(a[i:i + 1, :4], a[close.squeeze(), :4].reshape(-1, 4), x1y1x2y2=False)
|
||||
bad = close[iou > thresh]
|
||||
bad = close[iou > nms_thres]
|
||||
|
||||
if len(bad) > 0:
|
||||
mask = torch.ones(len(a)).type(torch.ByteTensor)
|
||||
|
@ -333,13 +331,12 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4):
|
|||
a = a[mask]
|
||||
pred = a
|
||||
|
||||
x, y, w, h = pred[:, 0], pred[:, 1], pred[:, 2], pred[:, 3]
|
||||
a = w * h # area
|
||||
ar = w / (h + 1e-16) # aspect ratio
|
||||
|
||||
log_w, log_h, log_a, log_ar = torch.log(w), torch.log(h), torch.log(a), torch.log(ar)
|
||||
|
||||
# Experiment: Prior class size rejection
|
||||
# x, y, w, h = pred[:, 0], pred[:, 1], pred[:, 2], pred[:, 3]
|
||||
# a = w * h # area
|
||||
# ar = w / (h + 1e-16) # aspect ratio
|
||||
# n = len(w)
|
||||
# log_w, log_h, log_a, log_ar = torch.log(w), torch.log(h), torch.log(a), torch.log(ar)
|
||||
# shape_likelihood = np.zeros((n, 60), dtype=np.float32)
|
||||
# x = np.concatenate((log_w.reshape(-1, 1), log_h.reshape(-1, 1)), 1)
|
||||
# from scipy.stats import multivariate_normal
|
||||
|
@ -348,7 +345,7 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4):
|
|||
|
||||
class_prob, class_pred = torch.max(F.softmax(pred[:, 5:], 1), 1)
|
||||
|
||||
v = ((pred[:, 4] > conf_thres) & (class_prob > .3))
|
||||
v = ((pred[:, 4] > conf_thres) & (class_prob > .3)) # TODO examine arbitrary 0.3 thres here
|
||||
v = v.nonzero().squeeze()
|
||||
if len(v.shape) == 0:
|
||||
v = v.unsqueeze(0)
|
||||
|
@ -375,44 +372,43 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4):
|
|||
nms_style = 'OR' # 'AND' or 'OR' (classical)
|
||||
for c in unique_labels:
|
||||
# Get the detections with the particular class
|
||||
detections_class = detections[detections[:, -1] == c]
|
||||
det_class = detections[detections[:, -1] == c]
|
||||
# Sort the detections by maximum objectness confidence
|
||||
_, conf_sort_index = torch.sort(detections_class[:, 4], descending=True)
|
||||
detections_class = detections_class[conf_sort_index]
|
||||
_, conf_sort_index = torch.sort(det_class[:, 4], descending=True)
|
||||
det_class = det_class[conf_sort_index]
|
||||
# Perform non-maximum suppression
|
||||
max_detections = []
|
||||
det_max = []
|
||||
|
||||
if nms_style == 'OR': # Classical NMS
|
||||
while detections_class.shape[0]:
|
||||
while det_class.shape[0]:
|
||||
# Get detection with highest confidence and save as max detection
|
||||
max_detections.append(detections_class[0].unsqueeze(0))
|
||||
det_max.append(det_class[0].unsqueeze(0))
|
||||
# Stop if we're at the last detection
|
||||
if len(detections_class) == 1:
|
||||
if len(det_class) == 1:
|
||||
break
|
||||
# Get the IOUs for all boxes with lower confidence
|
||||
ious = bbox_iou(max_detections[-1], detections_class[1:])
|
||||
ious = bbox_iou(det_max[-1], det_class[1:])
|
||||
|
||||
# Remove detections with IoU >= NMS threshold
|
||||
detections_class = detections_class[1:][ious < nms_thres]
|
||||
det_class = det_class[1:][ious < nms_thres]
|
||||
|
||||
elif nms_style == 'AND': # 'AND'-style NMS, at least two boxes must share commonality to pass, single boxes erased
|
||||
while detections_class.shape[0]:
|
||||
if len(detections_class) == 1:
|
||||
elif nms_style == 'AND': # 'AND'-style NMS: >=2 boxes must share commonality to pass, single boxes erased
|
||||
while det_class.shape[0]:
|
||||
if len(det_class) == 1:
|
||||
break
|
||||
|
||||
ious = bbox_iou(detections_class[:1], detections_class[1:])
|
||||
ious = bbox_iou(det_class[:1], det_class[1:])
|
||||
|
||||
if ious.max() > 0.5:
|
||||
max_detections.append(detections_class[0].unsqueeze(0))
|
||||
det_max.append(det_class[0].unsqueeze(0))
|
||||
|
||||
# Remove detections with IoU >= NMS threshold
|
||||
detections_class = detections_class[1:][ious < nms_thres]
|
||||
det_class = det_class[1:][ious < nms_thres]
|
||||
|
||||
if len(max_detections) > 0:
|
||||
max_detections = torch.cat(max_detections).data
|
||||
if len(det_max) > 0:
|
||||
det_max = torch.cat(det_max).data
|
||||
# Add max detections to outputs
|
||||
output[image_i] = max_detections if output[image_i] is None else torch.cat(
|
||||
(output[image_i], max_detections))
|
||||
output[image_i] = det_max if output[image_i] is None else torch.cat((output[image_i], det_max))
|
||||
|
||||
return output
|
||||
|
||||
|
@ -426,6 +422,7 @@ def strip_optimizer_from_checkpoint(filename='weights/best.pt'):
|
|||
|
||||
|
||||
def coco_class_count(path='../coco/labels/train2014/'):
|
||||
# histogram of occurrences per class
|
||||
import glob
|
||||
|
||||
nC = 80 # number classes
|
||||
|
@ -443,6 +440,7 @@ def plot_results():
|
|||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
# import os; os.system('rm -rf results.txt && wget https://storage.googleapis.com/ultralytics/results_v1_0.txt')
|
||||
|
||||
plt.figure(figsize=(16, 8))
|
||||
s = ['X', 'Y', 'Width', 'Height', 'Objectness', 'Classification', 'Total Loss', 'Precision', 'Recall', 'mAP']
|
||||
files = sorted(glob.glob('results*.txt'))
|
||||
|
|
Loading…
Reference in New Issue