def main():
args = parse_args()
logger, final_output_dir, tb_log_dir = \
utils.create_logger(config, args.cfg, 'train')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.determinstic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
model = models.get_face_alignment_net(config)
# copy model files
writer_dict = {
'writer': SummaryWriter(log_dir=tb_log_dir),
'train_global_steps': 0,
'valid_global_steps': 0,
}
gpus = list(config.GPUS)
model = nn.DataParallel(model, device_ids=gpus).cuda()
# loss
criterion = torch.nn.MSELoss(size_average=True).cuda()
optimizer = utils.get_optimizer(config, model)
best_nme = 100
last_epoch = config.TRAIN.BEGIN_EPOCH
if config.TRAIN.RESUME:
model_state_file = os.path.join(final_output_dir, 'latest.pth')
if os.path.islink(model_state_file):
checkpoint = torch.load(model_state_file)
last_epoch = checkpoint['epoch']
best_nme = checkpoint['best_nme']
model.module.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
else:
print("=> no checkpoint found")
if isinstance(config.TRAIN.LR_STEP, list):
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR,
last_epoch - 1)
else:
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
config.TRAIN.LR_STEP,
config.TRAIN.LR_FACTOR,
last_epoch - 1)
dataset_type = get_dataset(config)
train_loader = DataLoader(dataset=dataset_type(config, is_train=True),
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU *
len(gpus),
shuffle=config.TRAIN.SHUFFLE,
num_workers=config.WORKERS,
pin_memory=config.PIN_MEMORY)
val_loader = DataLoader(dataset=dataset_type(config, is_train=False),
batch_size=config.TEST.BATCH_SIZE_PER_GPU *
len(gpus),
shuffle=False,
num_workers=config.WORKERS,
pin_memory=config.PIN_MEMORY)
for epoch in range(last_epoch, config.TRAIN.END_EPOCH):
lr_scheduler.step()
function.train(config, train_loader, model, criterion, optimizer,
epoch, writer_dict)
# evaluate
nme, predictions = function.validate(config, val_loader, model,
criterion, epoch, writer_dict)
is_best = nme < best_nme
best_nme = min(nme, best_nme)
logger.info('=> saving checkpoint to {}'.format(final_output_dir))
print("best:", is_best)
utils.save_checkpoint(
{
"state_dict": model.module.state_dict(),
"epoch": epoch + 1,
"best_nme": best_nme,
"optimizer": optimizer.state_dict(),
}, predictions, is_best, final_output_dir,
'checkpoint_{}.pth'.format(epoch))
final_model_state_file = os.path.join(final_output_dir, 'final_state.pth')
logger.info(
'saving final model state to {}'.format(final_model_state_file))
torch.save(model.module.state_dict(), final_model_state_file)
writer_dict['writer'].close()
def main():
args = parse_args()
logger, final_output_dir, tb_log_dir = \
utils.create_logger(config, args.cfg, 'train')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn.benchmark = config.CUDNN.BENCHMARK
# cudnn.determinstic = config.CUDNN.DETERMINISTIC
# cudnn.enabled = config.CUDNN.ENABLED
# if isinstance(config.TRAIN.LR_STEP, list):
# lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
# optimizer, config.TRAIN.LR_STEP,
# # config.TRAIN.LR_FACTOR, last_epoch-1
# config.TRAIN.LR_FACTOR, 0
# )
# else:
# lr_scheduler = torch.optim.lr_scheduler.StepLR(
# optimizer, config.TRAIN.LR_STEP,
# # config.TRAIN.LR_FACTOR, last_epoch-1
# config.TRAIN.LR_FACTOR, 0
# )
dataset_type = get_dataset(config)
train_dataset = dataset_type(config,
is_train=True)
# train_dataset[0]
# return 0
train_loader = DataLoader(
dataset=dataset_type(config,
is_train=True),
# batch_size=config.TRAIN.BATCH_SIZE_PER_GPU*len(gpus),
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU,
shuffle=config.TRAIN.SHUFFLE,
num_workers=config.WORKERS)
# val_loader = DataLoader(
# dataset=dataset_type(config,
# is_train=True),
# # batch_size=config.TEST.BATCH_SIZE_PER_GPU*len(gpus),
# batch_size=config.TEST.BATCH_SIZE_PER_GPU,
# shuffle=False,
# num_workers=config.WORKERS,
# # pin_memory=config.PIN_MEMORY
# )
model = models.get_face_alignment_net(config)
# copy model files
writer_dict = {
'writer': SummaryWriter(log_dir=tb_log_dir),
'train_global_steps': 0,
'valid_global_steps': 0,
}
gpus = list(config.GPUS)
# model = nn.DataParallel(model, device_ids=gpus).cuda()
model.to("cuda")
# loss
criterion = torch.nn.MSELoss(size_average=True).cuda()
# criterion = fnn.mse_loss
# criterion = WingLoss()
# criterion = Loss_weighted()
optimizer = utils.get_optimizer(config, model)
best_nme = 100
last_epoch = config.TRAIN.BEGIN_EPOCH
if config.TRAIN.RESUME:
model_state_file = os.path.join(final_output_dir,
'latest.pth')
if os.path.isfile(model_state_file):
with open(model_state_file, "rb") as fp:
state_dict = torch.load(fp)
model.load_state_dict(state_dict)
last_epoch = 1
# checkpoint = torch.load(model_state_file)
# last_epoch = checkpoint['epoch']
# best_nme = checkpoint['best_nme']
# model.load_state_dict(checkpoint['state_dict'])
# optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint (epoch {})"
.format(last_epoch))
else:
print("=> no checkpoint found")
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
for epoch in range(last_epoch, config.TRAIN.END_EPOCH):
lr_scheduler.step()
function.train(config, train_loader, model, criterion,
optimizer, epoch, writer_dict)
# evaluate
nme = 0
# nme, predictions = function.validate(config, val_loader, model,
# criterion, epoch, writer_dict)
is_best = True
# is_best = nme < best_nme
best_nme = min(nme, best_nme)
logger.info('=> saving checkpoint to {}'.format(final_output_dir))
print("best:", is_best)
torch.save(model.state_dict(), os.path.join(final_output_dir, 'mse_relu_checkpoint_{}.pth'.format(epoch)))
# utils.save_checkpoint(
# {"state_dict": model,
# "epoch": epoch + 1,
# "best_nme": best_nme,
# "optimizer": optimizer.state_dict(),
# }, predictions, is_best, final_output_dir, 'checkpoint_{}.pth'.format(epoch))
final_model_state_file = os.path.join(final_output_dir,
'final_state.pth')
logger.info('saving final model state to {}'.format(
final_model_state_file))
torch.save(model.state_dict(), final_model_state_file)
writer_dict['writer'].close()
def main():
#
args = parse_args()
#
logger, final_output_dir, tb_log_dir = \
utils.create_logger(config, args.cfg, 'test')
#
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
#
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.determinstic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
#
config.defrost()
config.MODEL.INIT_WEIGHTS = False
config.freeze()
model = models.get_face_alignment_net(config)
#
gpus = list(config.GPUS)
#
# # load model
state_dict = torch.load(args.model_file)
if 'state_dict' in state_dict.keys():
state_dict = state_dict['state_dict']
model.load_state_dict(state_dict)
else:
model.module.load_state_dict(state_dict)
model = nn.DataParallel(model, device_ids=gpus).cuda()
#
dataset_type = get_dataset(config)
test_loader = DataLoader(dataset=dataset_type(config, is_train=False),
batch_size=config.TEST.BATCH_SIZE_PER_GPU *
len(gpus),
shuffle=False,
num_workers=config.WORKERS,
pin_memory=config.PIN_MEMORY)
nme, predictions = function.inference(config, test_loader, model)
torch.save(predictions, os.path.join(final_output_dir, 'predictions.pth'))
target = test_loader.dataset.load_all_pts()
pred = 16 * predictions
l = len(pred)
res = 0.0
res_tmp = [0.0 for i in range(config.MODEL.NUM_JOINTS)]
res_tmp = np.array(res_tmp)
res_temp_x = target - pred
res_temp_x = res_temp_x[:, :, 0]
res_temp_y = target - pred
res_temp_y = res_temp_y[:, :, 1]
# csv_file_test_x = pd.DataFrame(np.transpose(np.array(pred[:, :, 0])), columns=test_loader.dataset.annotation_files)
# csv_file_test_y = pd.DataFrame(np.transpose(np.array(pred[:, :, 1])), columns=test_loader.dataset.annotation_files)
# csv_file_target_x = pd.DataFrame(np.transpose(np.array(target[:, :, 0])), columns=test_loader.dataset.annotation_files)
# csv_file_target_y = pd.DataFrame(np.transpose(np.array(target[:, :, 1])), columns=test_loader.dataset.annotation_files)
for i in range(l):
trans = np.sqrt(
pow(target[i][0][0] - target[i][1][0], 2) +
pow(target[i][0][1] - target[i][1][1], 2)) / 30.0
res_temp_x[i] = res_temp_x[i] / trans
res_temp_y[i] = res_temp_y[i] / trans
for j in range(len(target[i])):
dist = np.sqrt(
np.power((target[i][j][0] - pred[i][j][0]), 2) +
np.power((target[i][j][1] - pred[i][j][1]), 2)) / trans
res += dist
res_tmp[j] += dist
res_t = np.sqrt(res_temp_x * res_temp_x + res_temp_y * res_temp_y)
# pd.DataFrame(data=res_temp_x.data.value).to_csv('res_x')
# pd.DataFrame(data=res_temp_y.data.value).to_csv('res_y')
# pd.DataFrame(data=res_t.data.value).to_csv('res_t')
res_tmp /= np.float(len(pred))
print(res_tmp)
print(np.mean(res_tmp))
res /= (len(pred) * len(pred[0]))
print(res)
def main():
# parsing script arguments
args = parse_args()
device = torch.device(args.device)
# initialize logger
logger, final_output_dir, tb_log_dir = utils.create_logger(config, args.cfg, "demo")
# log arguments and config values
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# init landmark model
model = models.get_face_alignment_net(config)
# get input size from the config
input_size = config.MODEL.IMAGE_SIZE
# load model
state_dict = torch.load(args.landmark_model, map_location=device)
# remove `module.` prefix from the pre-trained weights
new_state_dict = OrderedDict()
for key, value in state_dict.items():
name = key[7:]
new_state_dict[name] = value
# load weights without the prefix
model.load_state_dict(new_state_dict)
# run model on device
model = model.to(device)
# init mean and std values for the landmark model's input
mean = config.MODEL.MEAN
mean = np.array(mean, dtype=np.float32)
std = config.MODEL.STD
std = np.array(std, dtype=np.float32)
# defining prototxt and caffemodel paths
detector_model = args.detector_model
detector_weights = args.detector_weights
# load model
detector = cv2.dnn.readNetFromCaffe(detector_model, detector_weights)
capture = cv2.VideoCapture(0)
frame_num = 0
while True:
# capture frame-by-frame
success, frame = capture.read()
# break if no frame
if not success:
break
frame_num += 1
print("frame_num: ", frame_num)
landmarks_img = frame.copy()
result = frame.copy()
result = result.astype(np.float32) / 255.0
# get frame's height and width
height, width = frame.shape[:2] # 640x480
# resize and subtract BGR mean values, since Caffe uses BGR images for input
blob = cv2.dnn.blobFromImage(
frame, scalefactor=1.0, size=(300, 300), mean=(104.0, 177.0, 123.0),
)
# passing blob through the network to detect faces
detector.setInput(blob)
# detector output format:
# [image_id, class, confidence, left, bottom, right, top]
face_detections = detector.forward()
# loop over the detections
for i in range(0, face_detections.shape[2]):
# extract confidence
confidence = face_detections[0, 0, i, 2]
# filter detections by confidence greater than the minimum threshold
if confidence > 0.5:
# get coordinates of the bounding box
box = face_detections[0, 0, i, 3:7] * np.array(
[width, height, width, height],
)
(x1, y1, x2, y2) = box.astype("int")
# show original image
cv2.imshow("original image", frame)
# crop to detection and resize
resized = crop(
frame,
torch.Tensor([x1 + (x2 - x1) / 2, y1 + (y2 - y1) / 2]),
1.5,
tuple(input_size),
)
# convert from BGR to RGB since HRNet expects RGB format
resized = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
img = resized.astype(np.float32) / 255.0
# normalize landmark net input
normalized_img = (img - mean) / std
# predict face landmarks
model = model.eval()
with torch.no_grad():
input = torch.Tensor(normalized_img.transpose([2, 0, 1]))
input = input.to(device)
output = model(input.unsqueeze(0))
score_map = output.data.cpu()
preds = decode_preds(
score_map,
[torch.Tensor([x1 + (x2 - x1) / 2, y1 + (y2 - y1) / 2])],
[1.5],
score_map.shape[2:4],
)
preds = preds.squeeze(0)
landmarks = preds.data.cpu().detach().numpy()
# draw landmarks
for k, landmark in enumerate(landmarks, 1):
landmarks_img = cv2.circle(
landmarks_img,
center=(landmark[0], landmark[1]),
radius=3,
color=(0, 0, 255),
thickness=-1,
)
# draw landmarks' labels
landmarks_img = cv2.putText(
img=landmarks_img,
text=str(k),
org=(int(landmark[0]) + 5, int(landmark[1]) + 5),
fontFace=cv2.FONT_HERSHEY_SIMPLEX,
fontScale=0.5,
color=(0, 0, 255),
)
# show results by drawing predicted landmarks and their labels
cv2.imshow("image with landmarks", landmarks_img)
# get chosen landmarks 2-16, 30 as destination points
# note that landmarks numbering starts from 0
dst_pts = np.array(
[
landmarks[1],
landmarks[2],
landmarks[3],
landmarks[4],
landmarks[5],
landmarks[6],
landmarks[7],
landmarks[8],
landmarks[9],
landmarks[10],
landmarks[11],
landmarks[12],
landmarks[13],
landmarks[14],
landmarks[15],
landmarks[29],
],
dtype="float32",
)
# load mask annotations from csv file to source points
mask_annotation = os.path.splitext(os.path.basename(args.mask_image))[0]
mask_annotation = os.path.join(
os.path.dirname(args.mask_image), mask_annotation + ".csv",
)
with open(mask_annotation) as csv_file:
csv_reader = csv.reader(csv_file, delimiter=",")
src_pts = []
for i, row in enumerate(csv_reader):
# skip head or empty line if it's there
try:
src_pts.append(np.array([float(row[1]), float(row[2])]))
except ValueError:
continue
src_pts = np.array(src_pts, dtype="float32")
# overlay with a mask only if all landmarks have positive coordinates:
if (landmarks > 0).all():
# load mask image
mask_img = cv2.imread(args.mask_image, cv2.IMREAD_UNCHANGED)
mask_img = mask_img.astype(np.float32)
mask_img = mask_img / 255.0
# get the perspective transformation matrix
M, _ = cv2.findHomography(src_pts, dst_pts)
# transformed masked image
transformed_mask = cv2.warpPerspective(
mask_img,
M,
(result.shape[1], result.shape[0]),
None,
cv2.INTER_LINEAR,
cv2.BORDER_CONSTANT,
)
# mask overlay
alpha_mask = transformed_mask[:, :, 3]
alpha_image = 1.0 - alpha_mask
for c in range(0, 3):
result[:, :, c] = (
alpha_mask * transformed_mask[:, :, c]
+ alpha_image * result[:, :, c]
)
# display the resulting frame
cv2.imshow("image with mask overlay", result)
# waiting for the escape button to exit
k = cv2.waitKey(1)
if k == 27:
break
# when everything done, release the capture
capture.release()
cv2.destroyAllWindows()
def main():
args = parse_args()
# set logger and dir
logger, final_output_dir, tb_log_dir = \
utils.create_logger(config, args.experiment_name, 'train')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# set cudnn
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.determinstic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
# 目前仅支持单gpu，todo：增加多gpu支持
# set model and loss and criterion
model = models.get_face_alignment_net(config)
model = model.cuda(config.GPUS[0])
criterion = torch.nn.MSELoss(size_average=True).cuda(config.GPUS[0])
# criterion = AdaptiveWingLoss()
optimizer = utils.get_optimizer(config, model)
# get dataset
dataset_type = get_dataset(config)
# get dataloader
train_loader = DataLoader(dataset=dataset_type(config, is_train=True),
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU,
shuffle=config.TRAIN.SHUFFLE,
num_workers=config.WORKERS,
pin_memory=config.PIN_MEMORY)
val_loader = DataLoader(dataset=dataset_type(config, is_train=False),
batch_size=config.TEST.BATCH_SIZE_PER_GPU,
shuffle=False,
num_workers=config.WORKERS,
pin_memory=config.PIN_MEMORY)
# set lr_scheduler
last_epoch = config.TRAIN.BEGIN_EPOCH
if isinstance(config.TRAIN.LR_STEP, list):
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR,
last_epoch - 1)
else:
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
config.TRAIN.LR_STEP,
config.TRAIN.LR_FACTOR,
last_epoch - 1)
# set training writer
writer_dict = {
'writer': SummaryWriter(log_dir=tb_log_dir),
'train_global_steps': 0,
'valid_global_steps': 0,
}
# set training resume function
if config.TRAIN.RESUME:
model_state_file = os.path.join(final_output_dir, 'latest.pth')
if os.path.islink(model_state_file):
checkpoint = torch.load(model_state_file)
last_epoch = checkpoint['epoch']
best_nme = checkpoint['best_nme']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
else:
print("=> no checkpoint found")
# starting training
best_nme = 10000
for epoch in range(last_epoch, config.TRAIN.END_EPOCH):
lr_scheduler.step()
# traing
function.train(config, train_loader, model, criterion, optimizer,
epoch, writer_dict)
# evaluating
nme, predictions = function.validate(config, val_loader, model,
criterion, epoch, writer_dict)
# saving
is_best = nme < best_nme
best_nme = min(nme, best_nme)
logger.info('=> saving checkpoint to {}'.format(final_output_dir))
print("best:", is_best)
utils.save_checkpoint(
{
"state_dict": model,
"epoch": epoch + 1,
"best_nme": best_nme,
"optimizer": optimizer.state_dict(),
}, predictions, is_best, final_output_dir,
'checkpoint_{}.pth'.format(epoch))
final_model_state_file = os.path.join(final_output_dir, 'final_state.pth')
logger.info(
'saving final model state to {}'.format(final_model_state_file))
torch.save(model.state_dict(), final_model_state_file)
writer_dict['writer'].close()
def main():
args = parse_args()
logger, final_output_dir, tb_log_dir = \
utils.create_logger(config, args.cfg, 'test')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn.benchmark = config.CUDNN.BENCHMARK
# cudnn.determinstic = config.CUDNN.DETERMINISTIC
# cudnn.enabled = config.CUDNN.ENABLED
config.defrost()
config.MODEL.INIT_WEIGHTS = False
config.freeze()
model = models.get_face_alignment_net(config)
gpus = list(config.GPUS)
# model = nn.DataParallel(model, device_ids=gpus).cuda()
model.to("cuda")
# print(model)
# load model
# state_dict = torch.load(args.model_file)
# print(state_dict)
# model = torch.load(args.model_file)
with open(args.model_file, "rb") as fp:
state_dict = torch.load(fp)
model.load_state_dict(state_dict)
# model.load_state_dict(state_dict['state_dict'])
# if 'state_dict' in state_dict.keys():
# state_dict = state_dict['state_dict']
# # print(state_dict)
# model.load_state_dict(state_dict)
# else:
# model.module.load_state_dict(state_dict)
dataset_type = get_dataset(config)
test_loader = DataLoader(
dataset=dataset_type(config,
is_train=False),
batch_size=config.TEST.BATCH_SIZE_PER_GPU*len(gpus),
shuffle=False,
num_workers=config.WORKERS,
pin_memory=config.PIN_MEMORY
)
predictions = function.inference(config, test_loader, model)
# print("len(predictions)", len(predictions))
# print(predictions[0])
df_predictions = []
for pred in predictions:
row = dict()
row['file_name'] = pred[0]
for id_point in range(194):
row[f'Point_M{id_point}_X'] = int(pred[1][id_point])
row[f'Point_M{id_point}_Y'] = int(pred[2][id_point])
df_predictions.append(row)
df_predictions = pd.DataFrame(df_predictions)
# print(predictions_meta[0])
df_predictions.to_csv('pred_test.csv', index=False)
def main():
args = parse_args()
logger, final_output_dir, tb_log_dir = \
utils.create_logger(config, args.cfg, 'train')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.determinstic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
gpus = list(config.GPUS)
dataset_type = get_dataset(config)
train_data = dataset_type(config, split="train")
train_loader = DataLoader(dataset=train_data,
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU *
len(gpus),
shuffle=config.TRAIN.SHUFFLE,
num_workers=config.WORKERS,
pin_memory=config.PIN_MEMORY)
val_data = dataset_type(config, split="valid")
val_loader = DataLoader(dataset=val_data,
batch_size=config.TEST.BATCH_SIZE_PER_GPU *
len(gpus),
shuffle=False,
num_workers=config.WORKERS,
pin_memory=config.PIN_MEMORY)
# config.MODEL.NUM_JOINTS = train_data.get_num_points()
model = models.get_face_alignment_net(config)
# copy model files
writer_dict = {
'writer': SummaryWriter(log_dir=tb_log_dir),
'train_global_steps': 0,
'valid_global_steps': 0,
}
model = nn.DataParallel(model, device_ids=gpus).cuda()
# loss
criterion = torch.nn.MSELoss(size_average=True).cuda()
optimizer = utils.get_optimizer(config, model)
best_nme = 100
last_epoch = config.TRAIN.BEGIN_EPOCH
if isinstance(config.TRAIN.LR_STEP, list):
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR,
last_epoch - 1)
else:
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
config.TRAIN.LR_STEP,
config.TRAIN.LR_FACTOR,
last_epoch - 1)
if config.TRAIN.RESUME:
model_state_file = os.path.join(final_output_dir, 'final.pth')
if os.path.islink(model_state_file):
checkpoint = torch.load(model_state_file)
last_epoch = checkpoint['epoch']
best_nme = checkpoint['best_nme']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
else:
print("=> no checkpoint found")
loss = []
for epoch in range(last_epoch, config.TRAIN.END_EPOCH):
losses, diff = function.train(config, train_loader, model, criterion,
optimizer, epoch, writer_dict)
loss.append(losses)
lr_scheduler.step()
np.save(
os.path.join(final_output_dir, "train_diff@epoch{}".format(epoch)),
diff)
# evaluate
nme, predictions, diff = function.validate(config, val_loader, model,
criterion, epoch,
writer_dict)
np.save(
os.path.join(final_output_dir, "valid_diff@epoch{}".format(epoch)),
diff)
is_best = nme < best_nme
best_nme = min(nme, best_nme)
logger.info('=> saving checkpoint to {}'.format(final_output_dir))
print("best:", is_best)
utils.save_checkpoint(
{
"state_dict": model,
"epoch": epoch + 1,
"best_nme": best_nme,
"optimizer": optimizer.state_dict(),
}, predictions, is_best, final_output_dir,
'checkpoint_{}.pth'.format(epoch))
if is_best:
for i in range(len(predictions)):
afile = val_data.annotation_files[i]
new_afile = '{}.{}.txt'.format(
afile,
os.path.basename(args.cfg).split('.')[0])
with open(new_afile, 'wt') as f:
pts = predictions[i].cpu().numpy()
for j in range(len(pts)):
f.write("{},{}\n".format(
pts[j][1] / val_data.factor[1],
pts[j][0] / val_data.factor[0]))
pd.DataFrame(data=loss).to_csv('loss2.csv')
final_model_state_file = os.path.join(final_output_dir, 'final_state.pth')
logger.info(
'saving final model state to {}'.format(final_model_state_file))
torch.save(model.module.state_dict(), final_model_state_file)
writer_dict['writer'].close()
def main_wflw():
args = parse_args()
logger, final_output_dir, tb_log_dir = \
utils.create_logger(config, args.cfg, 'test')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.determinstic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
config.defrost()
config.MODEL.INIT_WEIGHTS = False
config.freeze()
model = models.get_face_alignment_net(config)
dataset_type = get_dataset(config)
dataset = dataset_type(config, is_train=True)
fp = open('data/wflw/face_landmarks70_wflw_train.csv', 'w')
for i in range(len(dataset)):
img, image_path, meta = dataset[i]
fold, name = image_path.split('/')[-2], image_path.split('/')[-1]
folder = osp.join('data/wflw/xximages', fold)
if not osp.exists(folder):
os.makedirs(folder)
fname = osp.join(folder, name)
scale = meta['scale']
center = meta['center']
tpts = meta['tpts']
selpts = []
for j in range(0, 33, 2):
selpts.append(tpts[j])
# eyebow
selpts.append(tpts[33])
selpts.append((tpts[34] + tpts[41]) / 2)
selpts.append((tpts[35] + tpts[40]) / 2)
selpts.append((tpts[36] + tpts[39]) / 2)
selpts.append((tpts[37] + tpts[38]) / 2)
selpts.append((tpts[42] + tpts[50]) / 2)
selpts.append((tpts[43] + tpts[49]) / 2)
selpts.append((tpts[44] + tpts[48]) / 2)
selpts.append((tpts[45] + tpts[47]) / 2)
selpts.append(tpts[46])
# nose
for j in range(51, 60):
selpts.append(tpts[j])
# eye
selpts.append(tpts[60])
selpts.append((tpts[61] + tpts[62]) / 2)
selpts.append(tpts[63])
selpts.append(tpts[64])
selpts.append(tpts[65])
selpts.append((tpts[66] + tpts[67]) / 2)
selpts.append(tpts[68])
selpts.append(tpts[69])
selpts.append((tpts[70] + tpts[71]) / 2)
selpts.append(tpts[72])
selpts.append((tpts[73] + tpts[74]) / 2)
selpts.append(tpts[75])
for j in range(76, 98):
selpts.append(tpts[j])
fp.write('%s,%.2f,%.1f,%.1f' %
(osp.join(fold, name), scale, center[0], center[1]))
for spt in selpts:
cv2.circle(img, (spt[0], spt[1]), 1, (0, 0, 255))
fp.write(',%f,%f' % (spt[0], spt[1]))
fp.write('\n')
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
cv2.imwrite(fname, img)
fp.close()
