def main(args): if args.backbone == "v2": from models.pfld import PFLDInference, AuxiliaryNet elif args.backbone == "v3": from models.mobilev3_pfld import PFLDInference, AuxiliaryNet elif args.backbone == "ghost": from models.ghost_pfld import PFLDInference, AuxiliaryNet elif args.backbone == "lite": from models.lite import PFLDInference, AuxiliaryNet else: raise ValueError("backbone is not implemented") checkpoint = torch.load(args.model_path, map_location=device) plfd_backbone = PFLDInference().to(device) plfd_backbone.load_state_dict(checkpoint['plfd_backbone'], strict=False) transform = transforms.Compose([transforms.ToTensor()]) wlfw_val_dataset = PFLDDatasets(args.test_dataset, transform, img_root=os.path.realpath('./data')) wlfw_val_dataloader = DataLoader(wlfw_val_dataset, batch_size=1, shuffle=False, num_workers=0) validate(wlfw_val_dataloader, plfd_backbone)
def main(args): checkpoint = torch.load(args.model_path, map_location=device) plfd_backbone = PFLDInference(args.r).to(device) plfd_backbone = nn.DataParallel(plfd_backbone) plfd_backbone.load_state_dict(checkpoint['plfd_backbone']) transform = transforms.Compose([transforms.ToTensor()]) wlfw_val_dataset = WLFWDatasets(args.test_dataset, transform) wlfw_val_dataloader = DataLoader(wlfw_val_dataset, batch_size=1, shuffle=False, num_workers=0) validate(wlfw_val_dataloader, plfd_backbone)
def InferLandmark(args): checkpoint = torch.load(args.model_path, map_location=device) plfd_backbone = PFLDInference().to(device) plfd_backbone.load_state_dict(checkpoint['plfd_backbone']) transform = transforms.Compose([transforms.ToTensor()]) wlfw_val_dataset = WLFWDatasetsInfer(args.test_dataset, transform) wlfw_val_dataloader = DataLoader(wlfw_val_dataset, batch_size=1, shuffle=False, num_workers=0) infer(wlfw_val_dataloader, plfd_backbone, args.presave)
def main(args): checkpoint = torch.load(args.model_path) plfd_backbone = PFLDInference().cuda() auxiliarynet = AuxiliaryNet().cuda() plfd_backbone.load_state_dict(checkpoint['plfd_backbone']) auxiliarynet.load_state_dict(checkpoint['auxiliarynet']) transform = transforms.Compose([transforms.ToTensor()]) wlfw_val_dataset = WLFWDatasets(args.test_dataset, transform) wlfw_val_dataloader = DataLoader( wlfw_val_dataset, batch_size=8, shuffle=False, num_workers=0) validate(wlfw_val_dataloader, plfd_backbone, auxiliarynet)
def main(args): checkpoint = torch.load(args.model_path, map_location=device) plfd_backbone = PFLDInference().to(device) plfd_backbone.load_state_dict(checkpoint['plfd_backbone']) plfd_backbone.eval() plfd_backbone = plfd_backbone.to(device) transform = transforms.Compose([transforms.ToTensor()]) cap = cv2.VideoCapture(0) while True: ret, img = cap.read() if not ret: break height, width = img.shape[:2] bounding_boxes, landmarks = detect_faces(img) for box in bounding_boxes: score = box[4] x1, y1, x2, y2 = (box[:4] + 0.5).astype(np.int32) w = x2 - x1 + 1 h = y2 - y1 + 1 size = int(max([w, h]) * 1.1) cx = x1 + w // 2 cy = y1 + h // 2 x1 = cx - size // 2 x2 = x1 + size y1 = cy - size // 2 y2 = y1 + size dx = max(0, -x1) dy = max(0, -y1) x1 = max(0, x1) y1 = max(0, y1) edx = max(0, x2 - width) edy = max(0, y2 - height) x2 = min(width, x2) y2 = min(height, y2) cropped = img[y1:y2, x1:x2] if dx > 0 or dy > 0 or edx > 0 or edy > 0: cropped = cv2.copyMakeBorder(cropped, dy, edy, dx, edx, cv2.BORDER_CONSTANT, 0) cropped = cv2.resize(cropped, (112, 112)) input = cv2.resize(cropped, (112, 112)) input = cv2.cvtColor(input, cv2.COLOR_BGR2RGB) input = transform(input).unsqueeze(0).to(device) _, landmarks = plfd_backbone(input) pre_landmark = landmarks[0] pre_landmark = pre_landmark.cpu().detach().numpy().reshape( -1, 2) * [size, size] for (x, y) in pre_landmark.astype(np.int32): cv2.circle(img, (x1 + x, y1 + y), 1, (0, 0, 255)) cv2.imshow('0', img) if cv2.waitKey(10) == 27: break
def main(args): checkpoint = torch.load(args.model_path, map_location=device) pfld_backbone = PFLDInference().to(device) pfld_backbone.load_state_dict(checkpoint['plfd_backbone']) pfld_backbone.eval() pfld_backbone = pfld_backbone.to(device) transform = transforms.Compose([transforms.ToTensor()]) im = Image.open(args.image_path) img = np.array(im) height, width = img.shape[:2] draw = ImageDraw.Draw(im) bounding_boxes, landmarks = detect_faces(img) print(bounding_boxes) for box in bounding_boxes: score = box[4] x1, y1, x2, y2 = (box[:4] + 0.5).astype(np.int32) w = x2 - x1 + 1 h = y2 - y1 + 1 size = int(max([w, h]) * 1.1) cx = x1 + w // 2 cy = y1 + h // 2 x1 = cx - size // 2 x2 = x1 + size y1 = cy - size // 2 y2 = y1 + size dx = max(0, -x1) dy = max(0, -y1) x1 = max(0, x1) y1 = max(0, y1) edx = max(0, x2 - width) edy = max(0, y2 - height) x2 = min(width, x2) y2 = min(height, y2) cropped = img[y1:y2, x1:x2] if (dx > 0 or dy > 0 or edx > 0 or edy > 0): cropped = cv2.copyMakeBorder(cropped, dy, edy, dx, edx, cv2.BORDER_CONSTANT, 0) cropped = cv2.resize(cropped, (112, 112)) input = cv2.resize(cropped, (112, 112)) input = cv2.cvtColor(input, cv2.COLOR_BGR2RGB) input = transform(input).unsqueeze(0).to(device) landmarks = pfld_backbone(input) pre_landmark = landmarks[0] pre_landmark = pre_landmark.cpu().detach().numpy().reshape( -1, 2) * [size, size] print(pre_landmark) for (x, y) in pre_landmark.astype(np.int32): # cv2.circle(img, (x1 + x, y1 + y), 1, (0, 0, 255)) draw.ellipse((x1 + x - 1, y1 + y - 1, x1 + x + 1, y1 + y + 1), fill=(0, 0, 255)) im.show()
def __init__(self, model_path): checkpoint = torch.load(model_path) plfd_backbone = PFLDInference().cuda() plfd_backbone.load_state_dict(checkpoint['plfd_backbone']) plfd_backbone.eval() self.plfd_backbone = plfd_backbone.cuda() self.transform = transforms.Compose([transforms.ToTensor()])
def extract_keypoints(img_path, model_path): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") checkpoint = torch.load(model_path, map_location=device) pfld_backbone = PFLDInference().to(device) pfld_backbone.load_state_dict(checkpoint['pfld_backbone']) pfld_backbone.eval() pfld_backbone = pfld_backbone.to(device) transform = torchvision.transforms.Compose( [torchvision.transforms.ToTensor()]) img = cv2.imread(img_path, 1) height, width = img.shape[:2] bounding_boxes, landmarks = detect_faces(img) for box in bounding_boxes: x1, y1, x2, y2 = (box[:4] + 0.5).astype(np.int32) w = x2 - x1 + 1 # 宽度 h = y2 - y1 + 1 # 高度 cx = x1 + w // 2 # 中心宽度 cy = y1 + h // 2 # 中心高度 size = int(max([w, h]) * 1.1) x1 = cx - size // 2 x2 = x1 + size y1 = cy - size // 2 y2 = y1 + size x1 = max(0, x1) y1 = max(0, y1) x2 = min(width, x2) y2 = min(height, y2) edx1 = max(0, -x1) edy1 = max(0, -y1) edx2 = max(0, x2 - width) edy2 = max(0, y2 - height) cropped = img[y1:y2, x1:x2] if edx1 > 0 or edy1 > 0 or edx2 > 0 or edy2 > 0: cropped = cv2.copyMakeBorder(cropped, edy1, edy2, edx1, edx2, cv2.BORDER_CONSTANT, 0) input = cv2.resize(cropped, (112, 112)) cv2.imwrite(img_path.replace(".png", "_1.png"), input) input = transform(input).unsqueeze(0).to(device) _, landmarks = pfld_backbone(input) pre_landmark = landmarks[0] key_points = pre_landmark.cpu().detach().numpy().reshape( -1, 2) * [size, size] - [edx1, edy1] keypoints = [] for (x, y) in key_points: cv2.circle(img, (x1 + int(x), y1 + int(y)), 2, (0, 255, 0), -1) keypoints.append([x1 + int(x), y1 + int(y)]) cv2.imshow('face_landmark_68', img) cv2.waitKey(1000) return keypoints
def __init__(self): self.device = torch.device( "cuda" if torch.cuda.is_available() else "cpu") model_path = os.path.join(PTH_DIR, 'checkpoint.pth.tar') checkpoint = torch.load(model_path, map_location=self.device) plfd_backbone = PFLDInference().to(self.device) plfd_backbone.load_state_dict(checkpoint['plfd_backbone']) plfd_backbone.eval() self.plfd_backbone = plfd_backbone.to(self.device) self.transform = transforms.Compose([transforms.ToTensor()])
def main(args): checkpoint = torch.load(args.model_path, map_location=device) if (args.backbone == 'VoVNet'): pfld_backbone = vovnet_pfld( num_landmarks=args.num_landmarks).to(device) elif (args.backbone == 'MobileNet'): pfld_backbone = PFLDInference( num_landmarks=args.num_landmarks).to(device) else: print(args.backbone) pfld_backbone.load_state_dict(checkpoint['plfd_backbone']) transform = transforms.Compose([transforms.ToTensor()]) if (args.num_landmarks == 68): val_dataset = A300WDatasets(args.test_dataset, transform, False) elif (args.num_landmarks == 98): val_dataset = WFLWDatasets(args.test_dataset, transform, False) val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=0) validate(val_dataloader, pfld_backbone)
def main(args): checkpoint = torch.load(args.model_path) plfd_backbone = PFLDInference().cuda() plfd_backbone.load_state_dict(checkpoint['plfd_backbone']) plfd_backbone.eval() plfd_backbone = plfd_backbone.cuda() transform = transforms.Compose([transforms.ToTensor()]) cap = cv2.VideoCapture(0) while True: ret, img = cap.read() if not ret: break bounding_boxes, landmarks = detect_faces(img) maxval_index = 0 if (len(bounding_boxes) > 0): maxval_index = np.argmax(bounding_boxes[:, 4]) if (len(bounding_boxes) > maxval_index - 1 and len(bounding_boxes) > 0): face = img[int(bounding_boxes[maxval_index][1] ):int(bounding_boxes[maxval_index][3]), int(bounding_boxes[maxval_index][0] ):int(bounding_boxes[maxval_index][2])] if face.size == 0: continue face_resized = cv2.resize(face, dsize=(112, 112), interpolation=cv2.INTER_LINEAR) face_resized = cv2.cvtColor(face_resized, cv2.COLOR_BGR2RGB) face_resized = transform(face_resized).unsqueeze(0).cuda() _, landmarks = plfd_backbone(face_resized) pre_landmark = landmarks.cpu().detach().numpy()[0].reshape( -1, 2) * [ int(bounding_boxes[maxval_index][2]) - int(bounding_boxes[maxval_index][0]), int(bounding_boxes[maxval_index][3]) - int(bounding_boxes[maxval_index][1]) ] for (x, y) in pre_landmark.astype(np.int32): cv2.circle(face, (x, y), 1, (255, 0, 0), -1) cv2.imshow("xxxx", face) if cv2.waitKey(10) == 27: break
import torch from models.pfld_vovnet import vovnet_pfld from models.pfld import PFLDInference from pthflops import count_ops device = 'cuda:0' model = PFLDInference().to(device) inp = torch.rand(1, 3, 112, 112).to(device) count_ops(model, inp)
def main(args): checkpoint = torch.load(args.model_path, map_location=device) pfld_backbone = PFLDInference().to(device) pfld_backbone.load_state_dict(checkpoint['plfd_backbone']) pfld_backbone.eval() pfld_backbone = pfld_backbone.to(device) transform = transforms.Compose([transforms.ToTensor()]) if not os.path.exists(args.video_path): print('Video not found.') exit() res_dir_path = os.path.splitext(os.path.basename( args.video_path))[0] if args.res_dir is None else args.res_dir os.makedirs(res_dir_path, exist_ok=True) cap = cv2.VideoCapture(args.video_path) frame_index = 0 start_time = time.time() while (cap.isOpened()): ret, frame = cap.read() if not ret: break frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) im = Image.fromarray(frame) img = np.array(im) height, width = img.shape[:2] draw = ImageDraw.Draw(im) bounding_boxes, landmarks = detect_faces(img) # print(bounding_boxes) for box in bounding_boxes: score = box[4] x1, y1, x2, y2 = (box[:4] + 0.5).astype(np.int32) w = x2 - x1 + 1 h = y2 - y1 + 1 size = int(max([w, h]) * 1.1) cx = x1 + w // 2 cy = y1 + h // 2 x1 = cx - size // 2 x2 = x1 + size y1 = cy - size // 2 y2 = y1 + size dx = max(0, -x1) dy = max(0, -y1) x1 = max(0, x1) y1 = max(0, y1) edx = max(0, x2 - width) edy = max(0, y2 - height) x2 = min(width, x2) y2 = min(height, y2) cropped = img[y1:y2, x1:x2] if (dx > 0 or dy > 0 or edx > 0 or edy > 0): cropped = cv2.copyMakeBorder(cropped, dy, edy, dx, edx, cv2.BORDER_CONSTANT, 0) cropped = cv2.resize(cropped, (112, 112)) input = cv2.resize(cropped, (112, 112)) input = cv2.cvtColor(input, cv2.COLOR_BGR2RGB) input = transform(input).unsqueeze(0).to(device) landmarks = pfld_backbone(input) pre_landmark = landmarks[0] pre_landmark = pre_landmark.cpu().detach().numpy().reshape( -1, 2) * [size, size] # print(pre_landmark) for (x, y) in pre_landmark.astype(np.int32): # cv2.circle(img, (x1 + x, y1 + y), 1, (0, 0, 255)) draw.ellipse((x1 + x - 1, y1 + y - 1, x1 + x + 1, y1 + y + 1), fill=(0, 0, 255)) # im.show() im.save(f'{res_dir_path}{os.sep}{frame_index:05}.jpg') if (frame_index + 1) == 80: print(f'{frame_index / (time.time()-start_time)} FPS') frame_index += 1 cap.release() print(f'{frame_index / (time.time()-start_time)} FPS') print('Done.')
import os import argparse from models.pfld import PFLDInference, CustomizedGhostNet from torch.autograd import Variable import torch import onnxsim parser = argparse.ArgumentParser(description='pytorch2onnx') parser.add_argument('--torch_model', default="./checkpoint/snapshot/checkpoint.pth.tar") parser.add_argument('--onnx_model', default="./output/pfld.onnx") parser.add_argument('--onnx_model_sim', help='Output ONNX model', default="./output/pfld-sim.onnx") args = parser.parse_args() print("=====> load pytorch checkpoint...") checkpoint = torch.load(args.torch_model, map_location=torch.device('cpu')) plfd_backbone = PFLDInference() # plfd_backbone = CustomizedGhostNet(width=1, dropout=0.2) plfd_backbone.load_state_dict(checkpoint['plfd_backbone']) print("PFLD bachbone:", plfd_backbone) print("=====> convert pytorch model to onnx...") dummy_input = Variable(torch.randn(1, 3, 112, 112)) input_names = ["input_1"] output_names = [ "output_1" ] torch.onnx.export(plfd_backbone, dummy_input, args.onnx_model, verbose=True, input_names=input_names, output_names=output_names) print("====> check onnx model...") import onnx model = onnx.load(args.onnx_model) onnx.checker.check_model(model)
def main(args): # Step 1: parse args config logging.basicConfig( format= '[%(asctime)s] [p%(process)s] [%(pathname)s:%(lineno)d] [%(levelname)s] %(message)s', level=logging.INFO, handlers=[ logging.FileHandler(args.log_file, mode='w'), logging.StreamHandler() ]) print_args(args) # Step 2: model, criterion, optimizer, scheduler plfd_backbone = PFLDInference().cuda() auxiliarynet = AuxiliaryNet().cuda() criterion = PFLDLoss() optimizer = torch.optim.Adam([{ 'params': plfd_backbone.parameters() }, { 'params': auxiliarynet.parameters() }], lr=args.base_lr, weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='min', patience=args.lr_patience, verbose=True) # step 3: data # argumetion transform = transforms.Compose([transforms.ToTensor()]) wlfwdataset = WLFWDatasets(args.dataroot, transform) dataloader = DataLoader(wlfwdataset, batch_size=args.train_batchsize, shuffle=True, num_workers=args.workers, drop_last=False) wlfw_val_dataset = WLFWDatasets(args.val_dataroot, transform) wlfw_val_dataloader = DataLoader(wlfw_val_dataset, batch_size=args.val_batchsize, shuffle=False, num_workers=args.workers) # step 4: run writer = SummaryWriter(args.tensorboard) for epoch in range(args.start_epoch, args.end_epoch + 1): weighted_train_loss, train_loss = train(dataloader, plfd_backbone, auxiliarynet, criterion, optimizer, epoch) filename = os.path.join(str(args.snapshot), "checkpoint_epoch_" + str(epoch) + '.pth.tar') save_checkpoint( { 'epoch': epoch, 'plfd_backbone': plfd_backbone.state_dict(), 'auxiliarynet': auxiliarynet.state_dict() }, filename) val_loss = validate(wlfw_val_dataloader, plfd_backbone, auxiliarynet, criterion, epoch) scheduler.step(val_loss) writer.add_scalar('data/weighted_loss', weighted_train_loss, epoch) writer.add_scalars('data/loss', { 'val loss': val_loss, 'train loss': train_loss }, epoch) writer.close()
def main(args): # Step 1: parse args config logging.basicConfig( format= '[%(asctime)s] [p%(process)s] [%(pathname)s:%(lineno)d] [%(levelname)s] %(message)s', level=logging.INFO, handlers=[ logging.FileHandler(args.log_file, mode='w'), logging.StreamHandler() ]) print_args(args) if args.backbone == "v2": from models.pfld import PFLDInference, AuxiliaryNet elif args.backbone == "v3": from models.mobilev3_pfld import PFLDInference, AuxiliaryNet elif args.backbone == "ghost": from models.ghost_pfld import PFLDInference, AuxiliaryNet elif args.backbone == "lite": from models.lite import PFLDInference, AuxiliaryNet else: raise ValueError("backbone is not implemented") plfd_backbone = PFLDInference() auxiliarynet = AuxiliaryNet() if os.path.exists(args.resume) and args.resume.endswith('.pth'): logging.info("loading the checkpoint from {}".format(args.resume)) check = torch.load(args.resume, map_location=torch.device('cpu')) plfd_backbone.load_state_dict(check["plfd_backbone"]) auxiliarynet.load_state_dict(check["auxiliarynet"]) args.start_epoch = check["epoch"] # Step 2: model, criterion, optimizer, scheduler plfd_backbone = plfd_backbone.to(device) auxiliarynet = auxiliarynet.to(device) criterion = LandMarkLoss() optimizer = torch.optim.Adam([{ 'params': plfd_backbone.parameters() }, { 'params': auxiliarynet.parameters() }], lr=args.base_lr, weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='min', patience=args.lr_patience, verbose=True) # step 3: data # argumetion transform = transforms.Compose([transforms.ToTensor()]) wlfwdataset = PFLDDatasets(args.dataroot, transform, img_root=os.path.realpath('./data'), img_size=args.img_size) dataloader = DataLoader(wlfwdataset, batch_size=args.train_batchsize, shuffle=True, num_workers=args.workers, drop_last=False) wlfw_val_dataset = PFLDDatasets(args.val_dataroot, transform, img_root=os.path.realpath('./data'), img_size=args.img_size) wlfw_val_dataloader = DataLoader(wlfw_val_dataset, batch_size=args.val_batchsize, shuffle=False, num_workers=args.workers) # step 4: run weighted_losses = [] train_losses = [] val_losses = [] val_nme = 1e6 for epoch in range(args.start_epoch, args.end_epoch + 1): weighted_train_loss, train_loss = train(dataloader, plfd_backbone, auxiliarynet, criterion, optimizer, epoch) if epoch % args.epoch_interval == 0: filename = os.path.join(str(args.snapshot), "checkpoint_epoch_" + str(epoch) + '.pth') save_checkpoint( { 'epoch': epoch, 'plfd_backbone': plfd_backbone.state_dict(), 'auxiliarynet': auxiliarynet.state_dict() }, filename) val_loss, cur_val_nme = validate(wlfw_val_dataloader, plfd_backbone, auxiliarynet, criterion) if cur_val_nme < val_nme: filename = os.path.join(str(args.snapshot), "checkpoint_min_nme.pth") save_checkpoint( { 'epoch': epoch, 'plfd_backbone': plfd_backbone.state_dict(), 'auxiliarynet': auxiliarynet.state_dict() }, filename) val_nme = cur_val_nme scheduler.step(val_loss) weighted_losses.append(weighted_train_loss.item()) train_losses.append(train_loss.item()) val_losses.append(val_loss.item()) logging.info( "epoch: {}, weighted_train_loss: {:.4f}, trainset loss: {:.4f} valset loss: {:.4f} best val " "nme: {:.4f}\n ".format(epoch, weighted_train_loss, train_loss, val_loss, val_nme)) weighted_losses = " ".join(list(map(str, weighted_losses))) train_losses = " ".join(list(map(str, train_losses))) val_losses = " ".join(list(map(str, val_losses))) logging.info(weighted_losses) logging.info(train_losses) logging.info(val_losses)
landmarks.append(landmark.tolist()) return landmarks ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'} app = Flask(__name__) # cfg = cfg_re50 cfg = cfg_mnet # trained_model = './weights/Resnet50_epoch_95.pth' retina_trained_model = './weights/mobilenet0.25_epoch_245.pth' mask_trained_model = './weights/net_21.pth' pfld_trained_model = './weights/checkpoint_epoch_500.pth.tar' # net and model retina_net = RetinaFace(cfg=cfg, phase='test') mask_net = resnet50(num_classes=2) pfld_backbone = PFLDInference() # load pre-trained model retina_net.load_state_dict(torch.load(retina_trained_model)) mask_net.load_state_dict(torch.load(mask_trained_model)) pfld_backbone.load_state_dict(torch.load(pfld_trained_model)['plfd_backbone']) retina_net = retina_net.cuda(0) mask_net = mask_net.cuda(0) pfld_net = pfld_backbone.cuda(0) retina_net.eval() mask_net.eval() pfld_net.eval() resize = 1 top_k = 5000 keep_top_k = 750
def main(args): checkpoint = torch.load(args.model_path) plfd_backbone = PFLDInference().cuda() plfd_backbone.load_state_dict(checkpoint['plfd_backbone']) plfd_backbone.eval() plfd_backbone = plfd_backbone.cuda() transform = transforms.Compose([transforms.ToTensor()]) img = cv2.imread(args.image_path) height, width = img.shape[:2] bounding_boxes, landmarks = detect_faces(img) for box in bounding_boxes: #score = box[4] x1, y1, x2, y2 = (box[:4] + 0.5).astype(np.int32) w = x2 - x1 + 1 h = y2 - y1 + 1 size = int(max([w, h]) * 1.1) cx = x1 + w // 2 cy = y1 + h // 2 x1 = cx - size // 2 x2 = x1 + size y1 = cy - size // 2 y2 = y1 + size dx = max(0, -x1) dy = max(0, -y1) x1 = max(0, x1) y1 = max(0, y1) edx = max(0, x2 - width) edy = max(0, y2 - height) x2 = min(width, x2) y2 = min(height, y2) cropped = img[y1:y2, x1:x2] if (dx > 0 or dy > 0 or edx > 0 or edy > 0): cropped = cv2.copyMakeBorder(cropped, dy, edy, dx, edx, cv2.BORDER_CONSTANT, 0) cv2.imwrite("cropped.jpg", cropped) print("cropped shape:", cropped.shape) #cropped shape: (668, 668, 3) cropped = cv2.resize(cropped, (112, 112)) input = cv2.resize(cropped, (112, 112)) input = cv2.cvtColor(input, cv2.COLOR_BGR2RGB) input = transform(input).unsqueeze(0).cuda() _, landmarks = plfd_backbone(input) pre_landmark = landmarks[0] print("size:", size) #size: 668 pre_landmark = pre_landmark.cpu().detach().numpy().reshape( -1, 2) * [size, size] for idx, (x, y) in enumerate(pre_landmark.astype(np.int32)): cv2.circle(img, (x1 + x, y1 + y), 1, (0, 0, 255)) font = cv2.FONT_HERSHEY_SIMPLEX cv2.putText(img, str(idx + 1), (x1 + x, y1 + y), font, 0.4, (0, 0, 255), 1, cv2.LINE_AA) points = transfer(pre_landmark) for point in points: cv2.circle(img, (int(round(x1 + point[0])), int(round(y1 + point[1]))), 5, (255, 0, 0), 2) cv2.imwrite("result.jpg", img)
def main(args): # Step 1: parse args config logging.basicConfig(format='[%(asctime)s] [%(levelname)s] %(message)s', level=logging.INFO, handlers=[ logging.FileHandler(args.log_file, mode='w'), logging.StreamHandler() ]) print_args(args) # Step 2: model, criterion, optimizer, scheduler plfd_backbone = PFLDInference().to(device) auxiliarynet = AuxiliaryNet().to(device) # criterion = PFLDLoss() criterion = LandmarkLoss() optimizer = torch.optim.Adam([{ 'params': plfd_backbone.parameters() }, { 'params': auxiliarynet.parameters() }], lr=args.base_lr, weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='min', patience=args.lr_patience, verbose=True) if args.resume_path: print('loading checkpoint {}'.format(args.resume_path)) checkpoint = torch.load(str(args.resume_path)) args.start_epoch = checkpoint['epoch'] plfd_backbone.load_state_dict(checkpoint['plfd_backbone']) auxiliarynet.load_state_dict(checkpoint['auxiliarynet']) if 'optimizer' in checkpoint.keys(): optimizer.load_state_dict(checkpoint['optimizer']) # step 3: data # argumetion train_transform = transforms.Compose([ AugCrop(output_size=112, is_training=True), HorizontalFlip(mirror=args.mirror_file), RandomRotate(max_angle=30), Affine(max_strength=30, output_size=112), ColorDistort() ]) val_transform = transforms.Compose([AugCrop(output_size=112)]) ibugdataset = IBUGDatasets(args.train_json, transform=train_transform, is_train=True) train_dataset_size = ibugdataset.get_dataset_size() sampler = RandomSampler(ibugdataset, replacement=True, num_samples=train_dataset_size) dataloader = DataLoader(ibugdataset, batch_size=args.train_batchsize, sampler=sampler, num_workers=args.workers, drop_last=False) ibug_val_dataset = IBUGDatasets(args.val_json, transform=val_transform) val_dataset_size = ibug_val_dataset.get_dataset_size() val_sampler = RandomSampler(ibug_val_dataset, replacement=True, num_samples=val_dataset_size) ibug_val_dataloader = DataLoader(ibug_val_dataset, batch_size=args.val_batchsize, sampler=val_sampler, num_workers=args.workers) # step 4: run writer = SummaryWriter(args.tensorboard) for epoch in range(args.start_epoch, args.end_epoch + 1): weighted_train_loss, train_loss = train(dataloader, plfd_backbone, auxiliarynet, criterion, optimizer, epoch) filename = os.path.join(str(args.snapshot), "checkpoint_epoch_" + str(epoch) + '.pth.tar') save_checkpoint( { 'epoch': epoch, 'plfd_backbone': plfd_backbone.state_dict(), 'auxiliarynet': auxiliarynet.state_dict(), 'optimizer': optimizer.state_dict() }, filename) val_loss = validate(ibug_val_dataloader, plfd_backbone, auxiliarynet, criterion, epoch) scheduler.step(val_loss) writer.add_scalar('data/weighted_loss', weighted_train_loss, epoch) writer.add_scalars('data/loss', { 'val loss': val_loss, 'train loss': train_loss }, epoch) writer.close()
import onnxsim parser = argparse.ArgumentParser(description='pytorch2onnx') parser.add_argument( '--torch_model', default="./checkpoint/snapshot/checkpoint.pth.tar") parser.add_argument('--onnx_model', default="./output/pfld.onnx") parser.add_argument( '--onnx_model_sim', help='Output ONNX model', default="./output/pfld-sim.onnx") args = parser.parse_args() print("=====> load pytorch checkpoint...") checkpoint = torch.load(args.torch_model, map_location=torch.device('cpu')) plfd_backbone = PFLDInference() plfd_backbone.load_state_dict(checkpoint['plfd_backbone'], strict=False) print("PFLD bachbone:", plfd_backbone) print("=====> convert pytorch model to onnx...") dummy_input = torch.randn(1, 3, 112, 112) input_names = ["input_1"] output_names = ["output_1"] torch.onnx.export( plfd_backbone, dummy_input, args.onnx_model, verbose=True, input_names=input_names, output_names=output_names)
def main(args): # face detection model model = init_detector(args.config, args.face_model) # landmark model checkpoint = torch.load(args.mark_model, map_location=device) plfd_backbone = PFLDInference().to(device) plfd_backbone.load_state_dict(checkpoint['plfd_backbone']) plfd_backbone.eval() plfd_backbone = plfd_backbone.to(device) transform = transforms.Compose([transforms.ToTensor()]) fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v') save_path = '/home/yang/mark.mp4' writer = cv2.VideoWriter(save_path, fourcc, 30.0, (1280, 720), True) if args.video_path: cap = cv2.VideoCapture(args.video_path) else: cap = cv2.VideoCapture(0) while True: ret, img = cap.read() if not ret: break height, width = img.shape[:2] # bounding_boxes, landmarks = detect_faces(img) results = inference_detector(model, img) bboxs = decode_detections(results[0], args.d_thresh) for i, bbox in enumerate(bboxs): # x1, y1, x2, y2 = (bbox[:4]+0.5).astype(np.int32) w = bbox[2] - bbox[0] h = bbox[3] - bbox[1] add = int(max(w, h)) bimg = cv2.copyMakeBorder(img, add, add, add, add, borderType=cv2.BORDER_CONSTANT, value=[127., 127., 127.]) bbox += add face_width = (1 + 0.4) * w center = [(bbox[0] + bbox[2]) // 2, (bbox[1] + bbox[3]) // 2] bbox[0] = center[0] - face_width // 2 bbox[1] = center[1] - face_width // 2 bbox[2] = center[0] + face_width // 2 bbox[3] = center[1] + face_width // 2 bbox = bbox.astype(np.int) crop_image = bimg[bbox[1]:bbox[3], bbox[0]:bbox[2], :] height, width, _ = crop_image.shape crop_image = cv2.resize(crop_image, (112, 112)) cv2.imshow('cropped face %d ' % i, crop_image) # input = cv2.resize(cropped, (112, 112)) input = cv2.cvtColor(crop_image, cv2.COLOR_BGR2RGB) input = transform(input).unsqueeze(0).to(device) _, landmarks = plfd_backbone(input) pre_landmark = landmarks[0] pre_landmark = pre_landmark.cpu().detach().numpy().reshape( -1, 2) * [width, height] for (x, y) in pre_landmark.astype(np.int32): cv2.circle(img, (bbox[0] + x - add, bbox[1] + y - add), 1, (0, 255, 0), -1) cv2.imshow('0', img) writer.write(img) if cv2.waitKey(1) == 27: break
def main(args): # Step 1: parse args config logging.basicConfig( format= '[%(asctime)s] [p%(process)s] [%(pathname)s:%(lineno)d] [%(levelname)s] %(message)s', level=logging.INFO, handlers=[ logging.FileHandler(args.log_file, mode='w'), logging.StreamHandler() ]) print_args(args) # Step 2: model, criterion, optimizer, scheduler if wandb.config.pfld_backbone == "GhostNet": plfd_backbone = CustomizedGhostNet(width=wandb.config.ghostnet_width, dropout=0.2) logger.info(f"Using GHOSTNET with width={wandb.config.ghostnet_width} as backbone of PFLD backbone") # If using pretrained weight from ghostnet model trained on image net if (wandb.config.ghostnet_with_pretrained_weight_image_net == True): logger.info(f"Using pretrained weights of ghostnet model trained on image net data ") plfd_backbone = load_pretrained_weight_imagenet_for_ghostnet_backbone( plfd_backbone, "./checkpoint_imagenet/state_dict_93.98.pth") else: plfd_backbone = PFLDInference().to(device) # MobileNet2 defaut logger.info("Using MobileNet2 as backbone of PFLD backbone") auxiliarynet = AuxiliaryNet().to(device) # Watch model by wandb wandb.watch(plfd_backbone) wandb.watch(auxiliarynet) criterion = PFLDLoss() optimizer = torch.optim.Adam( [{ 'params': plfd_backbone.parameters() }, { 'params': auxiliarynet.parameters() }], lr=args.base_lr, weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=args.lr_patience, verbose=True) # step 3: data # argumetion transform = transforms.Compose([transforms.ToTensor()]) wlfwdataset = WLFWDatasets(args.dataroot, transform) dataloader = DataLoader( wlfwdataset, batch_size=args.train_batchsize, shuffle=True, num_workers=args.workers, drop_last=False) wlfw_val_dataset = WLFWDatasets(args.val_dataroot, transform) wlfw_val_dataloader = DataLoader( wlfw_val_dataset, batch_size=args.val_batchsize, shuffle=False, num_workers=args.workers) # step 4: run writer = SummaryWriter(args.tensorboard) for epoch in range(args.start_epoch, args.end_epoch + 1): weighted_train_loss, train_loss = train(dataloader, plfd_backbone, auxiliarynet, criterion, optimizer, epoch) filename = os.path.join( str(args.snapshot), "checkpoint_epoch_" + str(epoch) + '.pth.tar') save_checkpoint({ 'epoch': epoch, 'plfd_backbone': plfd_backbone.state_dict(), 'auxiliarynet': auxiliarynet.state_dict() }, filename) val_loss = validate(wlfw_val_dataloader, plfd_backbone, auxiliarynet, criterion) wandb.log({"metric/val_loss": val_loss}) scheduler.step(val_loss) writer.add_scalar('data/weighted_loss', weighted_train_loss, epoch) writer.add_scalars('data/loss', {'val loss': val_loss, 'train loss': train_loss}, epoch) writer.close()