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()
