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)