def __init__(self, min_score=0.9, use_gpu=True):
# Minimum score to consider as a detection.
self.score_min = min_score
self.net = FaceBoxes(phase='test', size=None, num_classes=2)
self.use_gpu = use_gpu
self.logger = Logger()
def __init__(self, **opt):
net = FaceBoxes(phase='test', size=None, num_classes=2) # initialize detector
net = load_model(net, opt.get("weights", 'weights/FaceBoxes.pth'), True)
net.eval()
self.net = net
self.top_k = opt.get('top_k', 5000)
self.confidence_threshold = opt.get('confidence_threshold', 0.05)
self.nms_threshold = opt.get('nms_threshold', 0.3)
self.keep_top_k = opt.get('keep_top_k', 750)
def load_face_model(self):
torch.set_grad_enabled(False)
# net and model
net = FaceBoxes(phase='test', size=None,
num_classes=2) # initialize detector
net = load_model(net, self.args.trained_model, self.args.cpu)
net.eval()
cudnn.benchmark = True
self.device = torch.device("cpu" if self.args.cpu else "cuda")
self.net = net.to(self.device)
class Detect(object):
def __init__(self, path, device):
self.net = FaceBoxes(phase='test', size=None, num_classes=2).to(device)
self.net = load_model(self.net, path, False)
self.net.eval()
self.device = device
def get_bbox(self, img_raw):
img = torch.FloatTensor(img_raw).to(self.device)
im_height, im_width, _ = img.size()
scale = torch.FloatTensor([im_width, im_height, im_width,
im_height]).to(self.device)
img -= torch.FloatTensor((104, 117, 123)).to(self.device)
img = img.permute(2, 0, 1).unsqueeze(0)
loc, conf = self.net(img) # forward pass
priorbox = PriorBox(cfg, image_size=(im_height, im_width))
priors = priorbox.forward()
priors = priors.to(self.device)
prior_data = priors.data
boxes = decode(loc.data.squeeze(0), prior_data, cfg['variance'])
boxes = boxes * scale
boxes = boxes.cpu().numpy()
scores = conf.squeeze(0).data.cpu().numpy()[:, 1]
# ignore low scores
inds = np.where(scores > 0.05)[0]
boxes = boxes[inds]
scores = scores[inds]
# keep top-K before NMS
order = scores.argsort()[::-1][:5000]
boxes = boxes[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, 0.3, force_cpu=False)
dets = dets[keep, :]
# keep top-K faster NMS
dets = dets[:750, :]
bboxes = []
for b in dets:
if b[4] < 0.65:
continue
b = list(map(int, b))
bboxes.append((b[0], b[1], b[2], b[3]))
return bboxes
def main():
num_classes = CONFIG["num_classes"]
max_epoch = CONFIG["max_epoch"]
save_directory = CONFIG["save_directory"]
phase = "train"
device = "cuda" if torch.cuda.is_available() else "cpu"
model = FaceBoxes(num_classes, phase)
model.to(device)
if not os.path.exists(save_directory):
os.makedirs(save_directory)
train_net(model, device, max_epoch, save_directory)
def load_faceboxes():
"""
Load FaceBoxes model and weight in pytorch
"""
print('---------------------------------------------------')
pretrained_path = 'weights/FaceBoxes.pth'
net = FaceBoxes(phase='test', size=None,
num_classes=2) # initialize detector
net = load_model(net, pretrained_path)
net.eval()
print('Finished loading model')
print('---------------------------------------------------')
return net.cpu()
def __init__(self, cfg):
self.net = FaceBoxes(phase='test', size=None,
num_classes=2) # initialize detector
self.model = None
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
self.cfg = cfg
def __init__(self,
trained_model,
cpu=True,
nms_threshold=0.5,
top_k=1000,
confidence_threshold=0.8,
keep_top_k=10):
super(FaceDetector, self).__init__()
self.trained_model = trained_model
self.net = FaceBoxes(phase='test', size=None, num_classes=2)
self.net = load_model(self.net, trained_model, cpu)
self.net.eval()
print('Finished loading model', trained_model)
self.device = torch.device("cpu" if cpu else "cuda")
self.net = self.net.to(self.device)
self.nms_threshold = nms_threshold
self.top_k = top_k
self.confidence_threshold = confidence_threshold
self.keep_top_k = keep_top_k
self.cpu = cpu
def get_model():
pretrain_model = "/Users/yangjiang/temp/gpu1/FaceBoxes.pth"
num_classes = CONFIG["num_classes"]
params = torch.load(pretrain_model, map_location="cpu")
model = FaceBoxes(num_classes, "test")
model.load_state_dict(params)
model.eval()
return model
def __init__(self):
self.cfg = cfg
self.model = opt["model"]
self.use_cpu = opt["use_cpu"]
self.mean = opt["mean"]
self.val = opt["val"]
self.confidence_threshold = opt["confidence_threshold"]
self.nms_threshold = opt["nms_threshold"]
self.top_k = opt["top_k"]
self.keep_top_k = opt["keep_top_k"]
self.yuzhi = opt["yuzhi"]
self.device = torch.device("cpu" if self.use_cpu else "cuda")
self.weights = FaceBoxes(phase="test", size=None, num_classes=2)
self.net = self.load_model(self.weights, self.model, self.use_cpu)
def load_model(model=FaceBoxes(phase='test', size=None, num_classes=2),
pretrained_path=args.trained_model,
load_to_cpu=args.cpu):
# print('Loading pretrained model from {}'.format(pretrained_path))
if load_to_cpu:
pretrained_dict = torch.load(pretrained_path,
map_location=lambda storage, loc: storage)
else:
device = torch.cuda.current_device()
pretrained_dict = torch.load(
pretrained_path,
map_location=lambda storage, loc: storage.cuda(device))
if "state_dict" in pretrained_dict.keys():
pretrained_dict = remove_prefix(pretrained_dict['state_dict'],
'module.')
else:
pretrained_dict = remove_prefix(pretrained_dict, 'module.')
check_keys(model, pretrained_dict)
model.load_state_dict(pretrained_dict, strict=False)
model.eval()
return model
parser.add_argument('--save_folder', default='./weights/', help='Location to inference checkpoint models')
args = parser.parse_args()
if not os.path.exists(args.save_folder):
os.mkdir(args.save_folder)
img_dim = 1024
rgb_means = (104, 117, 123) #bgr order
num_classes = 2
batch_size = args.batch_size
weight_decay = args.weight_decay
gamma = args.gamma
momentum = args.momentum
gpu_train = cfg['gpu_train']
net = FaceBoxes('train', img_dim, num_classes)
print("Printing net...")
print(net)
if args.resume_net is not None:
print('Loading resume network...')
state_dict = torch.load(args.resume_net)
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
head = k[:7]
if head == 'module.':
name = k[7:] # remove `module.`
else:
name = k
img_dim = 1024 # only 1024 is supported
rgb_mean = (104, 117, 123) # bgr order
num_classes = 2
num_gpu = args.ngpu
num_workers = args.num_workers
batch_size = args.batch_size
momentum = args.momentum
weight_decay = args.weight_decay
initial_lr = args.lr
gamma = args.gamma
max_epoch = args.max_epoch
training_dataset = args.training_dataset
save_folder = args.save_folder
gpu_train = cfg['gpu_train']
#net为FaceBoxes网络
net = FaceBoxes('train', img_dim, num_classes)
print("Printing net...")
#print(net)
if args.resume_net is not None:
print('Loading resume network...')
state_dict = torch.load(args.resume_net)
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
head = k[:7]
if head == 'module.':
name = k[7:] # remove `module.`
else:
name = k
# chkpt = torch.load(ops.landmarks_model, map_location=device)
# landmarks_model.load_state_dict(chkpt)
chkpt = torch.load(ops.landmarks_model,
map_location=lambda storage, loc: storage)
landmarks_model.load_state_dict(chkpt)
landmarks_model.eval() # 设置为前向推断模式
print('load landmarks model : {}'.format(ops.landmarks_model))
print(
'\n/******************* landmarks model acc ******************/')
acc_model(ops, landmarks_model)
landmarks_model = landmarks_model.to(device)
#--------------------------------------------------------------------------- 构建人脸检测模型
# detect_model
detect_model = FaceBoxes(phase='test', size=None,
num_classes=2) # initialize detector
detect_model = load_model(detect_model, ops.detect_model, True)
detect_model.eval()
print('\n/******************* detect model acc ******************/')
acc_model(ops, detect_model)
detect_model = detect_model.to(device)
print('Finished loading model!')
# print(detect_model)
detect_model = detect_model.to(device)
video_capture = cv2.VideoCapture(ops.test_path)
resize = 1
with torch.no_grad():
class FaceBoxDetector:
def __init__(self, min_score=0.9, use_gpu=True):
# Minimum score to consider as a detection.
self.score_min = min_score
self.net = FaceBoxes(phase='test', size=None, num_classes=2)
self.use_gpu = use_gpu
self.logger = Logger()
def load_model(self, path_to_model):
self.logger.field('Loading pretrained model from', path_to_model)
device = torch.cuda.current_device()
pretrained_dict = torch.load(path_to_model, map_location=lambda storage, loc: storage.cuda(device))
if "state_dict" in pretrained_dict.keys():
pretrained_dict = self.remove_prefix(pretrained_dict['state_dict'], 'module.')
else:
pretrained_dict = self.remove_prefix(pretrained_dict, 'module.')
self.check_keys(self.net, pretrained_dict)
self.net.load_state_dict(pretrained_dict, strict=False)
self.net.eval()
if self.use_gpu:
self.net.cuda()
def detect(self, images) -> List[List[TrackingRegion]]:
frames = []
for img in images:
img = np.float32(img)
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)
if self.use_gpu:
img = img.cuda()
scale = scale.cuda()
out = self.net(img)
face_regions = self.nms_process(out, scale, im_height, im_width)
frames.append(face_regions)
return frames
def nms_process(self, network_output, scale, im_height, im_width) -> List[TrackingRegion]:
priorbox = PriorBox(cfg, network_output[2], (im_height, im_width), phase='test')
priors = priorbox.forward()
if self.use_gpu:
priors = priors.cuda()
loc, conf, _ = network_output
prior_data = priors.data
boxes = decode(loc.data.squeeze(0), prior_data, cfg['variance'])
boxes = boxes * scale
boxes = boxes.cpu().numpy()
scores = conf.data.cpu().numpy()[:, 1]
# ignore low scores
inds = np.where(scores > self.score_min)[0]
boxes = boxes[inds]
scores = scores[inds]
# keep top-K before NMS, top_k = 5
order = scores.argsort()[::-1][:5000]
boxes = boxes[order]
scores = scores[order]
# do NMS
dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
keep = nms(dets, 0.3, force_cpu=False)
dets = dets[keep, :]
# keep top-K faster NMS
dets = dets[:750, :]
regions = []
for i in range(dets.shape[0]):
face_region = TrackingRegion()
face_region.set_rect(left=dets[i, 0], top=dets[i, 1], right=dets[i, 2], bottom=dets[i, 3])
face_region.confidence = dets[i, 4]
face_region.data["class_id"] = "face"
regions.append(face_region)
return regions
def check_keys(self, model, pretrained_state_dict):
ckpt_keys = set(pretrained_state_dict.keys())
model_keys = set(model.state_dict().keys())
used_pretrained_keys = model_keys & ckpt_keys
unused_pretrained_keys = ckpt_keys - model_keys
missing_keys = model_keys - ckpt_keys
self.logger.field('Missing keys', len(missing_keys))
self.logger.field('Unused checkpoint keys', len(unused_pretrained_keys))
self.logger.field('Used keys', len(used_pretrained_keys))
assert len(used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint'
return True
@staticmethod
def remove_prefix(state_dict, prefix):
""" Old style model is stored with all names of parameters sharing common prefix 'module.' """
print('remove prefix \'{}\''.format(prefix))
f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else x
return {f(key): value for key, value in state_dict.items()}
class FaceDetector(object):
'''
Class for face detection
'''
def __init__(self,
trained_model,
cpu=True,
nms_threshold=0.5,
top_k=1000,
confidence_threshold=0.8,
keep_top_k=10):
super(FaceDetector, self).__init__()
self.trained_model = trained_model
self.net = FaceBoxes(phase='test', size=None, num_classes=2)
self.net = load_model(self.net, trained_model, cpu)
self.net.eval()
print('Finished loading model', trained_model)
self.device = torch.device("cpu" if cpu else "cuda")
self.net = self.net.to(self.device)
self.nms_threshold = nms_threshold
self.top_k = top_k
self.confidence_threshold = confidence_threshold
self.keep_top_k = keep_top_k
self.cpu = cpu
def predict(self, img_name):
img = np.float32(cv2.imread(img_name, cv2.IMREAD_COLOR))
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(self.device)
scale = scale.to(self.device)
_t = {'forward_pass': Timer(), 'misc': Timer()}
_t['forward_pass'].tic()
loc, conf = self.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(self.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.data.cpu().numpy()[:, 1]
# ignore low scores
inds = np.where(scores > self.confidence_threshold)[0]
boxes = boxes[inds]
scores = scores[inds]
# keep top-K before NMS
order = scores.argsort()[::-1][:self.top_k]
boxes = boxes[order]
scores = scores[order]
# do NMS
dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32,
copy=False)
#keep = py_cpu_nms(dets, self.nms_threshold)
keep = nms(dets, self.nms_threshold, force_cpu=self.cpu)
dets = dets[keep, :]
# keep top-K faster NMS
dets = dets[:self.keep_top_k, :]
_t['misc'].toc()
return dets
else:
device = torch.cuda.current_device()
pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage.cuda(device))
if "state_dict" in pretrained_dict.keys():
pretrained_dict = remove_prefix(pretrained_dict['state_dict'], 'module.')
else:
pretrained_dict = remove_prefix(pretrained_dict, 'module.')
check_keys(model, pretrained_dict)
model.load_state_dict(pretrained_dict, strict=False)
return model
if __name__ == '__main__':
torch.set_grad_enabled(False)
# net and model
net = FaceBoxes(phase='test', size=None, num_classes=2) # initialize detector
net = load_model(net, 'weights/Final_HandBoxes.pth', False)
net = torch.nn.DataParallel(net, [1,0]).cuda() # multiprocessing edler look here
net.to(f'cuda:{net.device_ids[0]}')
net.eval()
print('Finished loading model!')
#print(net)
cudnn.benchmark = True
device = torch.device(f'cuda:{net.device_ids[0]}')
#net = net.to(device)
# testing scale
resize = 2
_t = {'forward_pass': Timer(), 'misc': Timer()}
def main():
global args
global minmum_loss
args.gpu = 0
args.world_size = 1
if args.distributed:
args.gpu = args.local_rank % torch.cuda.device_count()
torch.cuda.set_device(args.gpu)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.total_batch_size = args.world_size * args.batch_size
model = FaceBoxes('train', args.num_classes)
print("Printing net...")
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(model)
model = model.cuda()
# optimizer and loss function
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
criterion = MultiBoxLoss(num_classes=args.num_classes,
overlap_thresh=0.35,
prior_for_matching=True,
bkg_label=0,
neg_mining=True,
neg_pos=7,
neg_overlap=0.35,
encode_target=False)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(
args.resume,
map_location=lambda storage, loc: storage.cuda(args.gpu))
args.start_epoch = checkpoint['epoch']
minmum_loss = checkpoint['minmum_loss']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
# Data loading code
print('Loading Dataset...')
dataset = VOCDetection(args.training_dataset,
preproc(args.img_dim, args.rgb_mean),
AnnotationTransform())
train_loader = data.DataLoader(dataset,
args.batch_size,
num_workers=args.num_workers,
shuffle=True,
collate_fn=detection_collate,
pin_memory=True)
priorbox = PriorBox(cfg, image_size=(args.img_dim, args.img_dim))
with torch.no_grad():
priors = priorbox.forward()
priors = priors.cuda()
for epoch in range(args.start_epoch, args.epochs):
# train for one epoch
end = time.time()
loss = train(train_loader, model, priors, criterion, optimizer, epoch)
if args.local_rank == 0:
is_best = loss < minmum_loss
minmum_loss = min(loss, minmum_loss)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': minmum_loss,
'optimizer': optimizer.state_dict(),
}, is_best, epoch)
epoch_time = time.time() - end
print('Epoch %s time cost %f' % (epoch, epoch_time))
def load_model(model, pretrained_path):
print('Loading pretrained model from {}'.format(pretrained_path))
device = torch.cuda.current_device()
pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage.cuda(device))
if "state_dict" in pretrained_dict.keys():
pretrained_dict = remove_prefix(pretrained_dict['state_dict'], 'module.')
else:
pretrained_dict = remove_prefix(pretrained_dict, 'module.')
check_keys(model, pretrained_dict)
model.load_state_dict(pretrained_dict, strict=False)
return model
model = FaceBoxes(phase='test', size=None, num_classes=2)
device = torch.device("cuda")
model = load_model(model, trained_model_path)
model = model.to(device)
model.eval()
image_path = '/media/haoxue/WD/FaceBoxes.PyTorch/data/FDDB/images/2002/08/26/big/img_265.jpg'
img = np.float32(cv2.imread(image_path, cv2.IMREAD_COLOR))
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)
# if args.cuda:
img = img.cuda()
def faceboxes(img_raw, cur_frame_counter):
img = np.float32(img_raw)
torch.set_grad_enabled(False)
# net and model
net = FaceBoxes(phase='test', size=None,
num_classes=2) # initialize detector
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 scale
resize = 2
_t = {'forward_pass': Timer(), 'misc': Timer()}
# testing begin
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]])
# print(img)
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)
loc, conf = net(img) # forward pass
# print(loc.size(),conf.size())
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]
# ignore low scores
inds = np.where(scores > args.confidence_threshold)[0]
boxes = boxes[inds]
scores = scores[inds]
# keep top-K before NMS
order = scores.argsort()[::-1][:args.top_k]
boxes = boxes[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, :]
# keep top-K faster NMS
dets = dets[:args.keep_top_k, :]
_t['misc'].toc()
outputs_useful = []
for b in dets:
output_traffic = {}
if b[4] < args.vis_thres:
continue
b = list(map(int, b))
(left, right, top, bottom) = (b[0], b[2], b[1], b[3])
label_str = 'face'
output_traffic[label_str] = [left, right, top, bottom]
outputs_useful.append(output_traffic)
# show image
if args.show_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))
cv2.imwrite('boxes/' + str(cur_frame_counter) + '.jpg', img_raw)
return outputs_useful
b = list(map(int, b))
crop_img = img_raw[b[1]:b[1]+b[3], b[0]:b[0]+b[2]]
crop_resized = cv2.resize(crop_img, (112, 112))
# 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))
cv2.imwrite(os.path.join(newdir, img_p, str(inx), '.jpg'), crop_img)
if __name__ == '__main__':
trained_model = ''
load_to_cpu = False
torch.set_grad_enabled(False)
net = FaceBoxes(phase='test', size=None, num_classes=2)
net = load_model(net, trained_model, load_to_cpu)
net.eval()
cudnn.benchmark = True
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = net.to(device)
root = "/aidata/dataset/faces/CASIA-FaceV5/train"
root_ = "/aidata/dataset/faces/CASIA-FaceV5_Crop"
pathinfo = getinfo(root)
img_crop_face(root_, pathinfo, mode="train")
def __init__(self, path, device):
self.net = FaceBoxes(phase='test', size=None, num_classes=2).to(device)
self.net = load_model(self.net, path, False)
self.net.eval()
self.device = device
import sys
from scipy.special import softmax
import torch.onnx
import onnxruntime as ort
import numpy as np
import tensorflow as tf
from tensorflow.keras import backend as K
from pytorch2keras.converter import pytorch_to_keras
from models.faceboxes import FaceBoxes
input_dim = 1024
num_classes = 2
model_path = "weights/FaceBoxesProd.pth"
net = FaceBoxes('train', input_dim, num_classes)
def check_keys(model, pretrained_state_dict):
ckpt_keys = set(pretrained_state_dict.keys())
model_keys = set(model.state_dict().keys())
used_pretrained_keys = model_keys & ckpt_keys
unused_pretrained_keys = ckpt_keys - model_keys
missing_keys = model_keys - ckpt_keys
print('Missing keys:{}'.format(len(missing_keys)))
print('Unused checkpoint keys:{}'.format(len(unused_pretrained_keys)))
print('Used keys:{}'.format(len(used_pretrained_keys)))
assert len(
used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint'
return True
pretrained_path,
map_location=lambda storage, loc: storage.cuda(device))
if "state_dict" in pretrained_dict.keys():
pretrained_dict = remove_prefix(pretrained_dict['state_dict'],
'module.')
else:
pretrained_dict = remove_prefix(pretrained_dict, 'module.')
check_keys(model, pretrained_dict)
model.load_state_dict(pretrained_dict, strict=False)
return model
if __name__ == '__main__':
torch.set_grad_enabled(False)
# net and model
net = FaceBoxes(phase='test', size=None,
num_classes=2) # initialize detector
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)
# save file
if not os.path.exists(args.save_folder):
os.makedirs(args.save_folder)
fw = open(os.path.join(args.save_folder, args.dataset + '_dets.txt'), 'w')
# testing dataset
testset_folder = os.path.join('data', args.dataset, 'images/')
gamma = args.gamma
max_epoch = args.max_epoch
training_dataset = args.training_dataset
save_folder = args.save_folder
gpu_train = cfg['gpu_train']
"""def get_model():
model = faceboxes.FacesBoxes(num_classes)
model.build(input_shape=(None, img_dim, image_width, channels))
model.summary()
return model
"""
net = FaceBoxes('train', img_dim, num_classes)
print("Printing net...")
print(net)
if __name__ == '__main__':
# GPU settings
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
# if args.resume_net is not None:
# print('Loading resume network...')
# state_dict = torch.load(args.resume_net)
# # create new OrderedDict that does not contain `module.`
