car-detection-bayes/test.py

135 lines
5.8 KiB
Python
Raw Normal View History

2018-08-26 08:51:39 +00:00
import argparse
2018-10-10 15:07:21 +00:00
2018-08-26 08:51:39 +00:00
from models import *
from utils.datasets import *
from utils.utils import *
2018-11-22 19:02:11 +00:00
parser = argparse.ArgumentParser(prog='test.py')
2018-11-23 18:45:39 +00:00
parser.add_argument('-batch_size', type=int, default=32, help='size of each image batch')
2018-08-26 15:09:10 +00:00
parser.add_argument('-cfg', type=str, default='cfg/yolov3.cfg', help='path to model config file')
parser.add_argument('-data_config_path', type=str, default='cfg/coco.data', help='path to data config file')
2018-10-26 22:42:34 +00:00
parser.add_argument('-weights_path', type=str, default='weights/yolov3.pt', help='path to weights file')
2018-08-26 15:09:10 +00:00
parser.add_argument('-class_path', type=str, default='data/coco.names', help='path to class label file')
parser.add_argument('-iou_thres', type=float, default=0.5, help='iou threshold required to qualify as detected')
2018-11-14 15:14:41 +00:00
parser.add_argument('-conf_thres', type=float, default=0.3, help='object confidence threshold')
2018-09-10 15:00:39 +00:00
parser.add_argument('-nms_thres', type=float, default=0.45, help='iou threshold for non-maximum suppression')
2018-08-26 15:09:10 +00:00
parser.add_argument('-n_cpu', type=int, default=0, help='number of cpu threads to use during batch generation')
parser.add_argument('-img_size', type=int, default=416, help='size of each image dimension')
2018-08-26 08:51:39 +00:00
opt = parser.parse_args()
2018-11-22 12:52:22 +00:00
print(opt, end='\n\n')
2018-08-26 08:51:39 +00:00
2018-09-24 01:06:04 +00:00
cuda = torch.cuda.is_available()
2018-08-26 08:51:39 +00:00
device = torch.device('cuda:0' if cuda else 'cpu')
2018-11-14 15:14:41 +00:00
def main(opt):
# Configure run
data_config = parse_data_config(opt.data_config_path)
nC = int(data_config['classes']) # number of classes (80 for COCO)
if platform == 'darwin': # MacOS (local)
test_path = data_config['valid']
else: # linux (cloud, i.e. gcp)
test_path = '../coco/5k.part'
# Initiate model
model = Darknet(opt.cfg, opt.img_size)
# Load weights
if opt.weights_path.endswith('.weights'): # darknet format
load_weights(model, opt.weights_path)
elif opt.weights_path.endswith('.pt'): # pytorch format
checkpoint = torch.load(opt.weights_path, map_location='cpu')
model.load_state_dict(checkpoint['model'])
del checkpoint
model.to(device).eval()
# Get dataloader
# dataset = load_images_with_labels(test_path)
# dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batch_size, shuffle=False, num_workers=opt.n_cpu)
dataloader = load_images_and_labels(test_path, batch_size=opt.batch_size, img_size=opt.img_size)
2018-11-22 12:52:22 +00:00
print('%11s' * 5 % ('Image', 'Total', 'P', 'R', 'mAP'))
outputs, mAPs, mR, mP, TP, confidence, pred_class, target_class = [], [], [], [], [], [], [], []
2018-11-14 15:14:41 +00:00
AP_accum, AP_accum_count = np.zeros(nC), np.zeros(nC)
for batch_i, (imgs, targets) in enumerate(dataloader):
with torch.no_grad():
2018-11-23 17:09:47 +00:00
output = model(imgs.to(device))
2018-11-14 15:14:41 +00:00
output = non_max_suppression(output, conf_thres=opt.conf_thres, nms_thres=opt.nms_thres)
# Compute average precision for each sample
2018-11-23 17:09:47 +00:00
for sample_i, (labels, detections) in enumerate(zip(targets, output)):
2018-11-14 15:14:41 +00:00
correct = []
if detections is None:
2018-11-23 17:09:47 +00:00
# If there are no detections but there are labels mask as zero AP
if labels.size(0) != 0:
mAPs.append(0), mR.append(0), mP.append(0)
2018-11-14 15:14:41 +00:00
continue
# Get detections sorted by decreasing confidence scores
2018-11-23 17:09:47 +00:00
detections = detections.cpu().numpy()
2018-11-14 15:14:41 +00:00
detections = detections[np.argsort(-detections[:, 4])]
2018-11-23 17:09:47 +00:00
# If no labels add number of detections as incorrect
if labels.size(0) == 0:
2018-11-14 15:14:41 +00:00
# correct.extend([0 for _ in range(len(detections))])
2018-11-23 17:09:47 +00:00
mAPs.append(0), mR.append(0), mP.append(0)
2018-11-14 15:14:41 +00:00
continue
else:
2018-11-23 17:09:47 +00:00
target_cls = labels[:, 0]
2018-11-14 15:14:41 +00:00
# Extract target boxes as (x1, y1, x2, y2)
2018-11-23 17:09:47 +00:00
target_boxes = xywh2xyxy(labels[:, 1:5]) * opt.img_size
2018-11-14 15:14:41 +00:00
detected = []
for *pred_bbox, conf, obj_conf, obj_pred in detections:
pred_bbox = torch.FloatTensor(pred_bbox).view(1, -1)
# Compute iou with target boxes
iou = bbox_iou(pred_bbox, target_boxes)
# Extract index of largest overlap
best_i = np.argmax(iou)
# If overlap exceeds threshold and classification is correct mark as correct
2018-11-23 17:09:47 +00:00
if iou[best_i] > opt.iou_thres and obj_pred == labels[best_i, 0] and best_i not in detected:
2018-11-14 15:14:41 +00:00
correct.append(1)
detected.append(best_i)
else:
correct.append(0)
# Compute Average Precision (AP) per class
2018-11-22 12:52:22 +00:00
AP, AP_class, R, P = ap_per_class(tp=correct, conf=detections[:, 4], pred_cls=detections[:, 6],
target_cls=target_cls)
2018-11-14 15:14:41 +00:00
# Accumulate AP per class
AP_accum_count += np.bincount(AP_class, minlength=nC)
AP_accum += np.bincount(AP_class, minlength=nC, weights=AP)
2018-11-22 12:52:22 +00:00
# Compute mean AP across all classes in this image, and append to image list
mAPs.append(AP.mean())
mR.append(R.mean())
mP.append(P.mean())
2018-11-14 15:14:41 +00:00
2018-11-22 12:52:22 +00:00
# Means of all images
2018-11-14 23:56:03 +00:00
mean_mAP = np.mean(mAPs)
2018-11-22 12:52:22 +00:00
mean_R = np.mean(mR)
mean_P = np.mean(mP)
2018-11-14 15:14:41 +00:00
# Print image mAP and running mean mAP
2018-11-23 17:09:47 +00:00
print(('%11s%11s' + '%11.3g' * 3) % (len(mAPs), dataloader.nF, mean_P, mean_R, mean_mAP))
2018-11-14 15:14:41 +00:00
# Print mAP per class
2018-11-22 15:42:58 +00:00
print('%11s' * 5 % ('Image', 'Total', 'P', 'R', 'mAP') + '\n\nmAP Per Class:')
2018-11-14 15:14:41 +00:00
classes = load_classes(opt.class_path) # Extracts class labels from file
for i, c in enumerate(classes):
print('%15s: %-.4f' % (c, AP_accum[i] / AP_accum_count[i]))
2018-11-22 15:42:58 +00:00
# Return mAP
2018-11-22 12:52:22 +00:00
return mean_mAP, mean_R, mean_P
2018-11-14 15:14:41 +00:00
if __name__ == '__main__':
mAP = main(opt)