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_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)
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/') testset_list = os.path.join('data', args.dataset, 'img_list.txt') with open(testset_list, 'r') as fr: test_dataset = fr.read().split() num_images = len(test_dataset) # testing scale if args.dataset == "FDDB":
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(): idx = 0 while True: ret, img_raw = video_capture.read() if ret:
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
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() scale = scale.cuda() out = model(img) priorbox = PriorBox(cfg, out[2], (im_height, im_width), phase='test')
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
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 net = load_model(net, model_path, False) net.eval() net.to("cuda") model_name = model_path.split("/")[-1].split(".")[0] onnx_model_path = f"models/onnx/base-model.onnx" # export ONNX model dummy_input = torch.randn(1, 3, input_dim, input_dim).to("cuda") torch.onnx.export(net, dummy_input, onnx_model_path, verbose=False, input_names=['input'], output_names=['output']) """ # try using pytorch2keras keras_model = pytorch_to_keras(net, dummy_input, [(3, input_dim, input_dim)])
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()} cap = cv2.VideoCapture(0)
# 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 new_state_dict[name] = v net.load_state_dict(new_state_dict) if args.ngpu > 1: net = torch.nn.DataParallel(net, device_ids=list(range(args.ngpu))) net.to(device) cudnn.benchmark = True optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) criterion_m = CrossEntropyLoss2d() criterion = MultiBoxLoss(num_classes, 0.35, True, 0, True, 7, 0.35, False) priorbox = PriorBox(cfg) with torch.no_grad(): priors = priorbox.forward() priors = priors.to(device)