def visualize_stn(original, transformed, before_stn, after_stn, filenames,
target):
target = os.path.join(target, "stnviz")
mkdir_if_not_exists(target)
transformed = transformed.detach().cpu().numpy().transpose(0, 2, 3, 1)
before_stn = before_stn.detach().cpu().numpy().transpose(0, 2, 3, 1)
after_stn = after_stn.detach().cpu().numpy().transpose(0, 2, 3, 1)
for orig, trans, before, after, filename in zip(original, transformed,
before_stn, after_stn,
filenames):
origimg = cv2.cvtColor(orig, cv2.COLOR_RGB2BGR)
before = cv2.cvtColor(
FaceLandmarksTrainingData.undo_normalization(before),
cv2.COLOR_RGB2BGR)
after = cv2.cvtColor(
FaceLandmarksTrainingData.undo_normalization(after),
cv2.COLOR_RGB2BGR)
trans = cv2.cvtColor(
FaceLandmarksTrainingData.undo_normalization(trans),
cv2.COLOR_RGB2BGR)
comb = np.hstack((origimg, trans, before, after))
fn = os.path.join(target, "%s.png" % filename)
cv2.imwrite(fn, comb)
def plot_curve(filename, key, split, target):
color_map = {
"e49": "red",
"h49": "blue",
"e68": "green",
"h68": "purple",
"e": None,
"h": None,
"_49": None,
"_68": None,
"_all": None
}
with open(filename) as csvfile:
r = csv.DictReader(csvfile)
data = defaultdict(list)
for row in r:
for col, val in row.items():
data[col].append(val)
x = [int(x) for x in data[key]]
y = [float(v) for v in data[split]]
depths = [int(x) for x in data["hg_depth"]]
x1 = [x[i] for i in range(len(x)) if depths[i] < 5]
x2 = [x[i] for i in range(len(x)) if depths[i] == 5]
y1 = [y[i] for i in range(len(y)) if depths[i] < 5]
y2 = [y[i] for i in range(len(y)) if depths[i] == 5]
fig = plt.gcf()
fig.set_size_inches(5.5, 3.5) #width,height
ax = plt.gca()
ax.yaxis.set_major_formatter(FormatStrFormatter('%0.2f'))
plt.scatter(x2, y2, s=10, c="black", label="hg_depth = 5")
plt.scatter(x1, y1, s=10, c=color_map[split], label="hg_depth < 5")
plt.xlabel(key.replace("_", " "))
plt.ylabel("IOD normalized RMSE in %")
target_ratio = 1
current_aspect = sub(*ax.get_ylim()) / sub(*ax.get_xlim())
#print(current_aspect, 1/current_aspect)
#print(sub(*ax.get_ylim()) / sub(*ax.get_xlim()))
ax.set_aspect(target_ratio / current_aspect)
#print(sub(*ax.get_ylim()) / sub(*ax.get_xlim()))
plt.legend()
targetdir = os.path.join(target, key)
mkdir_if_not_exists(targetdir)
plt.savefig(os.path.join(targetdir, "%s_%s.png" % (key, split)),
bbox_inches='tight')
#plt.show()
plt.clf()
def __init__(self, config, output_dir, gpu_id, sub_gpu_id, data, gridsearch):
super(ModelTrainer, self).__init__(gpu_id, sub_gpu_id)
if isinstance(config, dict):
self.config = config
else:
self.config = load_config(config, verbose=False)
self.data = data
self.output_dir = output_dir
mkdir_if_not_exists(self.output_dir)
self.model_dir = os.path.join(self.output_dir, "models")
mkdir_if_not_exists(self.model_dir)
self.result_dir = os.path.join(self.output_dir, "results")
mkdir_if_not_exists(self.result_dir)
self.plot_dir = os.path.join(self.output_dir, "plots")
mkdir_if_not_exists(self.plot_dir)
self.is_gridsearch = gridsearch
if self.is_gridsearch:
self.loss_log = []
self.l2d_log = []
self.gts = []
self.hg_coords_log = []
def plot_curve(filename, key, split, target):
color_map = {
"e49" : "red",
"h49" : "blue",
"e68" : "green",
"h68" : "purple",
"e": None,
"h": None,
"_49": None,
"_68": None,
"_all": None
}
with open(filename) as csvfile:
r = csv.DictReader(csvfile)
data = defaultdict(list)
for row in r:
for col,val in row.items():
data[col].append(val)
x = data[key]
y = [float(v) for v in data[split]]
ax = plt.gca()
ax.yaxis.set_major_formatter(FormatStrFormatter('%0.3f'))
plt.plot(x, y, '--ro', c=color_map[split])
fig = plt.gcf()
fig.set_size_inches(4,2.5) #width,height
plt.xticks(x)
plt.xlabel(key)
plt.ylabel("IOD normalized RMSE in %")
plt.grid()
target_ratio = 1
current_aspect = sub(*ax.get_ylim()) / sub(*ax.get_xlim())
#print(current_aspect, 1/current_aspect)
#print(sub(*ax.get_ylim()) / sub(*ax.get_xlim()))
ax.set_aspect(target_ratio/current_aspect)
#print(sub(*ax.get_ylim()) / sub(*ax.get_xlim()))
#plt.legend()
targetdir = os.path.join(target, key)
mkdir_if_not_exists(targetdir)
plt.savefig(os.path.join(targetdir, "%s_%s.png" % (key, split)), bbox_inches='tight')
#plt.show()
plt.clf()
def plot(untransformed, transformation, target):
mkdir_if_not_exists(target)
transformed = PDM.transform(untransformed, transformation)
for i in range(len(transformed)):
fig = pyplot.figure()
ax = Axes3D(fig)
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")
ax.scatter(transformed[i, :, 0], transformed[i, :, 1] * -1,
transformed[i, :, 2])
ax.view_init(elev=70, azim=-90)
pyplot.savefig(os.path.join(target, "%d.png" % i))
pyplot.close()
def __init__(self, config, output_dir, gpu_id, sub_gpu_id, data,
gridsearch):
super(ModelTrainer, self).__init__(gpu_id, sub_gpu_id)
if isinstance(config, dict):
self.config = config
else:
self.config = load_config(config, verbose=False)
self.data = data
self.output_dir = output_dir
mkdir_if_not_exists(self.output_dir)
self.model_dir = os.path.join(self.output_dir, "models")
mkdir_if_not_exists(self.model_dir)
self.result_dir = os.path.join(self.output_dir, "results")
mkdir_if_not_exists(self.result_dir)
self.plot_dir = os.path.join(self.output_dir, "plots")
mkdir_if_not_exists(self.plot_dir)
self.is_gridsearch = gridsearch
if self.is_gridsearch:
self.loss_log = {"train": [], "test": []}
self.metrics_log = {}
self.category_calculator = {
"e49":
lambda metrics: metrics["e49"],
"h49":
lambda metrics: metrics["h49"],
"e68":
lambda metrics: metrics["e68"],
"h68":
lambda metrics: metrics["h68"],
"49":
lambda metrics: (metrics["e49"] + metrics["h49"]) / 2,
"68":
lambda metrics: (metrics["e68"] + metrics["h68"]) / 2,
"e":
lambda metrics: (metrics["e49"] + metrics["e68"]) / 2,
"h":
lambda metrics: (metrics["h49"] + metrics["h68"]) / 2,
"all":
lambda metrics: (metrics["e49"] + metrics["h49"] + metrics[
"e68"] + metrics["h68"]) / 4
}
self.categories = self.category_calculator.keys()
self.best_epoch = {k: 0 for k in self.categories}
self.lowest_error = {k: np.Inf for k in self.categories}
def __init__(self, config, output_dir, gpu_id, sub_gpu_id, data,
gridsearch):
super(ModelTrainer, self).__init__(gpu_id, sub_gpu_id)
self.config = load_config(config, verbose=False)
self.data = data
self.output_dir = output_dir
mkdir_if_not_exists(self.output_dir)
self.model_dir = os.path.join(self.output_dir, "models")
mkdir_if_not_exists(self.model_dir)
self.result_dir = os.path.join(self.output_dir, "results")
mkdir_if_not_exists(self.result_dir)
self.plot_dir = os.path.join(self.output_dir, "plots")
mkdir_if_not_exists(self.plot_dir)
self.is_gridsearch = gridsearch # TODO use this to print verbose
help="Path to pretrained hourglass (.torch)")
parser.add_argument("pdm",
type=str,
help="Path to pretrained PDM (.torch)")
parser.add_argument("data",
type=str,
help="all_data_valid_w_profile_pts.h5")
parser.add_argument("target", type=str, help="Where to store")
parser.add_argument("--gpu", type=int, default=0, help="GPU ID")
parser.add_argument("--n",
type=int,
default=10,
help="N biggest improvements")
args = parser.parse_args()
mkdir_if_not_exists(args.target)
gpu_id = args.gpu
location = 'cpu' if gpu_id < 0 else "cuda:%d" % gpu_id
hg, hg_config = load_hg(args.hg, location)
pdm, pdm_config = load_pdm(args.pdm, location)
pdm.print_losses = False
#pdm.test_epochs = 100
torch.autograd.set_detect_anomaly(True) # This makes debugging much easier
if location is not 'cpu':
torch.cuda.set_device(torch.device(location))
normMean, normStd = FaceLandmarksTrainingData.TRAIN_MEAN, FaceLandmarksTrainingData.TRAIN_STD
def run(self):
torch.autograd.set_detect_anomaly(
True) # This makes debugging much easier
make_deterministic(self.config['random_seed'])
encoders = None
if self.config["encoder"]:
# This assumes that an encoder has already been trained for the PDM
# Example: pdm path is my/dir/models/pdm_4.torch
# Then the encoder is loaded from my/dir/encoders/encoder_4.torch (if it does not exists, the code crashes)
pdm_filename = os.path.basename(self.config["pdm"])
if "final" in pdm_filename:
pdm_id = int(pdm_filename.split(".")[0].split("_")[-1])
else:
pdm_id = int(pdm_filename.split("_")[0])
encoders = {
49:
os.path.join(
os.path.dirname(os.path.dirname(self.config["pdm"])),
"encoders", "encoder_49_%d.torch" % pdm_id),
68:
os.path.join(
os.path.dirname(os.path.dirname(self.config["pdm"])),
"encoders", "encoder_68_%d.torch" % pdm_id),
}
if not self.is_gridsearch:
print("encoder", encoders)
if "prediction_target" in self.config and self.config[
"prediction_target"] is not None:
pred_target = self.config["prediction_target"]
pred_target_dir = os.path.dirname(pred_target)
mkdir_if_not_exists(pred_target_dir)
else:
pred_target = None
success = False
tries = 0
maxtries = 75
while not success:
tries += 1
try:
res, hg_config, pdm_config = load_and_run(
hg_src=self.config["hg"],
pdm_src=self.config["pdm"],
data_src=self.data,
gpu_id=self.gpu_id,
random_seed=self.config["random_seed"],
pdm_configurator=self.configure_pdm,
verbose=not self.is_gridsearch,
var_thresh=self.config["variance_threshold"],
encoders=encoders)
success = True
except RuntimeError as e:
txt = str(e)
if "out of memory" in txt:
if tries <= maxtries:
waittime = tries * random.randint(1, 5)
print(
"ERROR! There was a OOM error, wait %d seconds and try again. Try nr. %d"
% (waittime, tries))
time.sleep(waittime)
else:
print("ERROR! maxtries (%d) exceeded" % maxtries)
raise e
else:
raise e
results = {
"hg_easy49":
res["hg"]["easy49"],
"hg_hard49":
res["hg"]["hard49"],
"hg_easy68":
res["hg"]["easy68"],
"hg_hard68":
res["hg"]["hard68"],
"pdm_easy49":
res["pdm"]["easy49"],
"pdm_hard49":
res["pdm"]["hard49"],
"pdm_easy68":
res["pdm"]["easy68"],
"pdm_hard68":
res["pdm"]["hard68"],
"pdm_encoder_easy49":
res["pdm_encoder"]["easy49"],
"pdm_encoder_hard49":
res["pdm_encoder"]["hard49"],
"pdm_encoder_easy68":
res["pdm_encoder"]["easy68"],
"pdm_encoder_hard68":
res["pdm_encoder"]["hard68"],
"easy49_factor":
res["hg"]["easy49"] / res["pdm"]["easy49"],
"hard49_factor":
res["hg"]["hard49"] / res["pdm"]["hard49"],
"easy68_factor":
res["hg"]["easy68"] / res["pdm"]["easy68"],
"hard68_factor":
res["hg"]["hard68"] / res["pdm"]["hard68"],
"enc_easy49_factor":
res["hg"]["easy49"] /
res["pdm_encoder"]["easy49"] if self.config["encoder"] else 0.0,
"enc_hard49_factor":
res["hg"]["hard49"] /
res["pdm_encoder"]["hard49"] if self.config["encoder"] else 0.0,
"enc_easy68_factor":
res["hg"]["easy68"] /
res["pdm_encoder"]["easy68"] if self.config["encoder"] else 0.0,
"enc_hard68_factor":
res["hg"]["hard68"] /
res["pdm_encoder"]["hard68"] if self.config["encoder"] else 0.0
}
print(
"Config: %d | factor e49: %0.4f | factor h49: %0.4f | factor e68: %0.4f | factor h68: %0.4f"
% (self.config["config_id"], results["easy49_factor"],
results["hard49_factor"], results["easy68_factor"],
results["hard68_factor"]))
if self.is_gridsearch:
return {**self.config, **results}
else:
for k, v in results.items():
print(k, v)
if pred_target:
output = {
"meta": {
"hg_model": self.config["hg"],
"pdm_model": self.config["pdm"],
"hg_config": hg_config,
"pdm_config": pdm_config,
"gapsearch_config": self.config
},
"results": results,
"predictions": {
"easy": {
"gt":
res["gt"]["easy"].cpu().detach().numpy().tolist(),
"pdm_pred":
res["pdm_pred"]
["easy"].cpu().detach().numpy().tolist(),
"hg_pred":
res["hg_pred"]["easy"].cpu().detach().numpy().tolist(),
"pdm_3d":
res["pdm_3d"]["easy"].cpu().detach().numpy().tolist()
},
"hard": {
"gt":
res["gt"]["hard"].cpu().detach().numpy().tolist(),
"pdm_pred":
res["pdm_pred"]
["hard"].cpu().detach().numpy().tolist(),
"hg_pred":
res["hg_pred"]["hard"].cpu().detach().numpy().tolist(),
"pdm_3d":
res["pdm_3d"]["hard"].cpu().detach().numpy().tolist()
}
}
}
if "pdm_encoder_pred" in res:
output["predictions"]["easy"]["pdm_encoder_pred"] = res[
"pdm_encoder_pred"]["easy"].cpu().detach().numpy().tolist(
)
output["predictions"]["hard"]["pdm_encoder_pred"] = res[
"pdm_encoder_pred"]["hard"].cpu().detach().numpy().tolist(
)
json.dump(output, open(pred_target, "w"), indent=2)
print("Predictions written to", pred_target)
def make_video(pdm_path,
hg_results,
target,
include=[
"gt", "hg", "pdm", "connection_gt", "connection_hg",
"confidence"
],
epoch_frequency=1,
encoder=None):
hg_coords, hg_coords_and_conf, gt, l2d, history = run_pdm(pdm_path,
hg_results,
encoder=encoder,
history=True)
if "pdm" not in include:
print(
"Warning: pdm output is not included in video, so each frame will be the same"
)
gt = gt.detach().cpu().numpy()
hg_out = hg_coords.detach().cpu().numpy()
losses = [
mean_squared_error(hg_out[i], history[-1]["l2d"][i])
for i in range(len(hg_out))
]
sortedlosses = list(sorted(losses, reverse=True))[:2]
interesting_losses_idx = [
i for i in range(len(losses)) if losses[i] in sortedlosses
]
pyplot.figure(figsize=(8, 8))
for sample_index in tqdm(interesting_losses_idx):
tqdm.write(str(sample_index))
evolution = np.array([
history[epoch]["l2d"][sample_index]
for epoch in range(len(history))
][::epoch_frequency])
mkdir_if_not_exists(os.path.join(target, "%d" % sample_index),
verbose=False)
framenr = 0
for epoch, pred in tqdm(list(enumerate(evolution))):
epoch *= epoch_frequency
mse = mean_squared_error(hg_out[sample_index], pred)
if "hg" in include:
pyplot.scatter(hg_out[sample_index, :, 0],
hg_out[sample_index, :, 1] * -1,
s=6,
label="HG",
color="green")
if "gt" in include:
pyplot.scatter(gt[sample_index, :, 0],
gt[sample_index, :, 1] * -1,
s=6,
label="GT",
color="orange")
if "pdm" in include:
pyplot.scatter(pred[:, 0],
pred[:, 1] * -1,
s=8,
label="PDM",
color="blue")
for i in range(68):
if "connection_gt" in include:
x1, y1 = [gt[sample_index, i, 0], pred[i, 0]], [
gt[sample_index, i, 1] * -1, pred[i, 1] * -1
]
pyplot.plot(x1,
y1,
marker=None,
color='orange',
linewidth=1,
linestyle='dashed')
if "connection_hg" in include:
x1, y1 = [hg_out[sample_index, i, 0], pred[i, 0]], [
hg_out[sample_index, i, 1] * -1, pred[i, 1] * -1
]
pyplot.plot(x1,
y1,
marker=None,
color='green',
linewidth=1,
linestyle='dashed')
if "pdm" in include and "confidence" in include:
pyplot.annotate(
"%0.1f" % sum(hg_coords_and_conf[sample_index, i, 2:]),
(pred[i, 0], pred[i, 1] * -1),
size=6)
pyplot.xlim([-1.0, 1.0])
pyplot.ylim([-1.0, 1.0])
pyplot.legend()
pyplot.figtext(0.5,
0.01,
"Epoch %04d - Loss %0.4f" % (epoch, mse),
wrap=True,
horizontalalignment='center',
fontsize=12)
pyplot.savefig(
os.path.join(target,
"%d/frame_%d.png" % (sample_index, framenr)))
pyplot.clf()
framenr += 1
if __name__ == "__main__":
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument(
"zs_and_nr_from_pdm",
type=str,
help="json file with zs, nr and coords after PDM training")
parser.add_argument("target",
type=str,
help="Where to store encoder (.torch)")
parser.add_argument("--is_49lm", default=False, action="store_true")
parser.add_argument("--gpu", type=int, default=-1, help="GPU id, -1 = CPU")
args = parser.parse_args()
device = torch.device("cuda:%d" % args.gpu if args.gpu >= 0 else "cpu")
#print(device)
data = json.load(open(args.zs_and_nr_from_pdm, "r"))
enc = train(data["train"], data["test"], is_49lm=args.is_49lm)
test_loss = test(enc, data["test"], is_49lm=args.is_49lm)
print("Final test loss", test_loss)
mkdir_if_not_exists(os.path.dirname(args.target))
torch.save(
{
"state_dict": enc.state_dict(),
"zs_size": enc.zs_size,
"nr_size": enc.nr_size
}, open(args.target, "wb"))
def plot(src, dst, dataset, split):
mkdir_if_not_exists(dst)
with open(src, "r") as f:
data = json.load(f)["predictions"][split]
enc_pred = torch.tensor(data["encoder_pred"])
hg_pred = torch.tensor(data["hg_pred"])
gt = torch.tensor(data["gt"])
n_lm = gt.shape[1]
#if n_lm == 49:
# gt = gt[:,NO_OUTLINE_MASK,:]
with h5py.File(dataset) as d:
if split == "easy":
dataset = FaceLandmarksEasyTestData(d, n_lm=n_lm)
else:
dataset = FaceLandmarksHardTestData(d, n_lm=n_lm)
for i in tqdm(range(len(dataset))):
_gt = gt[i]
img = cv2.cvtColor(dataset[i]["original_image"], cv2.COLOR_RGB2BGR)
_enc_pred = enc_pred[i]
_hg_pred = hg_pred[i]
cv2.imwrite(os.path.join(dst, "%d_raw.png" % i), img)
gt_img = draw_landmarks(img, _gt, color=COLOR_GT, size=2)
cv2.imwrite(os.path.join(dst, "%d_gt.png" % i), gt_img)
hg_img = draw_landmarks(img, _hg_pred, color=COLOR_HG, size=2)
cv2.imwrite(os.path.join(dst, "%d_hg.png" % i), hg_img)
hg_gt_img = draw_landmarks(gt_img, _hg_pred, color=COLOR_HG, size=2)
cv2.imwrite(os.path.join(dst, "%d_hg_gt.png" % i), hg_gt_img)
enc_img = draw_landmarks(img, _enc_pred, color=COLOR_ENCODER, size=2)
cv2.imwrite(os.path.join(dst, "%d_encoder.png" % i), enc_img)
enc_gt_img = draw_landmarks(gt_img,
_enc_pred,
color=COLOR_ENCODER,
size=2)
cv2.imwrite(os.path.join(dst, "%d_enc_gt.png" % i), enc_gt_img)
enc_hg_img = draw_landmarks(hg_img,
_enc_pred,
color=COLOR_ENCODER,
size=2)
cv2.imwrite(os.path.join(dst, "%d_enc_hg.png" % i), enc_hg_img)
res = {
"gt": _gt.cpu().detach().numpy().tolist(),
"encoder": _enc_pred.cpu().detach().numpy().tolist(),
"hg": _hg_pred.cpu().detach().numpy().tolist()
}
json.dump(res,
open(os.path.join(dst, "%d_predictions.json" % i), "w"),
indent=2)
def run(model,
src_300w,
src_menpo,
target,
gpu=None,
override_norm_params=False,
bs_factor=1):
location = 'cpu' if gpu is None else "cuda:%d" % gpu
if location is not 'cpu':
# This fixes the problem that pytorch is always allocating memory on GPU 0 even if this is not included
# in the list of GPUs to use
torch.cuda.set_device(torch.device(location))
# cudnn.benchmark improves training speed when input sizes do not change
# https://discuss.pytorch.org/t/what-does-torch-backends-cudnn-benchmark-do/5936
# It selects the best algorithms as the training iterates over the dataset
#cudnn.benchmark = True # disable for deterministic behavior
print("Location: ", location)
data = torch.load(model, map_location=location)
state_dict = data['state_dict']
config = data['config']
n_lm = config["n_lm"]
if n_lm == 49:
print("WARNING! THIS IS A 49 LM model!!!!", n_lm)
num_workers = multiprocessing.cpu_count()
batch_size = config[
'batch_size'] * bs_factor if gpu is not None else num_workers
pin_memory = gpu is not None
print("Workers: ", num_workers)
print("Batchsize: ", batch_size)
net = ModelTrainer.create_net(config, verbose=False)
net.load_state_dict(state_dict)
net.eval()
net.to(location)
mkdir_if_not_exists(os.path.dirname(target))
normMean, normStd = FaceLandmarksTrainingData.TRAIN_MEAN, FaceLandmarksTrainingData.TRAIN_STD
if override_norm_params:
normMean = tuple(
np.array([133.0255852472676, 101.61684197664563, 87.4134193236219])
/ 255.0)
normStd = tuple(
np.array([71.91047346327116, 62.94368776888253, 61.56865329427311])
/ 255.0)
normTransform = transforms.Normalize(normMean, normStd)
transform = transforms.Compose([
ImageTransform(transforms.ToPILImage()),
ImageTransform(transforms.ToTensor()),
ImageTransform(normTransform)
])
with h5py.File(src_300w, 'r') as f:
print("Run on easy")
easy_d = FaceLandmarksEasyTestData(f, transform=transform, n_lm=n_lm)
easy_loader = DataLoader(dataset=easy_d,
shuffle=False,
num_workers=num_workers,
pin_memory=pin_memory,
batch_size=batch_size)
easy_results = evaluate_split(net,
easy_loader,
location=location,
n_lm=n_lm)
print("Run on hard")
hard_d = FaceLandmarksHardTestData(f, transform=transform, n_lm=n_lm)
hard_loader = DataLoader(dataset=hard_d,
shuffle=False,
num_workers=num_workers,
pin_memory=pin_memory,
batch_size=batch_size)
hard_results = evaluate_split(net,
hard_loader,
location=location,
n_lm=n_lm)
print("Run on train")
train = FaceLandmarksTrainingData(f, transform=transform, n_lm=n_lm)
train_loader = DataLoader(dataset=train,
shuffle=False,
num_workers=num_workers,
pin_memory=pin_memory,
batch_size=batch_size)
train_results = evaluate_split(net,
train_loader,
location=location,
n_lm=n_lm)
with h5py.File(src_menpo, "r") as f:
print("Run on menpo")
menpo = Menpo(f, transform=transform, n_lm=n_lm)
menpo_loader = DataLoader(dataset=menpo,
shuffle=False,
num_workers=num_workers,
pin_memory=pin_memory,
batch_size=batch_size)
menpo_results = evaluate_split(net,
menpo_loader,
location=location,
n_lm=n_lm)
res = {
"easy": easy_results,
"hard": hard_results,
"train": train_results,
"menpo": menpo_results,
"model_src": model,
"config": config
}
if target is not None:
json.dump(res, open(target, "w"))
else:
return res
def plot(src, dst, dataset, split):
mkdir_if_not_exists(dst)
with open(src, "r") as f:
data = json.load(f)["predictions"][split]
pdm_pred = torch.tensor(data["pdm_pred"])
hg_pred = torch.tensor(data["hg_pred"])
if "pdm_encoder_pred" in data:
pdm_enc_pred = torch.tensor(data["pdm_encoder_pred"])
else:
pdm_enc_pred = None
pdm_3d = torch.tensor(data["pdm_3d"])
gt = torch.tensor(data["gt"])
n_lm = pdm_pred.shape[1]
if n_lm == 49:
gt = gt[:, NO_OUTLINE_MASK, :]
with h5py.File(dataset) as d:
if split == "easy":
dataset = FaceLandmarksEasyTestData(d, n_lm=n_lm)
else:
dataset = FaceLandmarksHardTestData(d, n_lm=n_lm)
for i in tqdm(range(len(dataset))):
_gt = gt[i]
img = cv2.cvtColor(dataset[i]["original_image"], cv2.COLOR_RGB2BGR)
_pdm_pred = pdm_pred[i]
_hg_pred = hg_pred[i]
if pdm_enc_pred is not None:
_pdm_init_pred = pdm_enc_pred[i]
_pdm_3d = pdm_3d[i]
cv2.imwrite(os.path.join(dst, "%d_raw.png" % i), img)
gt_img = draw_landmarks(img, _gt, color=COLOR_GT, size=2)
cv2.imwrite(os.path.join(dst, "%d_gt.png" % i), gt_img)
pdm_img = draw_landmarks(img, _pdm_pred, color=COLOR_PDM, size=2)
cv2.imwrite(os.path.join(dst, "%d_pdm.png" % i), pdm_img)
pdm_gt_img = draw_landmarks(gt_img, _pdm_pred, color=COLOR_PDM, size=2)
cv2.imwrite(os.path.join(dst, "%d_pdm_gt.png" % i), pdm_gt_img)
hg_img = draw_landmarks(img, _hg_pred, color=COLOR_HG, size=2)
cv2.imwrite(os.path.join(dst, "%d_hg.png" % i), hg_img)
hg_gt_img = draw_landmarks(gt_img, _hg_pred, color=COLOR_HG, size=2)
cv2.imwrite(os.path.join(dst, "%d_hg_gt.png" % i), hg_gt_img)
hg_pdm_img = draw_landmarks(hg_img, _pdm_pred, color=COLOR_PDM, size=2)
cv2.imwrite(os.path.join(dst, "%d_hg_pdm.png" % i), hg_pdm_img)
hg_pdm_gt_img = draw_landmarks(hg_gt_img,
_pdm_pred,
color=COLOR_PDM,
size=2)
cv2.imwrite(os.path.join(dst, "%d_hg_pdm_gt.png" % i), hg_pdm_gt_img)
if pdm_enc_pred is not None:
pdm_init_img = draw_landmarks(img,
_pdm_init_pred,
color=COLOR_PDM_INIT,
size=2)
cv2.imwrite(os.path.join(dst, "%d_pdm_init.png" % i), pdm_init_img)
pdm_init_pdm_img = draw_landmarks(pdm_init_img,
_pdm_pred,
color=COLOR_PDM,
size=2)
cv2.imwrite(os.path.join(dst, "%d_pdm_init_and_pdm.png" % i),
pdm_init_pdm_img)
pdm_init_gt_img = draw_landmarks(gt_img,
_pdm_init_pred,
color=COLOR_PDM_INIT,
size=2)
cv2.imwrite(os.path.join(dst, "%d_pdm_init_and_gt.png" % i),
pdm_init_gt_img)
fig = pyplot.figure()
ax = Axes3D(fig)
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")
ax.scatter(_pdm_3d[:, 0], _pdm_3d[:, 1] * -1, _pdm_3d[:, 2])
ax.view_init(elev=70, azim=-90)
pyplot.savefig(os.path.join(dst, "%d_3d.png" % i))
pyplot.close()
res = {
"gt": _gt.cpu().detach().numpy().tolist(),
"pdm_2d": _pdm_pred.cpu().detach().numpy().tolist(),
"pdm_3d": _pdm_3d.cpu().detach().numpy().tolist(),
"hg": _hg_pred.cpu().detach().numpy().tolist()
}
if pdm_enc_pred is not None:
res["pdm_init_only"] = _pdm_init_pred.cpu().detach().numpy(
).tolist()
json.dump(res,
open(os.path.join(dst, "%d_predictions.json" % i), "w"),
indent=2)
def plot_best_improvements(hg_pred, pdm_pred, gt, images, target, n=10):
mkdir_if_not_exists(target)
improvements = get_improvements(hg_pred, pdm_pred, gt)[-n:]
all_imgs = []
all_imgs_desc = []
for rank, (i, hg_err, pdm_err, ratio) in enumerate(improvements[::-1]):
print(target, i, hg_err, pdm_err, ratio)
gt_img = draw_landmarks(cv2.cvtColor(images[i], cv2.COLOR_RGB2BGR),
gt[i],
color=COLOR_GT,
size=1)
cv2.imwrite(
os.path.join(target,
"%d_%d_gt_%0.4f_plain.png" % (rank, i, ratio)),
gt_img)
gt_img_desc = add_description(gt_img, "GT")
cv2.imwrite(
os.path.join(target, "%d_%d_gt_%0.4f_desc.png" % (rank, i, ratio)),
gt_img_desc)
hg_img = draw_landmarks(cv2.cvtColor(images[i], cv2.COLOR_RGB2BGR),
hg_pred[i],
color=COLOR_HG,
size=1)
cv2.imwrite(
os.path.join(target,
"%d_%d_hg_%0.4f_plain.png" % (rank, i, ratio)),
hg_img)
hg_img_desc = add_description(hg_img, "HG %0.6f" % hg_err)
cv2.imwrite(
os.path.join(target, "%d_%d_hg_%0.4f_desc.png" % (rank, i, ratio)),
hg_img_desc)
pdm_img = draw_landmarks(cv2.cvtColor(images[i], cv2.COLOR_RGB2BGR),
pdm_pred[i],
color=COLOR_PDM,
size=1)
cv2.imwrite(
os.path.join(target,
"%d_%d_pdm_%0.4f_plain.png" % (rank, i, ratio)),
pdm_img)
pdm_img_desc = add_description(pdm_img, "PDM %0.6f" % pdm_err)
cv2.imwrite(
os.path.join(target,
"%d_%d_pdm_%0.4f_desc.png" % (rank, i, ratio)),
pdm_img_desc)
gal_input = [add_border(x, 5) for x in [gt_img, hg_img, pdm_img]]
cur_all = gallery(np.array(gal_input), 1)
cv2.imwrite(
os.path.join(target, "%d_%d_all_%0.4f.png" % (rank, i, ratio)),
cur_all)
gal_input_desc = [
add_border(x, 5) for x in [gt_img_desc, hg_img_desc, pdm_img_desc]
]
cur_all_desc = gallery(np.array(gal_input_desc), 1)
cv2.imwrite(
os.path.join(target,
"%d_%d_all_desc_%0.4f.png" % (rank, i, ratio)),
cur_all_desc)
all_imgs.append(cur_all)
all_imgs_desc.append(
add_description(cur_all_desc, " impr. %0.4f" % ratio))
all_imgs = gallery(np.array(all_imgs), len(all_imgs))
cv2.imwrite(os.path.join(target, "gallery.png"), all_imgs)
all_imgs_desc = gallery(np.array(all_imgs_desc), len(all_imgs_desc))
cv2.imwrite(os.path.join(target, "gallery_desc.png"), all_imgs_desc)
sample_pdm = np.array(pdm_pred[epoch][sample_index])
sample_gt = np.array(gt[sample_index])
pyplot.scatter(sample_gt[:, 0],
sample_gt[:, 1] * -1,
s=3,
label="GT",
color="orange")
pyplot.scatter(sample_hg[:, 0],
sample_hg[:, 1] * -1,
s=3,
label="HG",
color="green")
#pyplot.scatter(sample_pdm[:, 0], sample_hg[:, 1] * -1, s=6, label="PDM", color="blue")
pyplot.xlim([-1.0, 1.0])
pyplot.ylim([-1.0, 1.0])
mkdir_if_not_exists(os.path.join(target, "%d" % sample_index))
pyplot.legend()
pyplot.figtext(0.5,
0.01,
"Epoch %04d" % (epoch, ),
wrap=True,
horizontalalignment='center',
fontsize=12)
pyplot.savefig(
os.path.join(target, "%d/frame_%d.png" % (sample_index, epoch)))
pyplot.clf()
args = parser.parse_args()
with h5py.File(args.source, 'r') as f:
if args.is_menpo:
splits = [Menpo(f)]
else:
splits = [FaceLandmarksTrainingData(f), FaceLandmarksEasyTestData(f), FaceLandmarksHardTestData(f)]
stats = defaultdict(int)
for data in splits:
split = data.split
print('Split: %s' % split)
directory = os.path.join(args.target, split)
mkdir_if_not_exists(directory)
for sample in tqdm(list(data)):
image = sample['original_image'][:,:,::-1].copy() # RGB o BGR
if args.draw_landmarks:
image = draw_landmarks(image, sample['landmarks'])
target_file = os.path.join(directory, sample['filename'])
cv2.imwrite(target_file, image)
stats[split] += 1
print("Dataset stats")
for split, count in stats.items():
print("%s : %d images" % (split, count))
def visualize_split(net, split, target, location, landmarks_in_heatmaps):
mkdir_if_not_exists(target)
evaluated_images = []
i = 0
curidx = 0
with torch.no_grad():
for batch in tqdm(split):
real_landmarks = batch['original_landmarks']
gt_lm = batch['landmarks']
original_images = batch['original_image']
transformed_images = batch['image'].to(location)
predicted_landmarks, heatmaps, var, unnormal_hms = net(
transformed_images)
#transformed_landmarks = net.stn.transform_coords(predicted_landmarks, affine_params)
predicted_landmarks = predicted_landmarks.cpu()
unnormal = FaceLandmarksTrainingData.undo_normalization
original_images = batch['image'].cpu().detach().permute(
0, 2, 3, 1).numpy()
predicted_heatmaps = heatmaps.cpu().detach().numpy()
outline_errors = with_outline_error(predicted_landmarks.float(),
real_landmarks.float())
real_landmarks = real_landmarks.detach().numpy()
predicted_landmarks = predicted_landmarks.detach().numpy()
COLOR_GT = [255, 0, 127][::-1]
COLOR_PRED = [0, 0, 255][::-1]
N_LM = 68
for i in range(len(original_images)):
original_image = cv2.cvtColor(
unnormal(original_images[i]),
cv2.COLOR_RGB2BGR).astype(dtype=np.uint8)
dotsize = 1
face_gt = draw_landmarks(original_image,
real_landmarks[i],
color=COLOR_GT,
size=dotsize)
face_prediction = draw_landmarks(original_image,
predicted_landmarks[i],
color=COLOR_PRED,
size=dotsize)
face_both = draw_landmarks(face_gt,
predicted_landmarks[i],
color=COLOR_PRED,
size=dotsize)
# upscale heatmaps from 32x32 to 128x128
imgsize = 128
scaled_hms = np.array([
cv2.resize(hm,
dsize=(imgsize, imgsize),
interpolation=cv2.INTER_CUBIC)
for hm in predicted_heatmaps[i]
])
# make highest value in each heatmap 1 and make it a RGB image
prob_maps = np.array([
scaled_hms[j] / np.max(scaled_hms[j])
for j in range(len(scaled_hms))
],
dtype=np.float)
prob_maps = np.stack([prob_maps, prob_maps, prob_maps], axis=3)
# define a blending max that is used to vizualize the values in the heatmap
blend = np.ones(prob_maps.shape)
blend[:, :, :, 2] = 0 # R
blend[:, :, :, 1] = 255 # G
blend[:, :, :, 0] = 0 # B
# interpolate original image and blending mask weighted by the heatmap
heatmaps = (1 - prob_maps) * original_image + prob_maps * blend
heatmaps = heatmaps.astype(np.uint8)
# draw GT and pred
if landmarks_in_heatmaps:
dotsize = 2
heatmaps = [
draw_landmarks(heatmaps[k], [real_landmarks[i][k]],
size=dotsize,
color=COLOR_GT) for k in range(N_LM)
]
heatmaps = [
draw_landmarks(heatmaps[k],
[predicted_landmarks[i][k]],
size=dotsize,
color=COLOR_PRED) for k in range(N_LM)
]
# label with variance and error
errors = [
get_scaled_error(real_landmarks[i][k],
predicted_landmarks[i][k], imgsize)
for k in range(N_LM)
]
variances = [var[i][k].mean() for k in range(N_LM)]
heatmaps = [
add_description(
heatmaps[k],
"%d V%0.1f E%0.2f" % (k + 1, variances[k], errors[k]))
for k in range(N_LM)
]
original_image = add_description(original_image, "input img")
face_gt = add_description(face_gt, "ground truth")
face_prediction = add_description(face_prediction,
"prediction")
face_both = add_description(face_both, "GT + pred")
output_images = np.array(
[original_image, face_gt, face_prediction, face_both])
to_show = np.append(output_images, heatmaps, axis=0)
to_show = np.array([add_border(img, 2) for img in to_show])
allimages = gallery(to_show, cols=12)
evaluated_images.append(
(outline_errors[i].data.item(), allimages, curidx))
curidx += 1
print("Write files...")
ranked_images = sorted(evaluated_images, key=lambda x: x[0])
#order = [(err, idx) for err,_, idx in ranked_images]
#print(order)
for rank, (err, img, _) in tqdm(list(enumerate(ranked_images))):
cv2.imwrite(os.path.join(target, "%d_%f.png" % (rank, err)), img)
def visualize(model,
dataset,
target,
gpu=None,
splits=["easy", "hard"],
landmarks_in_heatmaps=True):
location = 'cpu' if gpu is None else "cuda:%d" % gpu
if location is not 'cpu':
# This fixes the problem that pytorch is always allocating memory on GPU 0 even if this is not included
# in the list of GPUs to use
torch.cuda.set_device(torch.device(location))
# cudnn.benchmark improves training speed when input sizes do not change
# https://discuss.pytorch.org/t/what-does-torch-backends-cudnn-benchmark-do/5936
# It selects the best algorithms as the training iterates over the dataset
cudnn.benchmark = True
print("Location: ", location)
data = torch.load(model, map_location=location)
state_dict = data['state_dict']
config = data['config']
num_workers = multiprocessing.cpu_count()
batch_size = config['batch_size'] if gpu is not None else num_workers
pin_memory = gpu is not None
print("Workers: ", num_workers)
print("Batchsize: ", batch_size)
net = ModelTrainer.create_net(config, verbose=False)
net.load_state_dict(state_dict)
net.eval()
net = net.to(location)
mkdir_if_not_exists(target)
normMean, normStd = FaceLandmarksTrainingData.TRAIN_MEAN, FaceLandmarksTrainingData.TRAIN_STD
normTransform = transforms.Normalize(normMean, normStd)
transform = transforms.Compose([
ImageTransform(transforms.ToPILImage()),
#ImageAndLabelTransform(RandomHorizontalFlip()),
#ImageAndLabelTransform(RandomRotation(min_angle=-0, max_angle=0, retain_scale=False)),
ImageTransform(transforms.ToTensor()),
ImageTransform(normTransform)
])
with h5py.File(dataset, 'r') as f:
if "easy" in splits:
print("Run on easy")
easy_d = FaceLandmarksEasyTestData(f, transform=transform)
#print(len(easy_d))
easy_loader = DataLoader(dataset=easy_d,
shuffle=False,
num_workers=num_workers,
pin_memory=pin_memory,
batch_size=batch_size)
visualize_split(net,
easy_loader,
os.path.join(target, "easy"),
location,
landmarks_in_heatmaps=landmarks_in_heatmaps)
if "hard" in splits:
print("Run on hard")
hard_d = FaceLandmarksHardTestData(f, transform=transform)
#print(len(hard_d))
hard_loader = DataLoader(dataset=hard_d,
shuffle=False,
num_workers=num_workers,
pin_memory=pin_memory,
batch_size=batch_size)
visualize_split(net,
hard_loader,
os.path.join(target, "hard"),
location,
landmarks_in_heatmaps=landmarks_in_heatmaps)
if "train" in splits:
print("Run on train")
train = FaceLandmarksTrainingData(f, transform=transform)
#print(len(train))
train_loader = DataLoader(dataset=train,
shuffle=False,
num_workers=num_workers,
pin_memory=pin_memory,
batch_size=batch_size)
visualize_split(net,
train_loader,
os.path.join(target, "train"),
location,
landmarks_in_heatmaps=landmarks_in_heatmaps)
args = parser.parse_args()
options = yaml.safe_load(open(args.template_file, "r"))
model = options['model']
hyperparameters = options['hyperparameters']
print("Read %s" % args.template_file)
print(">> Model: %s" % model)
print(">> HYPERPARAMETERS:")
max_key_len = max(map(len, hyperparameters.keys()))
for k, v in hyperparameters.items():
print("%s %s -> %s" % (k, " " * (max_key_len - len(k)), str(v)))
root_folder = args.output_directory
mkdir_if_not_exists(root_folder)
print("\nWriting everything to '%s'" % root_folder)
num_workers = args.num_workers
extra_fields = ['model', 'config_id']
hparam_names = list(hyperparameters.keys())
hparam_values = list(hyperparameters.values())
combinations = [(i, combi)
for i, combi in enumerate(itertools.product(*hparam_values))]
if args.count_only:
print(len(combinations), "combinations")
exit()
type=str,
nargs="+",
default=["gt", "hg", "pdm", "connection", "confidence"],
choices=[
"gt", "hg", "pdm", "connection_gt", "connection_hg", "confidence"
],
help="What to include in frames")
parser.add_argument("--epoch_frequency",
type=int,
default=1,
metavar="N",
help="Plot only every N epochs")
parser.add_argument("--encoder", type=str, help="Path to encoder (.torch)")
args = parser.parse_args()
# TODO offer option to specify which images to plot
target_dir = args.target
mkdir_if_not_exists(target_dir)
hg_results = json.load(open(args.hourglass_results, "r"))
for split in set(args.splits):
target = os.path.join(target_dir, split)
mkdir_if_not_exists(os.path.join(target_dir, target))
make_video(pdm_path=args.pdm,
hg_results=hg_results[split],
target=target,
include=list(set(args.plot)),
epoch_frequency=args.epoch_frequency,
encoder=args.encoder)