263 lines
10 KiB
Python
263 lines
10 KiB
Python
import argparse
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from sklearn import metrics
|
|
from scipy import interpolate
|
|
import torch.nn.functional as F
|
|
from models import *
|
|
from utils.utils import *
|
|
from torchvision.transforms import transforms as T
|
|
from utils.datasets import LoadImagesAndLabels, JointDataset, collate_fn
|
|
|
|
def test(
|
|
cfg,
|
|
data_cfg,
|
|
weights,
|
|
batch_size=16,
|
|
iou_thres=0.5,
|
|
conf_thres=0.3,
|
|
nms_thres=0.45,
|
|
print_interval=40,
|
|
):
|
|
|
|
# Configure run
|
|
f = open(data_cfg)
|
|
data_cfg_dict = json.load(f)
|
|
f.close()
|
|
#nC = int(data_cfg_dict['classes']) # number of classes (80 for COCO)
|
|
nC = 1
|
|
test_path = data_cfg_dict['test']
|
|
dataset_root = data_cfg_dict['root']
|
|
cfg_dict = parse_model_cfg(cfg)
|
|
img_size = [int(cfg_dict[0]['width']), int(cfg_dict[0]['height'])]
|
|
|
|
# Initialize model
|
|
model = Darknet(cfg_dict, test_emb=False)
|
|
|
|
# Load weights
|
|
if weights.endswith('.pt'): # pytorch format
|
|
model.load_state_dict(torch.load(weights, map_location='cpu')['model'], strict=False)
|
|
else: # darknet format
|
|
load_darknet_weights(model, weights)
|
|
|
|
model = torch.nn.DataParallel(model)
|
|
model.cuda().eval()
|
|
|
|
# Get dataloader
|
|
transforms = T.Compose([T.ToTensor()])
|
|
dataset = JointDataset(dataset_root, test_path, img_size, augment=False, transforms=transforms)
|
|
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False,
|
|
num_workers=8, drop_last=False, collate_fn=collate_fn)
|
|
|
|
mean_mAP, mean_R, mean_P, seen = 0.0, 0.0, 0.0, 0
|
|
print('%11s' * 5 % ('Image', 'Total', 'P', 'R', 'mAP'))
|
|
outputs, mAPs, mR, mP, TP, confidence, pred_class, target_class, jdict = \
|
|
[], [], [], [], [], [], [], [], []
|
|
AP_accum, AP_accum_count = np.zeros(nC), np.zeros(nC)
|
|
for batch_i, (imgs, targets, paths, shapes, targets_len) in enumerate(dataloader):
|
|
t = time.time()
|
|
output = model(imgs.cuda())
|
|
output = non_max_suppression(output, conf_thres=conf_thres, nms_thres=nms_thres)
|
|
for i, o in enumerate(output):
|
|
if o is not None:
|
|
output[i] = o[:, :6]
|
|
|
|
# Compute average precision for each sample
|
|
targets = [targets[i][:int(l)] for i,l in enumerate(targets_len)]
|
|
for si, (labels, detections) in enumerate(zip(targets, output)):
|
|
seen += 1
|
|
|
|
if detections is None:
|
|
# If there are labels but no detections mark as zero AP
|
|
if labels.size(0) != 0:
|
|
mAPs.append(0), mR.append(0), mP.append(0)
|
|
continue
|
|
|
|
# Get detections sorted by decreasing confidence scores
|
|
detections = detections.cpu().numpy()
|
|
detections = detections[np.argsort(-detections[:, 4])]
|
|
|
|
|
|
# If no labels add number of detections as incorrect
|
|
correct = []
|
|
if labels.size(0) == 0:
|
|
# correct.extend([0 for _ in range(len(detections))])
|
|
mAPs.append(0), mR.append(0), mP.append(0)
|
|
continue
|
|
else:
|
|
target_cls = labels[:, 0]
|
|
|
|
# Extract target boxes as (x1, y1, x2, y2)
|
|
target_boxes = xywh2xyxy(labels[:, 2:6])
|
|
target_boxes[:, 0] *= img_size[0]
|
|
target_boxes[:, 2] *= img_size[0]
|
|
target_boxes[:, 1] *= img_size[1]
|
|
target_boxes[:, 3] *= img_size[1]
|
|
|
|
detected = []
|
|
for *pred_bbox, conf, obj_conf in detections:
|
|
obj_pred = 0
|
|
pred_bbox = torch.FloatTensor(pred_bbox).view(1, -1)
|
|
# Compute iou with target boxes
|
|
iou = bbox_iou(pred_bbox, target_boxes, x1y1x2y2=True)[0]
|
|
# Extract index of largest overlap
|
|
best_i = np.argmax(iou)
|
|
# If overlap exceeds threshold and classification is correct mark as correct
|
|
if iou[best_i] > iou_thres and obj_pred == labels[best_i, 0] and best_i not in detected:
|
|
correct.append(1)
|
|
detected.append(best_i)
|
|
else:
|
|
correct.append(0)
|
|
|
|
# Compute Average Precision (AP) per class
|
|
AP, AP_class, R, P = ap_per_class(tp=correct,
|
|
conf=detections[:, 4],
|
|
pred_cls=np.zeros_like(detections[:, 5]), # detections[:, 6]
|
|
target_cls=target_cls)
|
|
|
|
# Accumulate AP per class
|
|
AP_accum_count += np.bincount(AP_class, minlength=nC)
|
|
AP_accum += np.bincount(AP_class, minlength=nC, weights=AP)
|
|
|
|
# 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())
|
|
|
|
# Means of all images
|
|
mean_mAP = np.sum(mAPs) / ( AP_accum_count + 1E-16)
|
|
mean_R = np.sum(mR) / ( AP_accum_count + 1E-16)
|
|
mean_P = np.sum(mP) / (AP_accum_count + 1E-16)
|
|
|
|
if batch_i % print_interval==0:
|
|
# Print image mAP and running mean mAP
|
|
print(('%11s%11s' + '%11.3g' * 4 + 's') %
|
|
(seen, dataloader.dataset.nF, mean_P, mean_R, mean_mAP, time.time() - t))
|
|
# Print mAP per class
|
|
print('%11s' * 5 % ('Image', 'Total', 'P', 'R', 'mAP'))
|
|
|
|
print('AP: %-.4f\n\n' % (AP_accum[0] / (AP_accum_count[0] + 1E-16)))
|
|
|
|
# Return mAP
|
|
return mean_mAP, mean_R, mean_P
|
|
|
|
|
|
def test_emb(
|
|
cfg,
|
|
data_cfg,
|
|
weights,
|
|
batch_size=16,
|
|
iou_thres=0.5,
|
|
conf_thres=0.3,
|
|
nms_thres=0.45,
|
|
print_interval=40,
|
|
):
|
|
|
|
# Configure run
|
|
f = open(data_cfg)
|
|
data_cfg_dict = json.load(f)
|
|
f.close()
|
|
test_paths = data_cfg_dict['test_emb']
|
|
dataset_root = data_cfg_dict['root']
|
|
cfg_dict = parse_model_cfg(cfg)
|
|
img_size = [int(cfg_dict[0]['width']), int(cfg_dict[0]['height'])]
|
|
|
|
# Initialize model
|
|
model = Darknet(cfg_dict, test_emb=True)
|
|
|
|
# Load weights
|
|
if weights.endswith('.pt'): # pytorch format
|
|
model.load_state_dict(torch.load(weights, map_location='cpu')['model'], strict=False)
|
|
else: # darknet format
|
|
load_darknet_weights(model, weights)
|
|
|
|
model = torch.nn.DataParallel(model)
|
|
model.cuda().eval()
|
|
|
|
# Get dataloader
|
|
transforms = T.Compose([T.ToTensor()])
|
|
dataset = JointDataset(dataset_root, test_paths, img_size, augment=False, transforms=transforms)
|
|
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False,
|
|
num_workers=8, drop_last=False, collate_fn=collate_fn)
|
|
embedding, id_labels = [], []
|
|
print('Extracting pedestrain features...')
|
|
for batch_i, (imgs, targets, paths, shapes, targets_len) in enumerate(dataloader):
|
|
t = time.time()
|
|
output = model(imgs.cuda(), targets.cuda(), targets_len.cuda()).squeeze()
|
|
|
|
for out in output:
|
|
feat, label = out[:-1], out[-1].long()
|
|
if label != -1:
|
|
embedding.append(feat)
|
|
id_labels.append(label)
|
|
|
|
if batch_i % print_interval==0:
|
|
print('Extracting {}/{}, # of instances {}, time {:.2f} sec.'.format(batch_i, len(dataloader), len(id_labels), time.time() - t))
|
|
|
|
print('Computing pairwise similairity...')
|
|
if len(embedding) <1 :
|
|
return None
|
|
embedding = torch.stack(embedding, dim=0).cuda()
|
|
id_labels = torch.LongTensor(id_labels)
|
|
n = len(id_labels)
|
|
print(n, len(embedding))
|
|
assert len(embedding) == n
|
|
|
|
embedding = F.normalize(embedding, dim=1)
|
|
pdist = torch.mm(embedding, embedding.t()).cpu().numpy()
|
|
gt = id_labels.expand(n,n).eq(id_labels.expand(n,n).t()).numpy()
|
|
|
|
up_triangle = np.where(np.triu(pdist)- np.eye(n)*pdist !=0)
|
|
pdist = pdist[up_triangle]
|
|
gt = gt[up_triangle]
|
|
|
|
far_levels = [ 1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1]
|
|
far,tar,threshold = metrics.roc_curve(gt, pdist)
|
|
interp = interpolate.interp1d(far, tar)
|
|
tar_at_far = [interp(x) for x in far_levels]
|
|
for f,fa in enumerate(far_levels):
|
|
print('TPR@FAR={:.7f}: {:.4f}'.format(fa, tar_at_far[f]))
|
|
return tar_at_far
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser(prog='test.py')
|
|
parser.add_argument('--batch-size', type=int, default=40, help='size of each image batch')
|
|
parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path')
|
|
parser.add_argument('--data-cfg', type=str, default='cfg/ccmcpe.json', help='data config')
|
|
parser.add_argument('--weights', type=str, default='weights/latest.pt', help='path to weights file')
|
|
parser.add_argument('--iou-thres', type=float, default=0.5, help='iou threshold required to qualify as detected')
|
|
parser.add_argument('--conf-thres', type=float, default=0.3, help='object confidence threshold')
|
|
parser.add_argument('--nms-thres', type=float, default=0.5, help='iou threshold for non-maximum suppression')
|
|
parser.add_argument('--print-interval', type=int, default=10, help='size of each image dimension')
|
|
parser.add_argument('--test-emb', action='store_true', help='test embedding')
|
|
opt = parser.parse_args()
|
|
print(opt, end='\n\n')
|
|
|
|
with torch.no_grad():
|
|
if opt.test_emb:
|
|
res = test_emb(
|
|
opt.cfg,
|
|
opt.data_cfg,
|
|
opt.weights,
|
|
opt.batch_size,
|
|
opt.iou_thres,
|
|
opt.conf_thres,
|
|
opt.nms_thres,
|
|
opt.print_interval,
|
|
)
|
|
else:
|
|
mAP = test(
|
|
opt.cfg,
|
|
opt.data_cfg,
|
|
opt.weights,
|
|
opt.batch_size,
|
|
opt.iou_thres,
|
|
opt.conf_thres,
|
|
opt.nms_thres,
|
|
opt.print_interval,
|
|
)
|
|
|