mirror of
https://github.com/deepinsight/insightface.git
synced 2026-05-14 04:12:35 +00:00
257 lines
9.7 KiB
Python
Executable File
257 lines
9.7 KiB
Python
Executable File
import argparse
|
|
import os
|
|
import os.path as osp
|
|
import pickle
|
|
import numpy as np
|
|
import warnings
|
|
|
|
import mmcv
|
|
import torch
|
|
from mmcv import Config, DictAction
|
|
from mmcv.cnn import fuse_conv_bn
|
|
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
|
|
from mmcv.runner import (get_dist_info, init_dist, load_checkpoint,
|
|
wrap_fp16_model)
|
|
|
|
from mmdet.apis import multi_gpu_test, single_gpu_test
|
|
from mmdet.datasets import (build_dataloader, build_dataset,
|
|
replace_ImageToTensor)
|
|
from mmdet.models import build_detector
|
|
from mmdet.core.evaluation import wider_evaluation, get_widerface_gts
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(
|
|
description='MMDet test (and eval) a model')
|
|
parser.add_argument('config', help='test config file path')
|
|
parser.add_argument('checkpoint', help='checkpoint file')
|
|
parser.add_argument('--out', default='wout', help='output folder')
|
|
parser.add_argument(
|
|
'--eval',
|
|
type=str,
|
|
nargs='+',
|
|
help='evaluation metrics, which depends on the dataset, e.g., "bbox",'
|
|
' "segm", "proposal" for COCO, and "mAP", "recall" for PASCAL VOC')
|
|
parser.add_argument('--show', action='store_true', help='show results')
|
|
parser.add_argument('--save-preds', action='store_true', help='save results')
|
|
parser.add_argument('--show-assign', action='store_true', help='show bbox assign')
|
|
parser.add_argument('--debug', action='store_true', help='debug flag')
|
|
parser.add_argument(
|
|
'--show-dir', help='directory where painted images will be saved')
|
|
parser.add_argument(
|
|
'--show-score-thr',
|
|
type=float,
|
|
default=0.3,
|
|
help='score threshold (default: 0.3)')
|
|
parser.add_argument(
|
|
'--thr',
|
|
type=float,
|
|
default=0.02,
|
|
help='score threshold')
|
|
parser.add_argument('--local_rank', type=int, default=0)
|
|
parser.add_argument('--mode', type=int, default=0)
|
|
args = parser.parse_args()
|
|
if 'LOCAL_RANK' not in os.environ:
|
|
os.environ['LOCAL_RANK'] = str(args.local_rank)
|
|
|
|
return args
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
|
|
|
|
cfg = Config.fromfile(args.config)
|
|
if cfg.get('custom_imports', None):
|
|
from mmcv.utils import import_modules_from_strings
|
|
import_modules_from_strings(**cfg['custom_imports'])
|
|
# set cudnn_benchmark
|
|
if cfg.get('cudnn_benchmark', False):
|
|
torch.backends.cudnn.benchmark = True
|
|
cfg.model.pretrained = None
|
|
if cfg.model.get('neck'):
|
|
if isinstance(cfg.model.neck, list):
|
|
for neck_cfg in cfg.model.neck:
|
|
if neck_cfg.get('rfp_backbone'):
|
|
if neck_cfg.rfp_backbone.get('pretrained'):
|
|
neck_cfg.rfp_backbone.pretrained = None
|
|
elif cfg.model.neck.get('rfp_backbone'):
|
|
if cfg.model.neck.rfp_backbone.get('pretrained'):
|
|
cfg.model.neck.rfp_backbone.pretrained = None
|
|
|
|
# in case the test dataset is concatenated
|
|
if isinstance(cfg.data.test, dict):
|
|
cfg.data.test.test_mode = True
|
|
elif isinstance(cfg.data.test, list):
|
|
for ds_cfg in cfg.data.test:
|
|
ds_cfg.test_mode = True
|
|
|
|
gt_path = os.path.join(os.path.dirname(cfg.data.test.ann_file), 'gt')
|
|
pipelines = cfg.data.test.pipeline
|
|
for pipeline in pipelines:
|
|
if pipeline.type=='MultiScaleFlipAug':
|
|
if args.mode==0: #640 scale
|
|
pipeline.img_scale = (640, 640)
|
|
elif args.mode==1: #for single scale in other pages
|
|
pipeline.img_scale = (1100, 1650)
|
|
elif args.mode==2: #original scale
|
|
pipeline.img_scale = None
|
|
pipeline.scale_factor = 1.0
|
|
elif args.mode>30:
|
|
pipeline.img_scale = (args.mode, args.mode)
|
|
transforms = pipeline.transforms
|
|
for transform in transforms:
|
|
if transform.type=='Pad':
|
|
if args.mode!=2:
|
|
transform.size = pipeline.img_scale
|
|
else:
|
|
transform.size = None
|
|
transform.size_divisor = 32
|
|
print(cfg.data.test.pipeline)
|
|
distributed = False
|
|
|
|
# build the dataloader
|
|
samples_per_gpu = cfg.data.test.pop('samples_per_gpu', 1)
|
|
if samples_per_gpu > 1:
|
|
# Replace 'ImageToTensor' to 'DefaultFormatBundle'
|
|
cfg.data.test.pipeline = replace_ImageToTensor(cfg.data.test.pipeline)
|
|
dataset = build_dataset(cfg.data.test)
|
|
data_loader = build_dataloader(
|
|
dataset,
|
|
samples_per_gpu=samples_per_gpu,
|
|
workers_per_gpu=cfg.data.workers_per_gpu,
|
|
dist=distributed,
|
|
shuffle=False)
|
|
|
|
cfg.test_cfg.score_thr = args.thr
|
|
|
|
# build the model and load checkpoint
|
|
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
|
|
fp16_cfg = cfg.get('fp16', None)
|
|
checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
|
|
if 'CLASSES' in checkpoint['meta']:
|
|
model.CLASSES = checkpoint['meta']['CLASSES']
|
|
else:
|
|
model.CLASSES = dataset.CLASSES
|
|
|
|
if args.show_assign:
|
|
gts_easy, gts_medium, gts_hard = get_widerface_gts(gt_path)
|
|
assign_stat = [0, 0]
|
|
gts_size = []
|
|
model = MMDataParallel(model, device_ids=[0])
|
|
model.eval()
|
|
results = {}
|
|
output_folder = args.out
|
|
if not os.path.exists(output_folder):
|
|
os.makedirs(output_folder)
|
|
dataset = data_loader.dataset
|
|
prog_bar = mmcv.ProgressBar(len(dataset))
|
|
for i, data in enumerate(data_loader):
|
|
with torch.no_grad():
|
|
result = model(return_loss=False, rescale=True, **data)
|
|
assert len(result)==1
|
|
batch_size = 1
|
|
result = result[0][0]
|
|
img_metas = data['img_metas'][0].data[0][0]
|
|
filepath = img_metas['ori_filename']
|
|
det_scale = img_metas['scale_factor'][0]
|
|
#print(img_metas)
|
|
ori_shape = img_metas['ori_shape']
|
|
img_width = ori_shape[1]
|
|
img_height = ori_shape[0]
|
|
_vec = filepath.split('/')
|
|
pa, pb = _vec[-2], _vec[1]
|
|
if pa not in results:
|
|
results[pa] = {}
|
|
xywh = result.copy()
|
|
w = xywh[:,2] - xywh[:,0]
|
|
h = xywh[:,3] - xywh[:,1]
|
|
xywh[:,2] = w
|
|
xywh[:,3] = h
|
|
|
|
event_name = pa
|
|
img_name = pb.rstrip('.jpg')
|
|
results[event_name][img_name] = xywh
|
|
if args.save_preds:
|
|
out_dir = os.path.join(output_folder, pa)
|
|
if not os.path.exists(out_dir):
|
|
os.makedirs(out_dir)
|
|
out_file = os.path.join(out_dir, pb.replace('jpg', 'txt'))
|
|
boxes = result
|
|
with open(out_file, 'w') as f:
|
|
name = '/'.join([pa, pb])
|
|
f.write("%s\n"%(name))
|
|
f.write("%d\n"%(boxes.shape[0]))
|
|
for b in range(boxes.shape[0]):
|
|
box = boxes[b]
|
|
f.write("%.5f %.5f %.5f %.5f %g\n"%(box[0], box[1], box[2]-box[0], box[3]-box[1], box[4]))
|
|
|
|
if args.show_assign:
|
|
assert args.mode==0
|
|
input_height, input_width = 640, 640
|
|
gt_hard = gts_hard[event_name][img_name]
|
|
#print(event_name, img_name, gt_hard.shape)
|
|
gt_bboxes = gt_hard * det_scale
|
|
bbox_width = gt_bboxes[:,2] - gt_bboxes[:,0]
|
|
bbox_height = gt_bboxes[:,3] - gt_bboxes[:,1]
|
|
bbox_area = bbox_width * bbox_height
|
|
gt_size = np.sqrt(bbox_area+0.0001)
|
|
gts_size += list(gt_size)
|
|
anchor_cxs = []
|
|
anchor_cys = []
|
|
for idx, stride in enumerate([8,16,32,64,128]):
|
|
height = input_height // stride
|
|
width = input_width // stride
|
|
anchor_centers = np.stack(np.mgrid[:height, :width][::-1], axis=-1).astype(np.float32)
|
|
anchor_centers = (anchor_centers * stride).reshape( (-1, 2) )
|
|
anchor_cx = anchor_centers[:,0]
|
|
anchor_cy = anchor_centers[:,1]
|
|
anchor_cxs += list(anchor_cx)
|
|
anchor_cys += list(anchor_cy)
|
|
anchor_cx = np.array(anchor_cxs, dtype=np.float32)
|
|
anchor_cy = np.array(anchor_cys, dtype=np.float32)
|
|
|
|
num_gts = gt_bboxes.shape[0]
|
|
num_anchors = anchor_cx.shape[0]
|
|
anchor_cx = np.broadcast_to(anchor_cx.reshape((1,-1)), (num_gts, num_anchors)).reshape(num_anchors, num_gts)
|
|
anchor_cy = np.broadcast_to(anchor_cy.reshape((1,-1)), (num_gts, num_anchors)).reshape(num_anchors, num_gts)
|
|
gt_x1 = gt_bboxes[:,0]
|
|
gt_y1 = gt_bboxes[:,1]
|
|
gt_x2 = gt_bboxes[:,2]
|
|
gt_y2 = gt_bboxes[:,3]
|
|
gt_cover = np.zeros( (gt_bboxes.shape[0], ), dtype=np.float32)
|
|
l_ = anchor_cx - gt_x1
|
|
t_ = anchor_cy - gt_y1
|
|
r_ = gt_x2 - anchor_cx
|
|
b_ = gt_y2 - anchor_cy
|
|
dist = np.stack([l_, t_, r_, b_], axis=1).min(axis=1)
|
|
gt_dist = dist.max(axis=0)
|
|
gt_dist = gt_dist / gt_size
|
|
center_thres = 0.01
|
|
#center_thres = -0.25
|
|
gt_cover_inds = np.where(gt_dist>center_thres)[0]
|
|
num_assigned = len(gt_cover_inds)
|
|
assign_stat[0] += num_gts
|
|
assign_stat[1] += num_assigned
|
|
|
|
|
|
|
|
|
|
|
|
for _ in range(batch_size):
|
|
prog_bar.update()
|
|
aps = wider_evaluation(results, gt_path, 0.5, args.debug)
|
|
with open(os.path.join(output_folder, 'aps'), 'w') as f:
|
|
f.write("%f,%f,%f\n"%(aps[0],aps[1],aps[2]))
|
|
print('APS:', aps)
|
|
if args.show_assign:
|
|
print('ASSIGN:', assign_stat)
|
|
gts_size = np.array(gts_size, dtype=np.float32)
|
|
gts_size = np.sort(gts_size)
|
|
assert len(gts_size)==assign_stat[0]
|
|
print(gts_size[assign_stat[0]//2])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|