import numpy as np
from PIL import Image
import torch
import torch.utils.data as data
import torchvision.transforms as transforms
import json
import random
sys.path.insert(0, os.getcwd())
from data.face_bbx import FACE
from plot.CommonPlottingOperations import *
from pylib import HumanPts, FaceAug, HumanAug, FacePts
index = 0
val_loader = torch.utils.data.DataLoader(FACE("dataset/wflw_train.json",
".",
is_train=True,
keep_pts_inside=True),
batch_size=1,
shuffle=True,
num_workers=1,
pin_memory=True)
# Visualize some images
#for i, (img, _, points, _, s, visible_multiclass, _) in enumerate(val_loader):
val_loader = torch.utils.data.DataLoader(FACE("dataset/aflw_test_all.json",
".",
is_train=False,
keep_pts_inside=True),
batch_size=1,
shuffle=True,
def main():
opt = TrainOptions().parse()
train_history = TrainHistoryFace()
checkpoint = Checkpoint()
visualizer = Visualizer(opt)
exp_dir = os.path.join(opt.exp_dir, opt.exp_id)
log_name = opt.vis_env + 'log.txt'
visualizer.log_name = os.path.join(exp_dir, log_name)
num_classes = opt.class_num
if not opt.slurm:
os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu_id
layer_num = opt.layer_num
order = opt.order
net = create_cu_net(neck_size= 4, growth_rate= 32, init_chan_num= 128,
class_num= num_classes, layer_num= layer_num, order= order,
loss_num= layer_num, use_spatial_transformer= opt.stn,
mlp_tot_layers= opt.mlp_tot_layers, mlp_hidden_units= opt.mlp_hidden_units,
get_mean_from_mlp= opt.get_mean_from_mlp)
# Load the pre-trained model
saved_wt_file = opt.saved_wt_file
if saved_wt_file == "":
print("=> Training from scratch")
else:
print("=> Loading weights from " + saved_wt_file)
checkpoint_t = torch.load(saved_wt_file)
state_dict = checkpoint_t['state_dict']
tt_names=[]
for names in net.state_dict():
tt_names.append(names)
for name, param in state_dict.items():
name = name[7:]
if name not in net.state_dict():
print("=> not load weights '{}'".format(name))
continue
if isinstance(param, Parameter):
param = param.data
if (net.state_dict()[name].shape[0] == param.shape[0]):
net.state_dict()[name].copy_(param)
else:
print("First dim different. Not loading weights {}".format(name))
if (opt.freeze):
print("\n\t\tFreezing basenet parameters\n")
for param in net.parameters():
param.requires_grad = False
"""
for i in range(layer_num):
net.choleskys[i].fc_1.bias.requires_grad = True
net.choleskys[i].fc_1.weight.requires_grad = True
net.choleskys[i].fc_2.bias.requires_grad = True
net.choleskys[i].fc_2.weight.requires_grad = True
net.choleskys[i].fc_3.bias.requires_grad = True
net.choleskys[i].fc_3.weight.requires_grad = True
"""
net.cholesky.fc_1.bias.requires_grad = True
net.cholesky.fc_1.weight.requires_grad = True
net.cholesky.fc_2.bias.requires_grad = True
net.cholesky.fc_2.weight.requires_grad = True
net.cholesky.fc_3.bias.requires_grad = True
net.cholesky.fc_3.weight.requires_grad = True
else:
print("\n\t\tNot freezing anything. Tuning every parameter\n")
for param in net.parameters():
param.requires_grad = True
net = torch.nn.DataParallel(net).cuda() # use multiple GPUs
# Optimizer
if opt.optimizer == "rmsprop":
optimizer = torch.optim.RMSprop(filter(lambda p: p.requires_grad, net.parameters()), lr=opt.lr, alpha=0.99,
eps=1e-8, momentum=0, weight_decay=0)
elif opt.optimizer == "adam":
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, net.parameters()), lr=opt.lr)
else:
print("Unknown Optimizer. Aborting!!!")
sys.exit(0)
print type(optimizer)
# Optionally resume from a checkpoint
if opt.resume_prefix != '':
# if 'pth' in opt.resume_prefix:
# trunc_index = opt.resume_prefix.index('pth')
# opt.resume_prefix = opt.resume_prefix[0:trunc_index - 1]
checkpoint.save_prefix = os.path.join(exp_dir, opt.resume_prefix)
checkpoint.load_prefix = os.path.join(exp_dir, opt.resume_prefix)[0:-1]
checkpoint.load_checkpoint(net, optimizer, train_history)
else:
checkpoint.save_prefix = exp_dir + '/'
print("Save prefix = {}".format(checkpoint.save_prefix))
# Load data
json_path = opt.json_path
train_json = opt.train_json
val_json = opt.val_json
print("Path added to each image path in JSON = {}".format(json_path))
print("Train JSON path = {}".format(train_json))
print("Val JSON path = {}".format(val_json))
if opt.bulat_aug:
# Use Bulat et al Augmentation Scheme
train_loader = torch.utils.data.DataLoader(
FACE(train_json, json_path, is_train= True, scale_factor= 0.2, rot_factor= 50, use_occlusion= True, keep_pts_inside= True),
batch_size=opt.bs, shuffle= True,
num_workers=opt.nThreads, pin_memory= True)
else:
train_loader = torch.utils.data.DataLoader(
FACE(train_json, json_path, is_train= True, keep_pts_inside= True),
batch_size=opt.bs, shuffle= True,
num_workers=opt.nThreads, pin_memory= True)
val_loader = torch.utils.data.DataLoader(
FACE(val_json, json_path, is_train=False),
batch_size=opt.bs, shuffle=False,
num_workers=opt.nThreads, pin_memory=True)
logger = Logger(os.path.join(opt.exp_dir, opt.exp_id, opt.resume_prefix+'face-training-log.txt'),
title='face-training-summary')
logger.set_names(['Epoch', 'LR', 'Train Loss', 'Val Loss', 'Train RMSE', 'Val RMSE', 'Train RMSE Box', 'Val RMSE Box', 'Train RMSE Meta', 'Val RMSE Meta'])
if not opt.is_train:
visualizer.log_path = os.path.join(opt.exp_dir, opt.exp_id, 'val_log.txt')
val_loss, val_rmse, predictions = validate(val_loader, net,
train_history.epoch[-1]['epoch'], visualizer, num_classes, flip_index)
checkpoint.save_preds(predictions)
return
global weights_HG
weights_HG = [float(x) for x in opt.hg_wt.split(",")]
if opt.is_covariance:
print("Covariance used from the heatmap")
else:
print("Covariance calculated from MLP")
if opt.stn:
print("Using spatial transformer on heatmaps")
print ("Postprocessing applied = {}".format(opt.pp))
if (opt.smax):
print("Scaled softmax used with tau = {}".format(opt.tau))
else:
print("No softmax used")
print("Individual Hourglass loss weights")
print(weights_HG)
print("wt_MSE (tradeoff between GLL and MSE in each hourglass)= " + str(opt.wt_mse))
print("wt_gauss_regln (tradeoff between GLL and Gaussian Regularisation in each hourglass)= " + str(opt.wt_gauss_regln))
if opt.bulat_aug:
print("Using Bulat et al, ICCV 2017 Augmentation Scheme")
print("Using Learning Policy {}".format(opt.lr_policy))
chosen_lr_policy = dict_of_functions[opt.lr_policy]
# Optionally resume from a checkpoint
start_epoch = 0
if opt.resume_prefix != '':
start_epoch = train_history.epoch[-1]['epoch'] + 1
# Training and validation
start_epoch = 0
if opt.resume_prefix != '':
start_epoch = train_history.epoch[-1]['epoch'] + 1
train_loss_orig_epoch = []
train_loss_gau_t1_epoch = []
train_loss_gau_t2_epoch = []
train_nme_orig_epoch = []
train_nme_gau_epoch = []
train_nme_new_epoch = []
val_loss_orig_epoch = []
val_loss_gau_t1_epoch = []
val_loss_gau_t2_epoch = []
val_nme_orig_epoch = []
val_nme_gau_epoch = []
val_nme_new_epoch = []
for epoch in range(start_epoch, opt.nEpochs):
chosen_lr_policy(opt, optimizer, epoch)
# Train for one epoch
train_loss, train_loss_mse,train_loss_gau_t1, train_loss_gau_t2,train_rmse_orig, train_rmse_gau, train_rmse_new_gd_box, train_rmse_new_meta_box = train(train_loader, net, optimizer, epoch, visualizer, opt)
#train_loss_gau_epoch.append(train_loss_gau)
train_loss_gau_t1_epoch.append(train_loss_gau_t1)
train_loss_gau_t2_epoch.append(train_loss_gau_t2)
train_nme_orig_epoch.append(train_rmse_orig)
train_nme_gau_epoch.append(train_rmse_gau)
train_loss_orig_epoch.append(train_loss_mse)
# Evaluate on validation set
val_loss, val_loss_mse, val_loss_gau_t1, val_loss_gau_t2 , val_rmse_orig, val_rmse_gau, val_rmse_new_gd_box, val_rmse_new_meta_box, predictions= validate(val_loader, net, epoch, visualizer, opt, num_classes, flip_index)
val_loss_orig_epoch.append(val_loss_mse)
val_loss_gau_t1_epoch.append(val_loss_gau_t1)
val_loss_gau_t2_epoch.append(val_loss_gau_t2)
val_nme_orig_epoch.append(val_rmse_orig)
val_nme_gau_epoch.append(val_rmse_gau)
# Update training history
e = OrderedDict( [('epoch', epoch)] )
lr = OrderedDict( [('lr', optimizer.param_groups[0]['lr'])] )
loss = OrderedDict( [('train_loss', train_loss),('val_loss', val_loss)] )
rmse = OrderedDict( [('val_rmse', val_rmse_gau)] )
train_history.update(e, lr, loss, rmse)
checkpoint.save_checkpoint(net, optimizer, train_history, predictions)
visualizer.plot_train_history_face(train_history)
logger.append([epoch, optimizer.param_groups[0]['lr'], train_loss, val_loss, train_rmse_gau, val_rmse_gau, train_rmse_new_gd_box, val_rmse_new_gd_box, train_rmse_new_meta_box, val_rmse_new_meta_box])
logger.close()
from PIL import Image
import torch
import torch.utils.data as data
import torchvision.transforms as transforms
import json
import random
sys.path.insert(0, os.getcwd())
from data.face_bbx import FACE
from plot.CommonPlottingOperations import *
from pylib import HumanPts, FaceAug, HumanAug, FacePts
index = 0
val_loader = torch.utils.data.DataLoader(FACE(
"dataset/aflw_ours_frontal_train.json",
".",
is_train=True,
keep_pts_inside=True),
batch_size=1,
shuffle=True,
num_workers=1,
pin_memory=True)
# Visualize some images
for i, (img, _, points, _, s, visible_multiclass, _) in enumerate(val_loader):
print(i)
image = img[index].numpy()
pts = points[index].numpy()
vis = visible_multiclass.clone()
vis[visible_multiclass > 1] = 1
import torch
import torch.utils.data as data
import torchvision.transforms as transforms
import json
import random
sys.path.insert(0, os.getcwd())
from data.face_bbx import FACE
from plot.CommonPlottingOperations import *
from pylib import HumanPts, FaceAug, HumanAug, FacePts
from torch.autograd import Variable
index = 0
val_loader = torch.utils.data.DataLoader(FACE(
"dataset/aflw_ours_left_val.json",
".",
is_train=False,
keep_pts_inside=True),
batch_size=1,
shuffle=True,
num_workers=1,
pin_memory=True)
# Visualize some images
for i, (img, _, points, _, _, _, _, s, visible_multiclass,
_) in enumerate(val_loader):
print(i)
vis = visible_multiclass.clone()
vis[vis > 1] = 1
# vis with zero points was producing weird error. Add a small constant
from PIL import Image
import torch
import torch.utils.data as data
import torchvision.transforms as transforms
import json
import random
sys.path.insert(0, os.getcwd())
from data.face_bbx import FACE
from plot.CommonPlottingOperations import *
from pylib import HumanPts, FaceAug, HumanAug, FacePts
batch_size = 24
index = 0
train_loader = torch.utils.data.DataLoader(FACE("dataset/aflw_train.json",
".",
is_train=True),
batch_size=batch_size,
shuffle=False,
num_workers=1,
pin_memory=True)
# Visualize some images
# inp, heatmap, pts_input_res, heatmap_mask, s_rand, visible_multiclass, meta_box_size_input_res
for i, (img, _, points, _, s, visible_multiclass,
meta_box_size_input_res) in enumerate(train_loader):
#print(i)
image = img[index].numpy()
pts = points[index].numpy()
#plt.figure(figsize=(16, 8))
def main():
opt = TrainOptions().parse()
train_history = TrainHistoryFace()
checkpoint = Checkpoint()
visualizer = Visualizer(opt)
exp_dir = os.path.join(opt.exp_dir, opt.exp_id)
log_name = opt.vis_env + 'log.txt'
visualizer.log_name = os.path.join(exp_dir, log_name)
os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu_id
num_classes = 68
layer_num = opt.layer_num
net = create_cu_net(neck_size=4,
growth_rate=32,
init_chan_num=128,
class_num=num_classes,
layer_num=layer_num,
order=1,
loss_num=layer_num)
#num1 = get_n_params(net)
#num2 = get_n_trainable_params(net)
#num3 = get_n_conv_params(net)
#print 'number of params: ', num1
#print 'number of trainalbe params: ', num2
#print 'number of conv params: ', num3
#exit()
net = torch.nn.DataParallel(net).cuda()
optimizer = torch.optim.RMSprop(net.parameters(),
lr=opt.lr,
alpha=0.99,
eps=1e-8,
momentum=0,
weight_decay=0)
"""optionally resume from a checkpoint"""
if opt.resume_prefix != '':
# if 'pth' in opt.resume_prefix:
# trunc_index = opt.resume_prefix.index('pth')
# opt.resume_prefix = opt.resume_prefix[0:trunc_index - 1]
checkpoint.save_prefix = os.path.join(exp_dir, opt.resume_prefix)
checkpoint.load_prefix = os.path.join(exp_dir, opt.resume_prefix)[0:-1]
checkpoint.load_checkpoint(net, optimizer, train_history)
else:
checkpoint.save_prefix = exp_dir + '/'
print 'save prefix: ', checkpoint.save_prefix
# model = {'state_dict': net.state_dict()}
# save_path = checkpoint.save_prefix + 'test-model-size.pth.tar'
# torch.save(model, save_path)
# exit()
"""load data"""
train_loader = torch.utils.data.DataLoader(FACE('dataset/face.json',
'/bigdata1/zt53/data/face',
is_train=True),
batch_size=opt.bs,
shuffle=True,
num_workers=opt.nThreads,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(FACE('dataset/face.json',
'/bigdata1/zt53/data/face',
is_train=False),
batch_size=opt.bs,
shuffle=False,
num_workers=opt.nThreads,
pin_memory=True)
"""optimizer"""
#optimizer = torch.optim.SGD( net.parameters(), lr=opt.lr,
# momentum=opt.momentum,
# weight_decay=opt.weight_decay )
# optimizer = torch.optim.RMSprop(net.parameters(), lr=opt.lr, alpha=0.99,
# eps=1e-8, momentum=0, weight_decay=0)
print type(optimizer)
# idx = range(0, 16)
# idx = [e for e in idx if e not in (6, 7, 8, 9, 12, 13)]
# idx = [0, 1, 2, 3, 4, 5, 10, 11, 14, 15]
logger = Logger(os.path.join(opt.exp_dir, opt.exp_id,
opt.resume_prefix + 'face-training-log.txt'),
title='face-training-summary')
logger.set_names(
['Epoch', 'LR', 'Train Loss', 'Val Loss', 'Train RMSE', 'Val RMSE'])
if not opt.is_train:
visualizer.log_path = os.path.join(opt.exp_dir, opt.exp_id,
'val_log.txt')
val_loss, val_rmse, predictions = validate(
val_loader, net, train_history.epoch[-1]['epoch'], visualizer,
num_classes, flip_index)
checkpoint.save_preds(predictions)
return
"""training and validation"""
start_epoch = 0
if opt.resume_prefix != '':
start_epoch = train_history.epoch[-1]['epoch'] + 1
for epoch in range(start_epoch, opt.nEpochs):
AdjustLR(opt, optimizer, epoch)
# # train for one epoch
train_loss, train_rmse = train(train_loader, net, optimizer, epoch,
visualizer, opt)
# evaluate on validation set
val_loss, val_rmse, predictions = validate(val_loader, net, epoch,
visualizer, num_classes,
flip_index)
# visualizer.display_imgpts(imgs, pred_pts, 4)
# exit()
# update training history
e = OrderedDict([('epoch', epoch)])
lr = OrderedDict([('lr', optimizer.param_groups[0]['lr'])])
loss = OrderedDict([('train_loss', train_loss),
('val_loss', val_loss)])
rmse = OrderedDict([('val_rmse', val_rmse)])
train_history.update(e, lr, loss, rmse)
checkpoint.save_checkpoint(net, optimizer, train_history, predictions)
visualizer.plot_train_history_face(train_history)
logger.append([
epoch, optimizer.param_groups[0]['lr'], train_loss, val_loss,
train_rmse, val_rmse
])
logger.close()
def main():
opt = TrainOptions().parse()
train_history = TrainHistoryFace()
checkpoint = Checkpoint()
visualizer = Visualizer(opt)
exp_dir = os.path.join(opt.exp_dir, opt.exp_id)
log_name = opt.vis_env + '_val_log.txt'
visualizer.log_name = os.path.join(exp_dir, log_name)
num_classes = opt.class_num
if not opt.slurm:
os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu_id
layer_num = opt.layer_num
order = opt.order
net = create_cu_net(neck_size=4,
growth_rate=32,
init_chan_num=128,
class_num=num_classes,
layer_num=layer_num,
order=order,
loss_num=layer_num,
use_spatial_transformer=opt.stn,
mlp_tot_layers=opt.mlp_tot_layers,
mlp_hidden_units=opt.mlp_hidden_units,
get_mean_from_mlp=opt.get_mean_from_mlp)
# Load the pre-trained model
saved_wt_file = opt.saved_wt_file
print("Loading weights from " + saved_wt_file)
checkpoint_t = torch.load(saved_wt_file)
state_dict = checkpoint_t['state_dict']
for name, param in state_dict.items():
name = name[7:]
if name not in net.state_dict():
print("=> not load weights '{}'".format(name))
continue
if isinstance(param, Parameter):
param = param.data
net.state_dict()[name].copy_(param)
net = torch.nn.DataParallel(net).cuda() # use multiple GPUs
# Optimizer
if opt.optimizer == "rmsprop":
optimizer = torch.optim.RMSprop(filter(lambda p: p.requires_grad,
net.parameters()),
lr=opt.lr,
alpha=0.99,
eps=1e-8,
momentum=0,
weight_decay=0)
elif opt.optimizer == "adam":
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
net.parameters()),
lr=opt.lr)
else:
print("Unknown Optimizer. Aborting!!!")
sys.exit(0)
print(type(optimizer))
# Optionally resume from a checkpoint
if opt.resume_prefix != '':
# if 'pth' in opt.resume_prefix:
# trunc_index = opt.resume_prefix.index('pth')
# opt.resume_prefix = opt.resume_prefix[0:trunc_index - 1]
checkpoint.save_prefix = os.path.join(exp_dir, opt.resume_prefix)
checkpoint.load_prefix = os.path.join(exp_dir, opt.resume_prefix)[0:-1]
checkpoint.load_checkpoint(net, optimizer, train_history)
else:
checkpoint.save_prefix = exp_dir + '/'
print("Save prefix = {}".format(
checkpoint.save_prefix))
# Load data
json_path = opt.json_path
train_json = opt.train_json
val_json = opt.val_json
print("Path added to each image path in JSON = {}".format(json_path))
print("Train JSON path = {}".format(train_json))
print("Val JSON path = {}".format(val_json))
# This train loader is useless
train_loader = torch.utils.data.DataLoader(FACE(train_json,
json_path,
is_train=True),
batch_size=opt.bs,
shuffle=True,
num_workers=opt.nThreads,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(FACE(val_json,
json_path,
is_train=False),
batch_size=opt.bs,
shuffle=False,
num_workers=opt.nThreads,
pin_memory=True)
if not opt.is_train:
visualizer.log_path = os.path.join(opt.exp_dir, opt.exp_id,
'val_log.txt')
val_loss, val_rmse, predictions = validate(
val_loader, net, train_history.epoch[-1]['epoch'], visualizer,
num_classes, flip_index)
checkpoint.save_preds(predictions)
return
global f_path
global weights_HG
f_path = exp_dir
weights_HG = [float(x) for x in opt.hg_wt.split(",")]
print("Postprocessing applied = {}".format(opt.pp))
if (opt.smax):
print("Scaled softmax used with tau = {}".format(opt.tau))
else:
print("No softmax used")
if opt.is_covariance:
print("Covariance used from the heatmap")
else:
print("Covariance calculated from MLP")
print("Individual Hourglass loss weights")
print(weights_HG)
print("wt_MSE (tradeoff between GLL and MSE in each hourglass)= " +
str(opt.wt_mse))
print(
"wt_gauss_regln (tradeoff between GLL and Gaussian Regularisation in each hourglass)= "
+ str(opt.wt_gauss_regln))
# Optionally resume from a checkpoint
start_epoch = 0
if opt.resume_prefix != '':
start_epoch = train_history.epoch[-1]['epoch'] + 1
# Training and validation
start_epoch = 0
if opt.resume_prefix != '':
start_epoch = train_history.epoch[-1]['epoch'] + 1
train_loss_orig_epoch = []
train_loss_gau_t1_epoch = []
train_loss_gau_t2_epoch = []
train_nme_orig_epoch = []
train_nme_gau_epoch = []
train_nme_new_epoch = []
val_loss_orig_epoch = []
val_loss_gau_t1_epoch = []
val_loss_gau_t2_epoch = []
val_nme_orig_epoch = []
val_nme_gau_epoch = []
val_nme_new_epoch = []
for epoch in range(1):
# Evaluate on validation set
val_loss, val_loss_mse, val_loss_gau_t1, val_loss_gau_t2, val_rmse_orig, val_rmse_gau, val_rmse_new_box, predictions = validate(
val_loader, net, epoch, visualizer, opt, num_classes, flip_index)
val_loss_orig_epoch.append(val_loss_mse)
val_loss_gau_t1_epoch.append(val_loss_gau_t1)
val_loss_gau_t2_epoch.append(val_loss_gau_t2)
val_nme_orig_epoch.append(val_rmse_orig)
val_nme_gau_epoch.append(val_rmse_gau)
import numpy as np
from PIL import Image
import torch
import torch.utils.data as data
import torchvision.transforms as transforms
import json
import random
sys.path.insert(0, os.getcwd())
from data.face_bbx import FACE
from plot.CommonPlottingOperations import *
from pylib import HumanPts, FaceAug, HumanAug, FacePts
index = 0
val_loader = torch.utils.data.DataLoader(FACE("dataset/cofw_68_val.json",
".",
is_train=False),
batch_size=1,
shuffle=True,
num_workers=1,
pin_memory=True)
# Visualize some images
for i, (img, _, points, _, _, _, _, s, visible_multiclass,
_) in enumerate(val_loader):
print(i)
image = img[index].numpy()
pts = points[index].numpy()
vis = visible_multiclass.clone().numpy()
#vis[vis > 1] = 1
import os, sys
import numpy as np
from PIL import Image
import torch
import torch.utils.data as data
import torchvision.transforms as transforms
import json
import random
sys.path.insert(0, os.getcwd())
from data.face_bbx import FACE
from plot.CommonPlottingOperations import *
from pylib import HumanPts, FaceAug, HumanAug, FacePts
index = 0
val_loader = torch.utils.data.DataLoader(FACE(
"dataset/all_300Wtest_train.json", ".", is_train=True),
batch_size=1,
shuffle=False,
num_workers=1,
pin_memory=True)
# Visualize some images
for i, (img, _, points, _, s) in enumerate(val_loader):
print(i)
image = img[index].numpy()
pts = points[index].numpy()
plt.figure(figsize=(16, 8))
plt.subplot(121)
plt.imshow(swap_channels(image))
plt_pts(plt, pts)
import numpy as np
from PIL import Image
import torch
import torch.utils.data as data
import torchvision.transforms as transforms
import json
import random
sys.path.insert(0, os.getcwd())
from data.face_bbx import FACE
from plot.CommonPlottingOperations import *
from pylib import HumanPts, FaceAug, HumanAug, FacePts
index = 0
val_loader = torch.utils.data.DataLoader(FACE("dataset/300W_LP_train.json",
".",
is_train=True,
keep_pts_inside=True),
batch_size=1,
shuffle=True,
num_workers=1,
pin_memory=True)
# Visualize some images
for i, (img, _, points, _, s, visible_multiclass, _) in enumerate(val_loader):
print(i)
image = img[index].numpy()
pts = points[index].numpy()
vis = visible_multiclass.clone().numpy()
print(vis.shape)
vis[vis > 1] = 1