# Apply NMS | # Apply NMS | ||||
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, | pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, | ||||
multi_label=False, classes=opt.classes, agnostic=opt.agnostic_nms) | |||||
fast=True, classes=opt.classes, agnostic=opt.agnostic_nms) | |||||
# Apply Classifier | # Apply Classifier | ||||
if classify: | if classify: |
augment=False, | augment=False, | ||||
model=None, | model=None, | ||||
dataloader=None, | dataloader=None, | ||||
multi_label=True, | |||||
fast=False, | |||||
verbose=False): # 0 fast, 1 accurate | verbose=False): # 0 fast, 1 accurate | ||||
# Initialize/load model and set device | # Initialize/load model and set device | ||||
if model is None: | if model is None: | ||||
# Run NMS | # Run NMS | ||||
t = torch_utils.time_synchronized() | t = torch_utils.time_synchronized() | ||||
output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, multi_label=multi_label) | |||||
output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, fast=fast) | |||||
t1 += torch_utils.time_synchronized() - t | t1 += torch_utils.time_synchronized() - t | ||||
# Statistics per image | # Statistics per image |
final_epoch = epoch + 1 == epochs | final_epoch = epoch + 1 == epochs | ||||
if not opt.notest or final_epoch: # Calculate mAP | if not opt.notest or final_epoch: # Calculate mAP | ||||
results, maps, times = test.test(opt.data, | results, maps, times = test.test(opt.data, | ||||
batch_size=batch_size, | |||||
imgsz=imgsz_test, | |||||
save_json=final_epoch and opt.data.endswith(os.sep + 'coco.yaml'), | |||||
model=ema.ema, | |||||
single_cls=opt.single_cls, | |||||
dataloader=testloader, | |||||
multi_label=ni > n_burn) | |||||
batch_size=batch_size, | |||||
imgsz=imgsz_test, | |||||
save_json=final_epoch and opt.data.endswith(os.sep + 'coco.yaml'), | |||||
model=ema.ema, | |||||
single_cls=opt.single_cls, | |||||
dataloader=testloader, | |||||
fast=ni > n_burn) | |||||
# Write | # Write | ||||
with open(results_file, 'a') as f: | with open(results_file, 'a') as f: | ||||
if save: | if save: | ||||
with open(results_file, 'r') as f: # create checkpoint | with open(results_file, 'r') as f: # create checkpoint | ||||
ckpt = {'epoch': epoch, | ckpt = {'epoch': epoch, | ||||
'best_fitness': best_fitness, | |||||
'training_results': f.read(), | |||||
'model': ema.ema.module if hasattr(model, 'module') else ema.ema, | |||||
'optimizer': None if final_epoch else optimizer.state_dict()} | |||||
'best_fitness': best_fitness, | |||||
'training_results': f.read(), | |||||
'model': ema.ema.module if hasattr(model, 'module') else ema.ema, | |||||
'optimizer': None if final_epoch else optimizer.state_dict()} | |||||
# Save last, best and delete | # Save last, best and delete | ||||
torch.save(ckpt, last) | torch.save(ckpt, last) |
from scipy.signal import butter, filtfilt | from scipy.signal import butter, filtfilt | ||||
from tqdm import tqdm | from tqdm import tqdm | ||||
from . import torch_utils, google_utils # torch_utils, google_utils | |||||
from . import torch_utils, google_utils # torch_utils, google_utils | |||||
# Set printoptions | # Set printoptions | ||||
torch.set_printoptions(linewidth=320, precision=5, profile='long') | torch.set_printoptions(linewidth=320, precision=5, profile='long') | ||||
return tcls, tbox, indices, anch | return tcls, tbox, indices, anch | ||||
def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, multi_label=True, classes=None, agnostic=False): | |||||
def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, fast=False, classes=None, agnostic=False): | |||||
""" | """ | ||||
Performs Non-Maximum Suppression on inference results | Performs Non-Maximum Suppression on inference results | ||||
Returns detections with shape: | Returns detections with shape: | ||||
nx6 (x1, y1, x2, y2, conf, cls) | nx6 (x1, y1, x2, y2, conf, cls) | ||||
""" | """ | ||||
nc = prediction[0].shape[1] - 5 # number of classes | |||||
# Settings | # Settings | ||||
merge = True # merge for best mAP | |||||
min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height | min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height | ||||
max_det = 300 # maximum number of detections per image | max_det = 300 # maximum number of detections per image | ||||
time_limit = 10.0 # seconds to quit after | time_limit = 10.0 # seconds to quit after | ||||
redundant = conf_thres == 0.001 # require redundant detections | |||||
redundant = True # require redundant detections | |||||
fast |= conf_thres > 0.001 # fast mode | |||||
if fast: | |||||
merge = False | |||||
multi_label = False | |||||
else: | |||||
merge = True # merge for best mAP (adds 0.5ms/img) | |||||
multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img) | |||||
t = time.time() | t = time.time() | ||||
nc = prediction[0].shape[1] - 5 # number of classes | |||||
multi_label &= nc > 1 # multiple labels per box | |||||
output = [None] * prediction.shape[0] | output = [None] * prediction.shape[0] | ||||
for xi, x in enumerate(prediction): # image index, image inference | for xi, x in enumerate(prediction): # image index, image inference | ||||
# Apply constraints | # Apply constraints | ||||
# x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height | |||||
x = x[x[:, 4] > conf_thres] # confidence | x = x[x[:, 4] > conf_thres] # confidence | ||||
# x = x[((x[:, 2:4] > min_wh) & (x[:, 2:4] < max_wh)).all(1)] # width-height | |||||
# If none remain process next image | # If none remain process next image | ||||
if not x.shape[0]: | if not x.shape[0]: |