def main_wflwe70():
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)
for i in range(len(dataset)):
# ipdb.set_trace()
img, fname, meta = dataset[i]
filename = osp.join('data/wflwe70/xximages', fname)
if not osp.exists(osp.dirname(filename)):
os.makedirs(osp.dirname(filename))
scale = meta['scale']
center = meta['center']
tpts = meta['tpts']
for spt in tpts:
img = cv2.circle(img, (4 * spt[0], 4 * spt[1]),
1 + center[0] // 400, (255, 0, 0))
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
cv2.imwrite(filename, img)
def main_cofw():
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)
ipdb.set_trace()
dataset_type = get_dataset(config)
dataset = dataset_type(config, is_train=True)
fp = open('data/cofw/test.csv', 'w')
for i in range(len(dataset)):
# ipdb.set_trace()
img, image_path, meta = dataset[i]
fname = osp.join('data/cofw/test', osp.basename(image_path))
fp.write('%s,1,128,128' % fname)
tpts = meta['tpts']
for j in range(tpts.shape[0]):
fp.write(',%d,%d' % (tpts[j, 0], tpts[j, 1]))
fp.write('\n')
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
cv2.imwrite(fname, img)
fp.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()
# 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)
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)
with open('../data/crabs/crabs_data_test.csv', 'r') as f:
data = np.loadtxt(f, str, delimiter=",", skiprows=1)
paths = data[:, 0]
for index, path in enumerate(paths):
img = cv2.imread("../data/crabs/images/{}".format(path))
a = predictions[index]
b = a.numpy()
for index, px in enumerate(b):
# print(tuple(px))
cv2.circle(img, tuple(px), 1, (0, 0, 255), 3, 8, 0)
# cv2.imwrite("/home/njtech/Jiannan/crabs/dataset/result_new/{}".format(path.split('/')[-1]), img)
cv2.imshow("img", img)
cv2.waitKey(1000) & 0xFF
torch.save(predictions, os.path.join(final_output_dir, 'predictions.pth'))
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)
if args.onnx_export:
torch_out = torch.onnx._export(model,
torch.rand(1, 3, config.IMAGE_SIZE),
osp.join(final_output_dir,
args.onnx_export),
export_params=True)
return
gpus = list(config.GPUS)
if gpus[0] > -1:
model = nn.DataParallel(model, device_ids=gpus).cuda()
# load model
if gpus[0] > -1:
state_dict = torch.load(args.model_file)
else:
state_dict = torch.load(args.model_file, map_location='cpu')
if 'state_dict' in state_dict.keys():
state_dict = state_dict['state_dict']
model.load_state_dict(state_dict)
else:
if gpus[0] > -1:
model.module.load_state_dict(state_dict)
else:
model.load_state_dict(state_dict)
dataset_type = get_dataset(config)
dataset = dataset_type(config, is_train=False)
test_loader = DataLoader(dataset=dataset,
batch_size=config.TEST.BATCH_SIZE_PER_GPU *
len(gpus),
shuffle=False,
num_workers=config.WORKERS,
pin_memory=config.PIN_MEMORY)
ipdb.set_trace()
nme, predictions = function.inference(config, test_loader, model)
torch.save(predictions, os.path.join(final_output_dir, 'predictions.pth'))
def main():
args = parse_args()
reset_config(config, args)
logger, final_output_dir = create_logger(config, args.cfg, 'valid')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = config.CUDNN.ENABLED
model = models.pose3d_resnet.get_pose_net(config, is_train=False)
gpus = [int(i) for i in config.GPUS.split(',')]
model = torch.nn.DataParallel(model, device_ids=gpus).cuda()
# define loss function (criterion) and optimizer
validate = validate_integral
evaluate = eval_integral
# Start from a model
if not (config.MODEL.RESUME is ''):
checkpoint = torch.load(config.MODEL.RESUME)
if 'epoch' in checkpoint.keys():
config.TRAIN.BEGIN_EPOCH = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
logger.info('=> resume from pretrained model {}'.format(
config.MODEL.RESUME))
else:
model.load_state_dict(checkpoint)
logger.info('=> resume from pretrained model {}'.format(
config.MODEL.RESUME))
ds = eval('dataset.' + config.DATASET.DATASET)
# Data loading code
valid_dataset = ds(cfg=config,
root=config.DATASET.ROOT,
image_set=config.DATASET.TEST_SET,
is_train=False)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=config.TEST.BATCH_SIZE * len(gpus),
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
# evaluate on validation set
preds_in_patch_with_score = validate(valid_loader, model)
acc = evaluate(0, preds_in_patch_with_score, valid_loader,
final_output_dir, config.DEBUG.DEBUG)
print(acc)
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)
model = eval('models.' + config.MODEL.NAME + '.get_face_alignment_net')(
config, is_train=True)
gpus = list(config.GPUS)
model = nn.DataParallel(model, device_ids=gpus).cuda()
# 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)
if args.model_file:
logger.info('=> loading model from {}'.format(args.model_file))
# model.load_state_dict(torch.load(args.model_file), strict=False)
model_state = torch.load(args.model_file)
model.module.load_state_dict(model_state.state_dict())
else:
model_state_file = os.path.join(final_output_dir, 'final_state.pth')
logger.info('=> loading model from {}'.format(model_state_file))
model_state = torch.load(model_state_file)
model.module.load_state_dict(model_state)
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'))
def main(cfg):
if cfg.SEED_VALUE >= 0:
print(f'Seed value for the experiment {cfg.SEED_VALUE}')
os.environ['PYTHONHASHSEED'] = str(cfg.SEED_VALUE)
random.seed(cfg.SEED_VALUE)
torch.manual_seed(cfg.SEED_VALUE)
np.random.seed(cfg.SEED_VALUE)
logger = create_logger(cfg.LOGDIR, phase='train')
'''
logger.info(f'GPU name -> {torch.cuda.get_device_name()}')
logger.info(f'GPU feat -> {torch.cuda.get_device_properties("cuda")}')
logger.info(pprint.pformat(cfg))
# cudnn related setting
cudnn.benchmark = cfg.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED
writer = SummaryWriter(log_dir=cfg.LOGDIR)
writer.add_text('config', pprint.pformat(cfg), 0)
'''
# ========= Dataloaders ========= #
data_loaders = get_data_loaders(cfg)
# ========= Compile Loss ========= #
loss = VIBELoss(
e_pose_loss_weight=cfg.LOSS.POSE_W,
e_shape_loss_weight=cfg.LOSS.SHAPE_W,
)
print(loss)
# ========= Initialize networks, optimizers and lr_schedulers ========= #
# temporal generator include the neural network ResNet50,temporal encoder and smpl regressor
# VIBE
generator = VIBE_Demo(
batch_size=cfg.TRAIN.BATCH_SIZE, # 小批量训练 64
)
print(generator)
a = torch.ones((1, 2048))
z = generator(a)
for c in z:
for key in c.keys():
print(key)
print(c[key].shape)
# 定义generator 模型优化算法 常见的优化算法有: sgd,adam
gen_optimizer = get_optimizer(
model=generator, # 模型
optim_type=cfg.TRAIN.GEN_OPTIM, # 使用优化算法的类型，有：sgd,adam
lr=cfg.TRAIN.GEN_LR, # 学习率
weight_decay=cfg.TRAIN.GEN_WD, # regularization 超参数设置
momentum=cfg.TRAIN.GEN_MOMENTUM, # 动量法 超参数
)
def main():
args = parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = '{}'.format(args.gpu)
reset_config(config, args)
logger, final_output_dir, tb_log_dir = create_logger(
config, args.cfg, 'valid')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = config.CUDNN.ENABLED
model = eval('models.posenet.get_pose_net')(config, is_train=False)
model.eval()
params = count_parameters_in_MB(model)
logger.info("Params = %.2fMB" % params)
mult_adds = comp_multadds(model,
input_size=(3, config.MODEL.IMAGE_SIZE[1],
config.MODEL.IMAGE_SIZE[0]))
logger.info("Mult-Adds = %.2fMB" % mult_adds)
model.train()
if config.TEST.MODEL_FILE:
logger.info('=> loading model from {}'.format(config.TEST.MODEL_FILE))
print(config.TEST.MODEL_FILE)
model.load_state_dict(torch.load(config.TEST.MODEL_FILE))
else:
model_state_file = os.path.join(final_output_dir,
'final_state.pth.tar')
logger.info('=> loading model from {}'.format(model_state_file))
model.load_state_dict(torch.load(model_state_file))
model = model.cuda()
# define loss function (criterion) and optimizer
criterion = JointsMSELoss(
use_target_weight=config.LOSS.USE_TARGET_WEIGHT).cuda()
# Data loading code
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
valid_dataset = eval('dataset.' + config.DATASET.DATASET)(
config, config.DATASET.ROOT, config.DATASET.TEST_SET, False,
transforms.Compose([
transforms.ToTensor(),
normalize,
]))
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=config.TEST.BATCH_SIZE,
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
# evaluate on validation set
validate(config, valid_loader, valid_dataset, model, criterion,
final_output_dir, tb_log_dir)
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()
# 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 = torch.load(args.model_file)
model.eval()
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)
import cv2
img = cv2.imread('data/wflw/images/my3.jpg')
print(predictions, predictions.shape)
for item in predictions[0]:
cv2.circle(img, (item[0], item[1]), 3, (0, 0, 255), -1)
cv2.imwrite('out.png', img)
torch.save(predictions, os.path.join(final_output_dir, 'predictions.pth'))
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)
if torch.cuda.is_available():
model = nn.DataParallel(model, device_ids=gpus).cuda()
# load model
state_dict = torch.load(args.model_file)
else:
# model = nn.DataParallel(model)
state_dict = torch.load(args.model_file,
map_location=lambda storage, loc: storage)
if 'state_dict' in state_dict.keys():
state_dict = state_dict['state_dict']
model.load_state_dict(state_dict)
else:
try:
model.module.load_state_dict(state_dict.state_dict())
except AttributeError:
state_dict ## remove first seven
model.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)
#nme, predictions = function.inference(config, test_loader, model, args.model_file) #### testing
function.test(config, test_loader, model, args.model_file) #### testing
def main():
args = parse_args()
update_config(cfg, args)
logger, final_output_dir, tb_log_dir = create_logger(
cfg, args.cfg, 'valid')
logger.info(pprint.pformat(args))
logger.info(cfg)
# cudnn related setting
cudnn.benchmark = cfg.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED
model = eval('models.' + cfg.MODEL.NAME + '.get_pose_net')(
cfg, is_train=False
)
if cfg.TEST.MODEL_FILE != '':
print(cfg.TEST.MODEL_FILE)
logger.info('=> loading model from {}'.format(cfg.TEST.MODEL_FILE))
model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE), strict=False)
else:
model_state_file = os.path.join(
final_output_dir, 'model_best.pth'
)
logger.info('=> loading model from {}'.format(model_state_file))
model.load_state_dict(torch.load(model_state_file))
model = torch.nn.DataParallel(model, device_ids=cfg.GPUS).cuda()
# define loss function (criterion) and optimizer
criterion = JointsMSELoss(
use_target_weight=cfg.LOSS.USE_TARGET_WEIGHT
).cuda()
# Data loading code
valid_dataset = get_valid_dataset(cfg)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=cfg.TEST.BATCH_SIZE_PER_GPU * len(cfg.GPUS),
shuffle=False,
num_workers=cfg.WORKERS,
pin_memory=True
)
# evaluate on validation set
validate(cfg, valid_loader, valid_dataset, model, criterion, final_output_dir)
def main():
# Config
args = parse_args()
update_config(cfg, args)
# Build logger
logger, final_output_dir, tb_log_dir = create_logger(
cfg, args.cfg, 'valid')
logger.info(pprint.pformat(args))
logger.info(cfg)
# cudnn related setting
cudnn.benchmark = cfg.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED
# Build model
model = eval('models.' + cfg.MODEL.NAME + '.get_pose_net')(cfg,
is_train=False)
if cfg.TEST.MODEL_FILE:
logger.info('=> loading model from {}'.format(cfg.TEST.MODEL_FILE))
model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE), strict=False)
else:
model_state_file = os.path.join(final_output_dir, 'final_state.pth')
logger.info('=> loading model from {}'.format(model_state_file))
model.load_state_dict(torch.load(model_state_file))
model = torch.nn.DataParallel(model, device_ids=cfg.GPUS).cuda()
# Build loss
criterion = JointsMSELoss(
use_target_weight=cfg.LOSS.USE_TARGET_WEIGHT).cuda()
# Build dataloader
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
valid_dataset = eval('dataset.' + cfg.DATASET.DATASET)(
cfg, cfg.DATASET.ROOT, cfg.DATASET.TEST_SET, False,
transforms.Compose([transforms.ToTensor(), normalize]))
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=cfg.TEST.BATCH_SIZE_PER_GPU * len(cfg.GPUS),
shuffle=False,
num_workers=cfg.WORKERS,
pin_memory=True)
# Evaluation
validate(cfg, valid_loader, valid_dataset, model, criterion,
final_output_dir, tb_log_dir)
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 = torchvision.models.resnet101(pretrained=config.MODEL.PRETRAINED,
progress=True)
num_ftrs = model.fc.in_features
model.fc = torch.nn.Linear(num_ftrs, config.MODEL.OUTPUT_SIZE[0])
gpus = list(config.GPUS)
model = nn.DataParallel(model, device_ids=gpus).cuda()
# 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)
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)
torch.save(predictions, os.path.join(final_output_dir, 'predictions.pth'))
def main_300w():
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=False)
fp = open('data/300w/face_landmarks70_300w_test.csv', 'w')
for i in range(len(dataset)):
# ipdb.set_trace()
img, fname, meta = dataset[i]
filename = osp.join('data/300w/xximages', fname)
if not osp.exists(osp.dirname(filename)):
os.makedirs(osp.dirname(filename))
scale = meta['scale']
center = meta['center']
tpts = meta['tpts']
selpts = []
for j in range(0, 68):
selpts.append(tpts[j])
selpts.append(tpts[36:42].mean(0))
selpts.append(tpts[42:48].mean(0))
fp.write('%s,%.2f,%.1f,%.1f' % (fname, scale, center[0], center[1]))
for spt in selpts:
img = cv2.circle(img, (spt[0], spt[1]), 1 + center[0] // 400,
(255, 0, 0))
fp.write(',%f,%f' % (spt[0], spt[1]))
fp.write('\n')
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
cv2.imwrite(filename, img)
fp.close()
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():
# Step 1: load model
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()
# 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)
# Step 2: detect face and predict landmark
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)
dataset = get_preprocess(config)
# print('--------bboxes: ', bounding_boxes)
for box in bounding_boxes:
# x1, y1, x2, y2, _ = list(map(int, box))
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.3)
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)
# center_w = (x1+x2)/2
# center_h = (y1+y2)/2
# center = torch.Tensor([center_w, center_h])
# input = img[y1:y2, x1:x2, :]
# input = dataset._preprocessing(dataset, img=img, center=center, scale=1.0)
landmarks = get_lmks_by_img(model, cropped)
img = cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
for (x, y) in landmarks.astype(np.int32):
cv2.circle(img, (x1 + x, y1 + y), 2, (255, 255, 255))
cv2.imshow('0', img)
if cv2.waitKey(10) == 27:
break
def main(cfg):
if cfg.SEED_VALUE >= 0:
print(f'Seed value for the experiment {cfg.SEED_VALUE}')
os.environ['PYTHONHASHSEED'] = str(cfg.SEED_VALUE)
random.seed(cfg.SEED_VALUE)
torch.manual_seed(cfg.SEED_VALUE)
np.random.seed(cfg.SEED_VALUE)
logger = create_logger(cfg.LOGDIR, phase='train')
logger.info(f'GPU name -> {torch.cuda.get_device_name()}')
logger.info(f'GPU feat -> {torch.cuda.get_device_properties("cuda")}')
logger.info(pprint.pformat(cfg))
# cudnn related setting
cudnn.benchmark = cfg.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED
writer = SummaryWriter(log_dir=cfg.LOGDIR)
writer.add_text('config', pprint.pformat(cfg), 0)
# ========= Dataloaders ========= #
data_loaders = get_data_loaders(cfg)
# ========= Compile Loss ========= #
loss = VIBELoss(
e_pose_loss_weight=cfg.LOSS.POSE_W,
e_shape_loss_weight=cfg.LOSS.SHAPE_W,
)
# ========= Initialize networks, optimizers and lr_schedulers ========= #
# temporal generator include the neural network ResNet50,temporal encoder and smpl regressor
# VIBE
generator = VIBE_Demo(
batch_size=cfg.TRAIN.BATCH_SIZE, # 小批量训练 64
).to(cfg.DEVICE)
# 定义generator 模型优化算法 常见的优化算法有: sgd,adam
gen_optimizer = get_optimizer(
model=generator, # 模型
optim_type=cfg.TRAIN.GEN_OPTIM, # 使用优化算法的类型，有：sgd,adam
lr=cfg.TRAIN.GEN_LR, # 学习率
weight_decay=cfg.TRAIN.GEN_WD, # regularization 超参数设置
momentum=cfg.TRAIN.GEN_MOMENTUM, # 动量法 超参数
)
# generator 中学习率的调整
lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
gen_optimizer,
mode='min',
factor=0.1,
patience=cfg.TRAIN.LR_PATIENCE,
verbose=True,
)
# ========= Start Training ========= #
Trainer(
data_loaders=data_loaders,
generator=generator,
criterion=loss,
gen_optimizer=gen_optimizer,
start_epoch=cfg.TRAIN.START_EPOCH,
end_epoch=cfg.TRAIN.END_EPOCH,
device=cfg.DEVICE,
writer=writer,
debug=cfg.DEBUG,
logdir=cfg.LOGDIR,
lr_scheduler=lr_scheduler,
num_iters_per_epoch=cfg.TRAIN.NUM_ITERS_PER_EPOCH,
debug_freq=cfg.DEBUG_FREQ,
).fit()
def main(cfg):
if cfg.SEED_VALUE >= 0:
print(f'Seed value for the experiment {cfg.SEED_VALUE}')
os.environ['PYTHONHASHSEED'] = str(cfg.SEED_VALUE)
random.seed(cfg.SEED_VALUE)
torch.manual_seed(cfg.SEED_VALUE)
np.random.seed(cfg.SEED_VALUE)
logger = create_logger(cfg.LOGDIR, phase='train')
logger.info(f'GPU name -> {torch.cuda.get_device_name()}')
logger.info(f'GPU feat -> {torch.cuda.get_device_properties("cuda")}')
logger.info(pprint.pformat(cfg))
# cudnn related setting
cudnn.benchmark = cfg.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED
writer = SummaryWriter(log_dir=cfg.LOGDIR)
writer.add_text('config', pprint.pformat(cfg), 0)
# ========= Dataloaders ========= #
data_loaders = get_data_loaders(cfg)
# ========= Compile Loss ========= #
loss = TCMRLoss(
e_loss_weight=cfg.LOSS.KP_2D_W,
e_3d_loss_weight=cfg.LOSS.KP_3D_W,
e_pose_loss_weight=cfg.LOSS.POSE_W,
e_shape_loss_weight=cfg.LOSS.SHAPE_W,
d_motion_loss_weight=cfg.LOSS.D_MOTION_LOSS_W,
)
# ========= Initialize networks, optimizers and lr_schedulers ========= #
generator = TCMR(n_layers=cfg.MODEL.TGRU.NUM_LAYERS,
batch_size=cfg.TRAIN.BATCH_SIZE,
seqlen=cfg.DATASET.SEQLEN,
hidden_size=cfg.MODEL.TGRU.HIDDEN_SIZE,
pretrained=cfg.TRAIN.PRETRAINED_REGRESSOR).to(cfg.DEVICE)
gen_optimizer = get_optimizer(
model=generator,
optim_type=cfg.TRAIN.GEN_OPTIM,
lr=cfg.TRAIN.GEN_LR,
weight_decay=cfg.TRAIN.GEN_WD,
momentum=cfg.TRAIN.GEN_MOMENTUM,
)
motion_discriminator = MotionDiscriminator(
rnn_size=cfg.TRAIN.MOT_DISCR.HIDDEN_SIZE,
input_size=69,
num_layers=cfg.TRAIN.MOT_DISCR.NUM_LAYERS,
output_size=1,
feature_pool=cfg.TRAIN.MOT_DISCR.FEATURE_POOL,
attention_size=None if cfg.TRAIN.MOT_DISCR.FEATURE_POOL != 'attention'
else cfg.TRAIN.MOT_DISCR.ATT.SIZE,
attention_layers=None
if cfg.TRAIN.MOT_DISCR.FEATURE_POOL != 'attention' else
cfg.TRAIN.MOT_DISCR.ATT.LAYERS,
attention_dropout=None
if cfg.TRAIN.MOT_DISCR.FEATURE_POOL != 'attention' else
cfg.TRAIN.MOT_DISCR.ATT.DROPOUT).to(cfg.DEVICE)
dis_motion_optimizer = get_optimizer(model=motion_discriminator,
optim_type=cfg.TRAIN.MOT_DISCR.OPTIM,
lr=cfg.TRAIN.MOT_DISCR.LR,
weight_decay=cfg.TRAIN.MOT_DISCR.WD,
momentum=cfg.TRAIN.MOT_DISCR.MOMENTUM)
motion_lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
dis_motion_optimizer,
mode='min',
factor=0.1,
patience=cfg.TRAIN.LR_PATIENCE,
verbose=True,
)
lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
gen_optimizer,
mode='min',
factor=0.1,
patience=cfg.TRAIN.LR_PATIENCE,
verbose=True,
)
# ========= Start Training ========= #
Trainer(
data_loaders=data_loaders,
generator=generator,
motion_discriminator=motion_discriminator,
criterion=loss,
dis_motion_optimizer=dis_motion_optimizer,
dis_motion_update_steps=cfg.TRAIN.MOT_DISCR.UPDATE_STEPS,
gen_optimizer=gen_optimizer,
start_epoch=cfg.TRAIN.START_EPOCH,
end_epoch=cfg.TRAIN.END_EPOCH,
device=cfg.DEVICE,
writer=writer,
debug=cfg.DEBUG,
logdir=cfg.LOGDIR,
lr_scheduler=lr_scheduler,
motion_lr_scheduler=motion_lr_scheduler,
resume=cfg.TRAIN.RESUME,
num_iters_per_epoch=cfg.TRAIN.NUM_ITERS_PER_EPOCH,
debug_freq=cfg.DEBUG_FREQ,
).fit()
def main():
args = parse_args()
reset_config(config, args)
logger, final_output_dir, tb_log_dir = create_logger(
config, args.cfg, 'train')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = config.CUDNN.ENABLED
model = eval('models.'+config.MODEL.NAME+'.get_seg_net')(
config, is_train=True
)
# copy model file
this_dir = os.path.dirname(__file__)
shutil.copy2(
os.path.join(this_dir, '../lib/models', config.MODEL.NAME + '.py'),
final_output_dir)
writer_dict = {
'writer': SummaryWriter(log_dir=tb_log_dir),
'train_global_steps': 0,
'valid_global_steps': 0,
'vis_global_steps': 0,
}
# dump_input = torch.rand((config.TRAIN.BATCH_SIZE,
# 3,
# config.MODEL.IMAGE_SIZE[1],
# config.MODEL.IMAGE_SIZE[0]))
# writer_dict['writer'].add_graph(model, (dump_input, ), verbose=False)
gpus = [int(i) for i in config.GPUS.split(',')]
model = torch.nn.DataParallel(model, device_ids=gpus).cuda()
optimizer = get_optimizer(config, model)
# Data loading code
if 'xception' in config.MODEL.NAME:
# Xception uses different mean std for input image
normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
else:
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_augs = aug.Compose([aug.RandomScale(0.5, 2.0),
aug.RandomHorizontallyFlip(0.5),
aug.RandomSizedCrop(config.MODEL.IMAGE_SIZE)])
test_augs = None
train_dataset = eval('dataset.'+config.DATASET.DATASET)(
config,
config.DATASET.ROOT,
config.DATASET.TRAIN_SET,
transform=transforms.Compose([
transforms.ToTensor(),
normalize,
]),
augmentations=train_augs
)
valid_dataset = eval('dataset.'+config.DATASET.DATASET)(
config,
config.DATASET.ROOT,
config.DATASET.TEST_SET,
transform=transforms.Compose([
transforms.ToTensor(),
normalize,
]),
augmentations=test_augs
)
# define loss function (criterion) and optimizer
criterion = CrossEntropy2D(ignore_index=255, weight=train_dataset.class_weights).cuda()
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.TRAIN.BATCH_SIZE*len(gpus),
shuffle=config.TRAIN.SHUFFLE,
num_workers=config.WORKERS,
pin_memory=True,
drop_last=True if len(gpus) > 2 else False # PyTorch's DataParallel model cannot handle 0 image on either of the GPUs
)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=config.TEST.BATCH_SIZE,
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True
)
if config.TRAIN.LR_SCHEDULER == 'multistep':
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR
)
elif config.TRAIN.LR_SCHEDULER == 'poly':
max_iter = config.TRAIN.END_EPOCH * len(train_loader)
lr_scheduler = PolynomialLR(optimizer, max_iter=max_iter, decay_iter=1)
elif config.TRAIN.LR_SCHEDULER == 'none':
lr_scheduler = None
else:
raise ValueError('Scheduler {} not supported'.format(config.TRAIN.LR_SCHEDULER))
best_perf = 0.0
best_model = False
for epoch in range(config.TRAIN.BEGIN_EPOCH, config.TRAIN.END_EPOCH):
if config.TRAIN.LR_SCHEDULER == 'multistep':
lr_scheduler.step()
# train for one epoch
train(config, train_loader, model, criterion, optimizer, lr_scheduler, epoch,
final_output_dir, tb_log_dir, writer_dict)
if (epoch + 1) % config.TRAIN.EVAL_INTERVAL == 0:
if not config.MODEL.LEARN_GAMMA:
if float(lr_scheduler.last_epoch) / (lr_scheduler.max_iter * config.TRAIN.NE_ITER_RATIO) <= 1:
gamma = (config.TRAIN.NE_GAMMA_U - config.TRAIN.NE_GAMMA_L) * \
(1 - float(lr_scheduler.last_epoch) / (lr_scheduler.max_iter * config.TRAIN.NE_ITER_RATIO) ) ** \
config.TRAIN.NE_GAMMA_EXP + config.TRAIN.NE_GAMMA_L
else:
gamma = config.TRAIN.NE_GAMMA_L
else:
gamma = None
# evaluate on validation set
perf_indicator = validate(config, valid_loader, valid_dataset, model,
criterion, final_output_dir, tb_log_dir,
writer_dict, gamma=gamma)
if perf_indicator > best_perf:
best_perf = perf_indicator
best_model = True
else:
best_model = False
logger.info('=> saving checkpoint to {}'.format(final_output_dir))
save_checkpoint({
'epoch': epoch + 1,
'model': get_model_name(config),
'state_dict': model.state_dict(),
'perf': perf_indicator,
'optimizer': optimizer.state_dict(),
}, best_model, final_output_dir)
else:
perf_indicator = 0.0
final_model_state_file = os.path.join(final_output_dir,
'final_state.pth.tar')
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()
os.environ['CUDA_VISIBLE_DEVICES'] = '{}'.format(args.gpu)
reset_config(config, args)
logger, final_output_dir, tb_log_dir = create_logger(
config, args.cfg, 'train')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = config.CUDNN.ENABLED
model = eval('{}.get_pose_net'.format(args.model))(config, is_train=True)
model.eval()
params = count_parameters_in_MB(model)
logger.info("Params = %.2fMB" % params)
mult_adds = comp_multadds(model,
input_size=(3, config.MODEL.IMAGE_SIZE[1],
config.MODEL.IMAGE_SIZE[0]))
logger.info("Mult-Adds = %.2fMB" % mult_adds)
model.train()
model = model.cuda()
# copy model file
# define loss function (criterion) and optimizer
criterion = JointsMSELoss(
use_target_weight=config.LOSS.USE_TARGET_WEIGHT).cuda()
optimizer = get_optimizer(config, model)
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR)
# Data loading code
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = eval('dataset.' + config.DATASET.DATASET)(
config, config.DATASET.ROOT, config.DATASET.TRAIN_SET, True,
transforms.Compose([
transforms.ToTensor(),
normalize,
]))
valid_dataset = eval('dataset.' + config.DATASET.DATASET)(
config, config.DATASET.ROOT, config.DATASET.TEST_SET, False,
transforms.Compose([
transforms.ToTensor(),
normalize,
]))
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.TRAIN.BATCH_SIZE,
shuffle=config.TRAIN.SHUFFLE,
num_workers=config.WORKERS,
pin_memory=True)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=config.TEST.BATCH_SIZE,
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
best_perf = 0.0
for epoch in range(config.TRAIN.BEGIN_EPOCH, config.TRAIN.END_EPOCH):
# train for one epoch
train(config, train_loader, model, criterion, optimizer, epoch,
final_output_dir, tb_log_dir)
# evaluate on validation set
perf_indicator = validate(config, valid_loader, valid_dataset, model,
criterion, final_output_dir, tb_log_dir)
if perf_indicator > best_perf:
best_perf = perf_indicator
best_model = True
else:
best_model = False
logger.info('=> saving checkpoint to {}'.format(final_output_dir))
save_checkpoint(
{
'epoch': epoch + 1,
'model': get_model_name(config),
'state_dict': model.state_dict(),
'perf': perf_indicator,
'optimizer': optimizer.state_dict(),
}, best_model, final_output_dir)
final_model_state_file = os.path.join(final_output_dir,
'final_state.pth.tar')
logger.info(
'saving final model state to {}'.format(final_model_state_file))
torch.save(model.state_dict(), final_model_state_file)
lr_scheduler.step()
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():
args = parse_args()
reset_config(config, args)
logger, final_output_dir, tb_log_dir = create_logger(
config, args.cfg, 'valid')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = config.CUDNN.ENABLED
model = eval('models.' + config.MODEL.NAME + '.get_seg_net')(
config, is_train=False)
if config.TEST.MODEL_FILE:
logger.info('=> loading model from {}'.format(config.TEST.MODEL_FILE))
checkpoint = torch.load(config.TEST.MODEL_FILE)
if isinstance(checkpoint, OrderedDict):
state_dict_old = checkpoint
elif isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
state_dict_old = checkpoint['state_dict']
else:
raise RuntimeError(
'No state_dict found in checkpoint file {}'.format(
config.TEST.MODEL_FILE))
state_dict = OrderedDict()
# delete 'module.' because it is saved from DataParallel module
for key in state_dict_old.keys():
if key.startswith('module.'):
# state_dict[key[7:]] = state_dict[key]
# state_dict.pop(key)
state_dict[key[7:]] = state_dict_old[key]
else:
state_dict[key] = state_dict_old[key]
model.load_state_dict(state_dict)
else:
model_state_file = os.path.join(final_output_dir,
'final_state.pth.tar')
logger.info('=> loading model from {}'.format(model_state_file))
model.load_state_dict(torch.load(model_state_file))
writer_dict = {
'writer': SummaryWriter(log_dir=tb_log_dir),
'valid_global_steps': 0,
'vis_global_steps': 0,
}
# dump_input = torch.rand((config.TEST.BATCH_SIZE,
# 3,
# config.MODEL.IMAGE_SIZE[1],
# config.MODEL.IMAGE_SIZE[0]))
# writer_dict['writer'].add_graph(model, (dump_input, ), verbose=False)
gpus = [int(i) for i in config.GPUS.split(',')]
model = torch.nn.DataParallel(model, device_ids=gpus).cuda()
# define loss function (criterion)
criterion = CrossEntropy2D(ignore_index=255).cuda()
# Data loading code
if 'xception' in config.MODEL.NAME:
# Xception uses different mean std for input image
normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
else:
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
if 'seg_fpn' in config.MODEL.NAME:
test_augs = aug.Compose([aug.PadByStride(32)])
else:
test_augs = None
# test_augs = aug.FreeScale((config.MODEL.IMAGE_SIZE[0], config.MODEL.IMAGE_SIZE[1]))
valid_dataset = eval('dataset.' + config.DATASET.DATASET)(
config,
config.DATASET.ROOT,
config.DATASET.TEST_SET,
transform=transforms.Compose([
transforms.ToTensor(),
normalize,
]),
augmentations=test_augs)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=config.TEST.BATCH_SIZE,
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
# evaluate on validation set
perf_indicator = validate(config, valid_loader, valid_dataset, model,
criterion, final_output_dir, tb_log_dir,
writer_dict)
writer_dict['writer'].close()
def main():
final_output_dir = 'output'
args = parse_args()
reset_config(config, args)
logger, final_output_dir, tb_log_dir = create_logger(config, 'train')
logger.info(pprint.pformat(config))
# CuDNN
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.deterministic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
# HRNet Model
mv_hrnet = get_pose_net(config, is_train=True)
#pose_hrnet = get_pose_net(config, is_train=True) # Pose estimation model
#pose_hrnet.load_state_dict(torch.load(config.NETWORK.PRETRAINED), strict=False) # Pretrained weight loading
#mv_hrnet = get_multiview_pose_net(pose_hrnet, config) # Multiview adopting
#depth_hrnet = get_pose_net(config, is_train=True) # 2.5D depth prediction model
# Multi GPUs Setting
gpus = [int(i) for i in config.GPUS.split(',')]
mv_hrnet = torch.nn.DataParallel(mv_hrnet, device_ids=gpus).cuda()
logger.info('=> init data parallel model')
# Loss
criterion = JointsMSELoss(
use_target_weight=config.LOSS.USE_TARGET_WEIGHT).cuda()
logger.info('=> init criterion')
# Optimizer
optimizer = get_optimizer(config, mv_hrnet)
logger.info('=> init {} optimizer'.format(config.TRAIN.OPTIMIZER))
# Loading checkpoint
start_epoch = config.TRAIN.BEGIN_EPOCH
if config.TRAIN.RESUME:
start_epoch, mv_hrnet, optimizer = load_checkpoint(
mv_hrnet, optimizer, final_output_dir)
# Scheduler
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR)
logger.info('=> init scheduler')
# Summary
writer_dict = {
'writer': SummaryWriter(log_dir=tb_log_dir),
'train_global_steps': 0,
'valid_global_steps': 0,
}
# Data loader
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
logger.info('=> loading train dataset')
train_dataset = H36MDataset(
config, config.DATASET.TRAIN_SUBSET, True,
transforms.Compose([transforms.ToTensor(), normalize]))
#train_dataset = MultiViewH36M(config, config.DATASET.TRAIN_SUBSET, True, transforms.Compose([transforms.ToTensor(), normalize]))
logger.info('=> loading validation dataset')
valid_dataset = H36MDataset(
config, config.DATASET.TEST_SUBSET, False,
transforms.Compose([transforms.ToTensor(), normalize]))
logger.info('=> loading train dataloader')
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.TRAIN.BATCH_SIZE * len(gpus),
shuffle=config.TRAIN.SHUFFLE,
num_workers=config.WORKERS,
pin_memory=True)
logger.info('=> loading valid dataloader')
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=config.TEST.BATCH_SIZE * len(gpus),
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
# Training loop
best_perf = 0.0
best_model = False
for epoch in range(start_epoch, config.TRAIN.END_EPOCH):
lr_scheduler.step()
# Trainer
train(config, train_loader, mv_hrnet, criterion, optimizer, epoch,
final_output_dir, tb_log_dir, writer_dict)
# Performance indicator
perf_indicator = validate(config, valid_loader, valid_dataset,
mv_hrnet, criterion, final_output_dir,
tb_log_dir, writer_dict)
if perf_indicator > best_perf:
best_perf = perf_indicator
best_model = True
else:
best_model = False
logger.info('=> saving checkpoint to {}'.format(final_output_dir))
save_checkpoint(
{
'epoch': epoch + 1,
'model': get_model_name(config),
'state_dict': mv_hrnet.module.state_dict(),
'perf': perf_indicator,
'optimizer': optimizer.state_dict(),
}, best_model, final_output_dir)
# End
final_model_state_file = os.path.join(final_output_dir,
'final_state.pth.tar')
logger.info(
'saving final model state to {}'.format(final_model_state_file))
torch.save(mv_hrnet.module.state_dict(), final_model_state_file)
writer_dict['writer'].close()
def main():
args = parse_args()
logger, final_output_dir, tb_log_dir = create_logger(
config, args.cfg, 'train')
logger.info(pprint.pformat(args))
logger.info(config)
writer_dict = {
'writer': SummaryWriter(tb_log_dir),
'train_global_steps': 0,
'valid_global_steps': 0,
}
# cudnn related setting
#cudnn.benchmark = config.CUDNN.BENCHMARK
#cudnn.deterministic = config.CUDNN.DETERMINISTIC
#cudnn.enabled = config.CUDNN.ENABLED
gpus = list(config.GPUS)
# build model
model = get_seg_model(config)
dump_input = torch.rand(
(1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0]))
logger.info(get_model_summary(model.cuda(), dump_input.cuda()))
# copy model file
this_dir = os.path.dirname(__file__)
models_dst_dir = os.path.join(final_output_dir, 'models')
if os.path.exists(models_dst_dir):
shutil.rmtree(models_dst_dir)
shutil.copytree(os.path.join(this_dir, '../lib/models'), models_dst_dir)
# prepare data
crop_size = (config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0])
trfm = A.Compose([
A.Resize(crop_size[0], crop_size[1]),
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.5),
A.OneOf([
A.RandomGamma(),
A.RandomBrightnessContrast(),
A.ColorJitter(brightness=0.07,
contrast=0.07,
saturation=0.1,
hue=0.1,
always_apply=False,
p=0.3),
],
p=0.3),
A.OneOf([
A.ElasticTransform(
alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03),
A.GridDistortion(),
A.OpticalDistortion(distort_limit=2, shift_limit=0.5),
],
p=0.0),
A.ShiftScaleRotate(),
])
ds = HubDataset(config.DATA_PATH,
window=config.WINDOW,
overlap=config.MIN_OVERLAP,
transform=trfm)
valid_idx, train_idx = [], []
for i in range(len(ds)):
if ds.slices[i][0] == 7:
valid_idx.append(i)
else:
train_idx.append(i)
train_dataset = torch.utils.data.Subset(ds, train_idx)
test_dataset = torch.utils.data.Subset(ds, valid_idx)
# define training and validation data loaders
trainloader = DataLoader(train_dataset,
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU *
len(gpus),
shuffle=True,
num_workers=config.WORKERS,
drop_last=True)
testloader = DataLoader(test_dataset,
batch_size=config.TEST.BATCH_SIZE_PER_GPU *
len(gpus),
shuffle=False,
num_workers=config.WORKERS,
drop_last=False)
"""
train_dataset = eval('datasets.'+config.DATASET.DATASET)(
root=config.DATASET.ROOT,
list_path=config.DATASET.TRAIN_SET,
num_samples=None,
num_classes=config.DATASET.NUM_CLASSES,
multi_scale=config.TRAIN.MULTI_SCALE,
flip=config.TRAIN.FLIP,
ignore_label=config.TRAIN.IGNORE_LABEL,
base_size=config.TRAIN.BASE_SIZE,
crop_size=crop_size,
downsample_rate=config.TRAIN.DOWNSAMPLERATE,
scale_factor=config.TRAIN.SCALE_FACTOR)
trainloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU*len(gpus),
shuffle=config.TRAIN.SHUFFLE,
num_workers=config.WORKERS,
pin_memory=True,
drop_last=True)
if config.DATASET.EXTRA_TRAIN_SET:
extra_train_dataset = eval('datasets.'+config.DATASET.DATASET)(
root=config.DATASET.ROOT,
list_path=config.DATASET.EXTRA_TRAIN_SET,
num_samples=None,
num_classes=config.DATASET.NUM_CLASSES,
multi_scale=config.TRAIN.MULTI_SCALE,
flip=config.TRAIN.FLIP,
ignore_label=config.TRAIN.IGNORE_LABEL,
base_size=config.TRAIN.BASE_SIZE,
crop_size=crop_size,
downsample_rate=config.TRAIN.DOWNSAMPLERATE,
scale_factor=config.TRAIN.SCALE_FACTOR)
extra_trainloader = torch.utils.data.DataLoader(
extra_train_dataset,
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU*len(gpus),
shuffle=config.TRAIN.SHUFFLE,
num_workers=config.WORKERS,
pin_memory=True,
drop_last=True)
test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0])
test_dataset = eval('datasets.'+config.DATASET.DATASET)(
root=config.DATASET.ROOT,
list_path=config.DATASET.TEST_SET,
num_samples=config.TEST.NUM_SAMPLES,
num_classes=config.DATASET.NUM_CLASSES,
multi_scale=False,
flip=False,
ignore_label=config.TRAIN.IGNORE_LABEL,
base_size=config.TEST.BASE_SIZE,
crop_size=test_size,
downsample_rate=1)
testloader = torch.utils.data.DataLoader(
test_dataset,
batch_size=config.TEST.BATCH_SIZE_PER_GPU*len(gpus),
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
"""
# criterion
if config.LOSS.USE_OHEM:
criterion = OhemCrossEntropy(ignore_label=config.TRAIN.IGNORE_LABEL,
thres=config.LOSS.OHEMTHRES,
min_kept=config.LOSS.OHEMKEEP,
weight=train_dataset.class_weights)
else:
"""
criterion = CrossEntropy(ignore_label=config.TRAIN.IGNORE_LABEL,
weight=train_dataset.class_weights)
"""
criterion = list()
criterion.append(nn.BCEWithLogitsLoss())
criterion.append(SoftDiceLoss())
model = FullModel(model, criterion).cuda()
#model = nn.DataParallel(model, device_ids=gpus).cuda()
# optimizer
if config.TRAIN.OPTIMIZER == 'sgd':
optimizer = torch.optim.SGD(
[{
'params': filter(lambda p: p.requires_grad,
model.parameters()),
'lr': config.TRAIN.LR
}],
lr=config.TRAIN.LR,
momentum=config.TRAIN.MOMENTUM,
weight_decay=config.TRAIN.WD,
nesterov=config.TRAIN.NESTEROV,
)
else:
raise ValueError('Only Support SGD optimizer')
epoch_iters = np.int(train_dataset.__len__() /
config.TRAIN.BATCH_SIZE_PER_GPU / len(gpus))
best_mIoU = 0
last_epoch = 0
if config.TRAIN.RESUME:
model_state_file = os.path.join(final_output_dir, 'checkpoint.pth.tar')
if os.path.isfile(model_state_file):
checkpoint = torch.load(model_state_file)
best_mIoU = checkpoint['best_mIoU']
last_epoch = checkpoint['epoch']
model.module.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger.info("=> loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
start = timeit.default_timer()
end_epoch = config.TRAIN.END_EPOCH + config.TRAIN.EXTRA_EPOCH
num_iters = config.TRAIN.END_EPOCH * epoch_iters
extra_iters = config.TRAIN.EXTRA_EPOCH * epoch_iters
for epoch in range(last_epoch, end_epoch):
if epoch >= config.TRAIN.END_EPOCH:
train(config, epoch - config.TRAIN.END_EPOCH,
config.TRAIN.EXTRA_EPOCH, epoch_iters, config.TRAIN.EXTRA_LR,
extra_iters, extra_trainloader, optimizer, model,
writer_dict)
else:
train(config, epoch, config.TRAIN.END_EPOCH, epoch_iters,
config.TRAIN.LR, num_iters, trainloader, optimizer, model,
writer_dict)
logger.info('=> saving checkpoint to {}'.format(final_output_dir +
'checkpoint.pth.tar'))
torch.save(
{
'epoch': epoch + 1,
'best_mIoU': best_mIoU,
'state_dict': model.module.state_dict(),
'optimizer': optimizer.state_dict(),
}, os.path.join(final_output_dir, 'checkpoint.pth.tar'))
valid_loss, mean_IoU, IoU_array = validate(config, testloader, model,
writer_dict)
if mean_IoU > best_mIoU:
best_mIoU = mean_IoU
torch.save(model.module.state_dict(),
os.path.join(final_output_dir, 'best.pth'))
msg = 'Loss: {:.3f}, MeanIU: {: 4.4f}, Best_mIoU: {: 4.4f}'.format(
valid_loss, mean_IoU, best_mIoU)
logging.info(msg)
logging.info(IoU_array)
torch.save(model.module.state_dict(),
os.path.join(final_output_dir, 'final_state.pth'))
writer_dict['writer'].close()
end = timeit.default_timer()
logger.info('Hours: %d' % np.int((end - start) / 3600))
logger.info('Done')
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, is_train=True)
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_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)
# 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, '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")
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,
"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():
# 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():
# set all the configurations
args = parse_args()
update_config(cfg, args)
# set the logger, tb_log_dir means tensorboard logdir
logger, final_output_dir, tb_log_dir = create_logger(
cfg, args.cfg, 'train')
logger.info(pprint.pformat(args))
logger.info(cfg)
writer_dict = {
'writer': SummaryWriter(log_dir=tb_log_dir),
'train_global_steps': 0,
'valid_global_steps': 0,
}
# cudnn related setting
cudnn.benchmark = cfg.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED
# bulid up model
model = get_net(cfg)
model = torch.nn.DataParallel(model, device_ids=cfg.GPUS).cuda()
# Data loading
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = eval('dataset.' + cfg.DATASET.DATASET)(
cfg, True, transforms.Compose([
transforms.ToTensor(),
normalize,
]))
valid_dataset = eval('dataset.' + cfg.DATASET.DATASET)(
cfg, False, transforms.Compose([
transforms.ToTensor(),
normalize,
]))
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE_PER_GPU * len(cfg.GPUS),
shuffle=cfg.TRAIN.SHUFFLE,
num_workers=cfg.WORKERS,
pin_memory=cfg.PIN_MEMORY)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=cfg.TEST.BATCH_SIZE_PER_GPU * len(cfg.GPUS),
shuffle=False,
num_workers=cfg.WORKERS,
pin_memory=cfg.PIN_MEMORY)
# define loss function (criterion) and optimizer
criterion = get_loss(cfg).cuda()
optimizer = get_optimizer(cfg, model)
# load checkpoint model
best_perf = 0.0
best_model = False
last_epoch = -1
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
cfg.TRAIN.LR_STEP,
cfg.TRAIN.LR_FACTOR,
last_epoch=last_epoch)
begin_epoch = cfg.TRAIN.BEGIN_EPOCH
checkpoint_file = os.path.join(final_output_dir, 'checkpoint.pth')
if cfg.AUTO_RESUME and os.path.exists(checkpoint_file):
logger.info("=> loading checkpoint '{}'".format(checkpoint_file))
checkpoint = torch.load(checkpoint_file)
begin_epoch = checkpoint['epoch']
best_perf = checkpoint['perf']
last_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger.info("=> loaded checkpoint '{}' (epoch {})".format(
checkpoint_file, checkpoint['epoch']))
# training
for epoch in range(begin_epoch + 1, cfg.TRAIN.END_EPOCH + 1):
# train for one epoch
train(cfg, train_loader, model, criterion, optimizer, epoch,
writer_dict)
lr_scheduler.step()
# evaluate on validation set
if epoch % cfg.TRAIN.VAL_FREQ == 0 or epoch == cfg.TRAIN.END_EPOCH + 1:
perf_indicator = validate(cfg, valid_loader, valid_dataset, model,
criterion, final_output_dir, tb_log_dir,
writer_dict)
if perf_indicator >= best_perf:
best_perf = perf_indicator
best_model = True
else:
best_model = False
# save checkpoint model and best model
logger.info('=> saving checkpoint to {}'.format(final_output_dir))
save_checkpoint(
{
'epoch': epoch,
'model': cfg.MODEL.NAME,
'state_dict': model.state_dict(),
'best_state_dict': model.module.state_dict(),
'perf': perf_indicator,
'optimizer': optimizer.state_dict(),
}, best_model, final_output_dir)
# save final model
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(cfg):
if cfg.SEED_VALUE >= 0:
print(f'Seed value for the experiment {cfg.SEED_VALUE}')
os.environ['PYTHONHASHSEED'] = str(cfg.SEED_VALUE)
random.seed(cfg.SEED_VALUE)
torch.manual_seed(cfg.SEED_VALUE)
np.random.seed(cfg.SEED_VALUE)
logger = create_logger(cfg.LOGDIR, phase='train')
logger.info(f'GPU name -> {torch.cuda.get_device_name()}')
logger.info(f'GPU feat -> {torch.cuda.get_device_properties("cuda")}')
logger.info(pprint.pformat(cfg))
# cudnn related setting
cudnn.benchmark = cfg.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED
writer = SummaryWriter(log_dir=cfg.LOGDIR)
writer.add_text('config', pprint.pformat(cfg), 0)
# ========= Dataloaders ========= #
data_loaders = get_data_loaders(cfg)
# ========= Compile Loss ========= #
loss = VIBELoss(
e_loss_weight=cfg.LOSS.KP_2D_W,
e_3d_loss_weight=cfg.LOSS.KP_3D_W,
e_pose_loss_weight=cfg.LOSS.POSE_W,
e_shape_loss_weight=cfg.LOSS.SHAPE_W,
d_motion_loss_weight=cfg.LOSS.D_MOTION_LOSS_W,
)
# ========= Initialize networks, optimizers and lr_schedulers ========= #
vibe = VIBE(
n_layers=cfg.MODEL.TGRU.NUM_LAYERS,
batch_size=cfg.TRAIN.BATCH_SIZE,
seqlen=cfg.DATASET.SEQLEN,
hidden_size=cfg.MODEL.TGRU.HIDDEN_SIZE,
pretrained=cfg.TRAIN.PRETRAINED_REGRESSOR,
add_linear=cfg.MODEL.TGRU.ADD_LINEAR,
bidirectional=cfg.MODEL.TGRU.BIDIRECTIONAL,
use_residual=cfg.MODEL.TGRU.RESIDUAL,
).to(cfg.DEVICE)
if cfg.TRAIN.PRETRAINED != '' and os.path.isfile(cfg.TRAIN.PRETRAINED):
checkpoint = torch.load(cfg.TRAIN.PRETRAINED)
best_performance = checkpoint['performance']
vibe.load_state_dict(checkpoint['gen_state_dict'])
print(f'==> Loaded pretrained model from {cfg.TRAIN.PRETRAINED}...')
print(f'Performance on 3DPW test set {best_performance}')
else:
print(f'{cfg.TRAIN.PRETRAINED} is not a pretrained model!!!!')
generator = REFINER(vibe=vibe).to(cfg.DEVICE)
gen_optimizer = get_optimizer(
model=generator,
optim_type=cfg.TRAIN.GEN_OPTIM,
lr=cfg.TRAIN.GEN_LR,
weight_decay=cfg.TRAIN.GEN_WD,
momentum=cfg.TRAIN.GEN_MOMENTUM,
)
# motion_discriminator = MotionDiscriminator(
# rnn_size=cfg.TRAIN.MOT_DISCR.HIDDEN_SIZE,
# input_size=69,
# num_layers=cfg.TRAIN.MOT_DISCR.NUM_LAYERS,
# output_size=1,
# feature_pool=cfg.TRAIN.MOT_DISCR.FEATURE_POOL,
# attention_size=None if cfg.TRAIN.MOT_DISCR.FEATURE_POOL !='attention' else cfg.TRAIN.MOT_DISCR.ATT.SIZE,
# attention_layers=None if cfg.TRAIN.MOT_DISCR.FEATURE_POOL !='attention' else cfg.TRAIN.MOT_DISCR.ATT.LAYERS,
# attention_dropout=None if cfg.TRAIN.MOT_DISCR.FEATURE_POOL !='attention' else cfg.TRAIN.MOT_DISCR.ATT.DROPOUT
# ).to(cfg.DEVICE)
# dis_motion_optimizer = get_optimizer(
# model=motion_discriminator,
# optim_type=cfg.TRAIN.MOT_DISCR.OPTIM,
# lr=cfg.TRAIN.MOT_DISCR.LR,
# weight_decay=cfg.TRAIN.MOT_DISCR.WD,
# momentum=cfg.TRAIN.MOT_DISCR.MOMENTUM
# )
# motion_lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
# dis_motion_optimizer,
# mode='min',
# factor=0.1,
# patience=cfg.TRAIN.LR_PATIENCE,
# verbose=True,
# )
lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
gen_optimizer,
mode='min',
factor=0.1,
patience=cfg.TRAIN.LR_PATIENCE,
verbose=True,
)
# ========= Start Training ========= #
motion_discriminator = None
dis_motion_optimizer = None
motion_lr_scheduler = None
Trainer(
data_loaders=data_loaders,
generator=generator,
motion_discriminator=motion_discriminator,
criterion=loss,
dis_motion_optimizer=dis_motion_optimizer,
dis_motion_update_steps=cfg.TRAIN.MOT_DISCR.UPDATE_STEPS,
gen_optimizer=gen_optimizer,
start_epoch=cfg.TRAIN.START_EPOCH,
end_epoch=cfg.TRAIN.END_EPOCH,
device=cfg.DEVICE,
writer=writer,
debug=cfg.DEBUG,
logdir=cfg.LOGDIR,
lr_scheduler=lr_scheduler,
motion_lr_scheduler=motion_lr_scheduler,
resume=cfg.TRAIN.RESUME,
num_iters_per_epoch=cfg.TRAIN.NUM_ITERS_PER_EPOCH,
debug_freq=cfg.DEBUG_FREQ,
).fit()
def main():
args = parse_args()
update_config(cfg, args)
logger, final_output_dir, tb_log_dir = create_logger(
cfg, args.cfg, 'train')
logger.info(pprint.pformat(args))
logger.info(cfg)
# cudnn related setting
cudnn.benchmark = cfg.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED
model = eval('models.' + cfg.MODEL.NAME + '.get_pose_net')(cfg,
is_train=True)
# copy model file
this_dir = os.path.dirname(__file__)
shutil.copy2(
os.path.join(this_dir, '../lib/models', cfg.MODEL.NAME + '.py'),
final_output_dir)
# logger.info(pprint.pformat(model))
writer_dict = {
'writer': SummaryWriter(log_dir=tb_log_dir),
'train_global_steps': 0,
'valid_global_steps': 0,
}
dump_input = torch.rand(
(1, 3, cfg.MODEL.IMAGE_SIZE[1], cfg.MODEL.IMAGE_SIZE[0]))
# writer_dict['writer'].add_graph(model, (dump_input))
logger.info(get_model_summary(model, dump_input))
model = torch.nn.DataParallel(model, device_ids=cfg.GPUS).cuda()
# define loss function (criterion) and optimizer
criterion = JointsMSELoss(
use_target_weight=cfg.LOSS.USE_TARGET_WEIGHT).cuda()
# Data loading code
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = eval('dataset.' + cfg.DATASET.DATASET)(
cfg, cfg.DATASET.ROOT, cfg.DATASET.TRAIN_SET, True,
transforms.Compose([transforms.ToTensor(), normalize]))
valid_dataset = eval('dataset.' + cfg.DATASET.DATASET)(
cfg, cfg.DATASET.ROOT, cfg.DATASET.TEST_SET, False,
transforms.Compose([transforms.ToTensor(), normalize]))
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE_PER_GPU * len(cfg.GPUS),
shuffle=cfg.TRAIN.SHUFFLE,
num_workers=cfg.WORKERS,
pin_memory=cfg.PIN_MEMORY)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=cfg.TEST.BATCH_SIZE_PER_GPU * len(cfg.GPUS),
shuffle=False,
num_workers=cfg.WORKERS,
pin_memory=cfg.PIN_MEMORY)
best_perf = 0.0
best_model = False
last_epoch = -1
optimizer = get_optimizer(cfg, model)
begin_epoch = cfg.TRAIN.BEGIN_EPOCH
checkpoint_file = os.path.join(final_output_dir, 'checkpoint.pth')
if cfg.AUTO_RESUME and os.path.exists(checkpoint_file):
logger.info("=> loading checkpoint '{}'".format(checkpoint_file))
checkpoint = torch.load(checkpoint_file)
begin_epoch = checkpoint['epoch']
best_perf = checkpoint['perf']
last_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger.info("=> loaded checkpoint '{}' (epoch {})".format(
checkpoint_file, checkpoint['epoch']))
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
cfg.TRAIN.LR_STEP,
cfg.TRAIN.LR_FACTOR,
last_epoch=last_epoch)
for epoch in range(begin_epoch, cfg.TRAIN.END_EPOCH):
# train for one epoch
train(cfg, train_loader, model, criterion, optimizer, epoch,
final_output_dir, tb_log_dir, writer_dict)
lr_scheduler.step()
# evaluate on validation set
perf_indicator = validate(cfg, valid_loader, valid_dataset, model,
criterion, final_output_dir, tb_log_dir,
writer_dict)
if perf_indicator >= best_perf:
best_perf = perf_indicator
best_model = True
else:
best_model = False
logger.info('=> saving checkpoint to {}'.format(final_output_dir))
save_checkpoint(
{
'epoch': epoch + 1,
'model': cfg.MODEL.NAME,
'state_dict': model.state_dict(),
'best_state_dict': model.module.state_dict(),
'perf': perf_indicator,
'optimizer': optimizer.state_dict(),
}, best_model, final_output_dir)
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()
