from __future__ import print_function

import os

import argparse

import torch

import torch.backends.cudnn as cudnn

import numpy as np

from data import cfg_mnet, cfg_re50

from layers.functions.prior_box import PriorBox

from utils.nms.py_cpu_nms import py_cpu_nms

import cv2

from models.retinaface import RetinaFace

from utils.box_utils import decode, decode_landm

from utils.timer import Timer

parser = argparse.ArgumentParser(description='Retinaface')

parser.add_argument('-m', '--trained_model', default='./weights/Resnet50_Final.pth',

type=str, help='Trained state_dict file path to open')

parser.add_argument('--network', default='resnet50', help='Backbone network mobile0.25 or resnet50')

parser.add_argument('--origin_size', default=True, type=str, help='Whether use origin image size to evaluate')

parser.add_argument('--save_folder', default='./widerface_evaluate/widerface_txt/', type=str, help='Dir to save txt results')

parser.add_argument('--cpu', action="store_true", default=False, help='Use cpu inference')

parser.add_argument('--dataset_folder', default='./data/widerface/val/images/', type=str, help='dataset path')

parser.add_argument('--confidence_threshold', default=0.02, type=float, help='confidence_threshold')

parser.add_argument('--top_k', default=5000, type=int, help='top_k')

parser.add_argument('--nms_threshold', default=0.4, type=float, help='nms_threshold')

parser.add_argument('--keep_top_k', default=750, type=int, help='keep_top_k')

parser.add_argument('-s', '--save_image', action="store_true", default=False, help='show detection results')

parser.add_argument('--vis_thres', default=0.5, type=float, help='visualization_threshold')

args = parser.parse_args()

def check_keys(model, pretrained_state_dict):

return True

def remove_prefix(state_dict, prefix):

return {f(key): value for key, value in state_dict.items()}

def load_model(model, pretrained_path, load_to_cpu):

return model

if __name__ == '__main__':

torch.set_grad_enabled(False)

cfg = None

if args.network == "mobile0.25":

cfg = cfg_mnet

elif args.network == "resnet50":

cfg = cfg_re50

# net and model

net = RetinaFace(cfg=cfg, phase = 'test')

net = load_model(net, args.trained_model, args.cpu)

net.eval()

print('Finished loading model!')

print(net)

cudnn.benchmark = True

device = torch.device("cpu" if args.cpu else "cuda")

net = net.to(device)

# testing dataset

testset_folder = args.dataset_folder

testset_list = args.dataset_folder[:-7] + "wider_val.txt"

with open(testset_list, 'r') as fr:

test_dataset = fr.read().split()

num_images = len(test_dataset)

_t = {'forward_pass': Timer(), 'misc': Timer()}

# testing begin

for i, img_name in enumerate(test_dataset):

image_path = testset_folder + img_name

img_raw = cv2.imread(image_path, cv2.IMREAD_COLOR)

img = np.float32(img_raw)

# testing scale

target_size = 1600

max_size = 2150

im_shape = img.shape

im_size_min = np.min(im_shape[0:2])

im_size_max = np.max(im_shape[0:2])

resize = float(target_size) / float(im_size_min)

# prevent bigger axis from being more than max_size:

if np.round(resize * im_size_max) > max_size:

resize = float(max_size) / float(im_size_max)

if args.origin_size:

resize = 1

if resize != 1:

img = cv2.resize(img, None, None, fx=resize, fy=resize, interpolation=cv2.INTER_LINEAR)

im_height, im_width, _ = img.shape

scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]])

img -= (104, 117, 123)

img = img.transpose(2, 0, 1)

img = torch.from_numpy(img).unsqueeze(0)

img = img.to(device)

scale = scale.to(device)

_t['forward_pass'].tic()

loc, conf, landms = net(img) # forward pass

_t['forward_pass'].toc()

_t['misc'].tic()

priorbox = PriorBox(cfg, image_size=(im_height, im_width))

priors = priorbox.forward()

priors = priors.to(device)

prior_data = priors.data

boxes = decode(loc.data.squeeze(0), prior_data, cfg['variance'])

boxes = boxes * scale / resize

boxes = boxes.cpu().numpy()

scores = conf.squeeze(0).data.cpu().numpy()[:, 1]

landms = decode_landm(landms.data.squeeze(0), prior_data, cfg['variance'])

scale1 = torch.Tensor([img.shape[3], img.shape[2], img.shape[3], img.shape[2],

img.shape[3], img.shape[2], img.shape[3], img.shape[2],

img.shape[3], img.shape[2]])

scale1 = scale1.to(device)

landms = landms * scale1 / resize

landms = landms.cpu().numpy()

# ignore low scores

inds = np.where(scores > args.confidence_threshold)[0]

boxes = boxes[inds]

landms = landms[inds]

scores = scores[inds]

# keep top-K before NMS

order = scores.argsort()[::-1]

# order = scores.argsort()[::-1][:args.top_k]

boxes = boxes[order]

landms = landms[order]

scores = scores[order]

# do NMS

dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)

keep = py_cpu_nms(dets, args.nms_threshold)

# keep = nms(dets, args.nms_threshold,force_cpu=args.cpu)

dets = dets[keep, :]

landms = landms[keep]

# keep top-K faster NMS

# dets = dets[:args.keep_top_k, :]

# landms = landms[:args.keep_top_k, :]

dets = np.concatenate((dets, landms), axis=1)

_t['misc'].toc()

# --------------------------------------------------------------------

save_name = args.save_folder + img_name[:-4] + ".txt"

dirname = os.path.dirname(save_name)

if not os.path.isdir(dirname):

os.makedirs(dirname)

with open(save_name, "w") as fd:

bboxs = dets

file_name = os.path.basename(save_name)[:-4] + "\n"

bboxs_num = str(len(bboxs)) + "\n"

fd.write(file_name)

fd.write(bboxs_num)

for box in bboxs:

x = int(box[0])

y = int(box[1])

w = int(box[2]) - int(box[0])

h = int(box[3]) - int(box[1])

confidence = str(box[4])

line = str(x) + " " + str(y) + " " + str(w) + " " + str(h) + " " + confidence + " \n"

fd.write(line)

print('im_detect: {:d}/{:d} forward_pass_time: {:.4f}s misc: {:.4f}s'.format(i + 1, num_images, _t['forward_pass'].average_time, _t['misc'].average_time))

# save image

if args.save_image:

for b in dets:

if b[4] < args.vis_thres:

continue

text = "{:.4f}".format(b[4])

b = list(map(int, b))

cv2.rectangle(img_raw, (b[0], b[1]), (b[2], b[3]), (0, 0, 255), 2)

cx = b[0]

cy = b[1] + 12

cv2.putText(img_raw, text, (cx, cy),

cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255))

# landms

cv2.circle(img_raw, (b[5], b[6]), 1, (0, 0, 255), 4)

cv2.circle(img_raw, (b[7], b[8]), 1, (0, 255, 255), 4)

cv2.circle(img_raw, (b[9], b[10]), 1, (255, 0, 255), 4)

cv2.circle(img_raw, (b[11], b[12]), 1, (0, 255, 0), 4)

cv2.circle(img_raw, (b[13], b[14]), 1, (255, 0, 0), 4)

# save image

if not os.path.exists("./results/"):

os.makedirs("./results/")

name = "./results/" + str(i) + ".jpg"

cv2.imwrite(name, img_raw)