Esempio n. 1
0
def epoch_pass(
    loader,
    model,
    optimizer=None,
    scheduler=None,
    epoch=0,
    img_folder=None,
    fig=None,
    display_freq=10,
    epoch_display_freq=1,
    lr_decay_gamma=0,
    freeze_batchnorm=True,
    dump_results_path=None,
    render_folder=None,
    render_freq=10,
    true_root=False,
):
    prefix = "val"
    reorder_idxs, unorder_idxs = get_order_idxs()
    evaluators = {
        # "joints2d_trans": EvalUtil(),
        "joints2d_base": EvalUtil(),
        "corners2d_base": EvalUtil(),
        "verts2d_base": EvalUtil(),
        "joints3d_cent": EvalUtil(),
        "joints3d": EvalUtil(),
    }
    model.eval()
    model.cuda()
    avg_meters = AverageMeters()
    all_joints = []
    all_verts = []
    for batch_idx, batch in enumerate(tqdm(loader)):
        with torch.no_grad():
            loss, results, losses = model(batch)
            # Collect hand joints
            if true_root:
                results["recov_joints3d"][:, 0] = batch[BaseQueries.JOINTS3D][:, 0]
            recov_joints = results["recov_joints3d"].cpu().detach()[:, unorder_idxs]
            recov_joints[:, :, 0] = -recov_joints[:, :, 0]
            new_joints = [-val.numpy()[0] for val in recov_joints.split(1)]
            all_joints.extend(new_joints)

            # Collect hand vertices
            recov_verts = results["recov_handverts3d"].cpu().detach()
            recov_verts[:, :, 0] = -recov_verts[:, :, 0]
            new_verts = [-val.numpy()[0] for val in recov_verts.split(1)]
            all_verts.extend(new_verts)

        evaluate.feed_avg_meters(avg_meters, batch, results)
        if batch_idx % display_freq == 0 and epoch % epoch_display_freq == 0:
            img_filepath = f"{prefix}_epoch{epoch:04d}_batch{batch_idx:06d}.png"
            save_img_path = os.path.join(img_folder, img_filepath)
            samplevis.sample_vis(batch, results, fig=fig, save_img_path=save_img_path)
        evaluate.feed_evaluators(evaluators, batch, results)
    ho3dv2utils.dump(dump_results_path, all_joints, all_verts, codalab=True)
def epoch_pass(
    loader,
    model,
    train=False,
    optimizer=None,
    scheduler=None,
    epoch=0,
    img_folder=None,
    fig=None,
    display_freq=10,
    epoch_display_freq=1,
    lr_decay_gamma=0,
    loader_nb=2,
    freeze_batchnorm=True,
    premodel=None,
):
    if train:
        prefix = "train"
    else:
        prefix = "val"
    evaluators = {
        # "joints2d_trans": EvalUtil(),
        "joints2d_base": EvalUtil(),
        "corners2d_base": EvalUtil(),
        "verts2d_base": EvalUtil(),
        "joints3d_cent": EvalUtil(),
        "joints3d": EvalUtil(),
    }
    consist_evaluators = {
        # "joints2d_trans": EvalUtil(),
        "joints2d_base": EvalUtil(),
        "corners2d_base": EvalUtil(),
        "verts2d_base": EvalUtil(),
        "joints3d_cent": EvalUtil(),
        "joints3d": EvalUtil(),
    }
    avg_meters = AverageMeters()
    consist_avg_meters = AverageMeters()
    if train and not freeze_batchnorm:
        model.train()
    else:
        model.eval()
    for batch_idx, batch in enumerate(tqdm(loader)):
        if batch_idx % loader_nb == 0:
            losses = []
        loss, all_losses, results, pair_results = premodel.forward(batch)
        losses.append(loss.flatten())
        if train and ((batch_idx % loader_nb) == (loader_nb - 1)):
            if torch.isnan(loss):
                raise ValueError(f"Loss made of {all_losses} became nan!")
            optimizer.zero_grad()
            loss = torch.stack(losses).sum()
            if loss.requires_grad:
                loss.backward()
                optimizer.step()
        if "data" in batch["supervision"]:
            for loss_name, loss_val in all_losses.items():
                if loss_val is not None:
                    avg_meters.add_loss_value(loss_name, loss_val.mean().item())
            for sample, res in zip(batch["data"], results):
                evaluate.feed_avg_meters(avg_meters, sample, res)
                evaluate.feed_evaluators(evaluators, sample, res)
        elif "consist" in batch["supervision"]:
            for loss_name, loss_val in all_losses.items():
                if loss_val is not None:
                    consist_avg_meters.add_loss_value(loss_name, loss_val.mean().item())
            # Only compute metrics for weakly supervised frames (e.g. the first in the sequence)
            evaluate.feed_avg_meters(consist_avg_meters, batch["data"][0], results[0])
            evaluate.feed_evaluators(consist_evaluators, batch["data"][0], results[0])
        else:
            raise ValueError(f"Supervision {batch['supervision']} not in [data|consist]")
        if (batch_idx % display_freq) < loader_nb and epoch % epoch_display_freq == 0:
            img_prefix = f"{prefix}_epoch{epoch:04d}_batch{batch_idx:06d}"
            save_img_path = os.path.join(img_folder, img_prefix)
            warpvis.sample_vis(batch, results, pair_results, fig=fig, save_img_prefix=save_img_path)
    if lr_decay_gamma and scheduler is not None:
        scheduler.step()
    save_dict = {}
    for loss_name, avg_meter in avg_meters.average_meters.items():
        save_dict[loss_name] = {}
        loss_val = avg_meter.avg
        save_dict[loss_name][prefix] = loss_val
    for loss_name, avg_meter in consist_avg_meters.average_meters.items():
        if loss_name not in save_dict:
            save_dict[loss_name] = {}
        loss_val = avg_meter.avg
        save_dict[loss_name]["consist"] = loss_val
    evaluator_results = evaluate.parse_evaluators(evaluators)
    show_metrics = ["epe_mean"]  # "auc"
    for eval_name, eval_res in evaluator_results.items():
        for met in show_metrics:
            loss_name = f"{eval_name}_{met}"
            # Filter nans
            if eval_res[met] == eval_res[met]:
                save_dict[loss_name] = {}
                save_dict[loss_name][prefix] = eval_res[met]
    consist_evaluator_results = evaluate.parse_evaluators(consist_evaluators)
    for eval_name, eval_res in consist_evaluator_results.items():
        for met in show_metrics:
            loss_name = f"{eval_name}_{met}"
            # Filter nans
            if eval_res[met] == eval_res[met]:
                if loss_name not in save_dict:
                    save_dict[loss_name] = {}
                save_dict[loss_name]["consist"] = eval_res[met]
    img_filepath = f"{prefix}_epoch{epoch:04d}_eval.png"
    save_img_path = os.path.join(img_folder, img_filepath)
    # Filter out Nan pck curves
    evaluator_results = {
        eval_name: res for eval_name, res in evaluator_results.items() if res["epe_mean"] == res["epe_mean"]
    }
    evalvis.eval_vis(evaluator_results, save_img_path, fig=fig)
    pickle_path = save_img_path.replace(".png", ".pkl")
    with open(pickle_path, "wb") as p_f:
        pickle.dump(evaluator_results, p_f)
    return save_dict, avg_meters, evaluator_results
Esempio n. 3
0
def epoch_pass(
    loader,
    model,
    train=False,
    optimizer=None,
    scheduler=None,
    epoch=0,
    img_folder=None,
    fig=None,
    display_freq=10,
    epoch_display_freq=1,
    lr_decay_gamma=0,
    freeze_batchnorm=True,
):
    if train:
        prefix = "train"
    else:
        prefix = "val"
    evaluators = {
        # "joints2d_trans": EvalUtil(),
        "joints2d_base": EvalUtil(),
        "corners2d_base": EvalUtil(),
        "verts2d_base": EvalUtil(),
        "joints3d_cent": EvalUtil(),
        "joints3d": EvalUtil(),
    }
    if train and not freeze_batchnorm:
        model.train()
    else:
        model.eval()
    avg_meters = AverageMeters()
    render_step = 0
    # Loop over dataset
    for batch_idx, batch in enumerate(tqdm(loader)):
        # Compute outputs and losses
        if train:
            loss, results, losses = model(batch)
        else:
            with torch.no_grad():
                loss, results, losses = model(batch)
        # Optimize model if needed
        if train:
            if torch.isnan(loss):
                raise ValueError(f"Loss made of {losses} became nan!")
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        for loss_name, loss_val in losses.items():
            if loss_val is not None:
                avg_meters.add_loss_value(loss_name, loss_val.mean().item())
        evaluate.feed_avg_meters(avg_meters, batch, results)

        # Visualize outputs
        if batch_idx % display_freq == 0 and epoch % epoch_display_freq == 0:
            img_filepath = f"{prefix}_epoch{epoch:04d}_batch{batch_idx:06d}.png"
            save_img_path = os.path.join(img_folder, img_filepath)
            samplevis.sample_vis(batch,
                                 results,
                                 fig=fig,
                                 save_img_path=save_img_path)
        evaluate.feed_evaluators(evaluators, batch, results)
    if lr_decay_gamma and scheduler is not None:
        scheduler.step()
    save_dict = {}
    for loss_name, avg_meter in avg_meters.average_meters.items():
        save_dict[loss_name] = {}
        loss_val = avg_meter.avg
        save_dict[loss_name][prefix] = loss_val
    evaluator_results = evaluate.parse_evaluators(evaluators)
    for eval_name, eval_res in evaluator_results.items():
        for met in ["epe_mean", "auc"]:
            loss_name = f"{eval_name}_{met}"
            # Filter nans
            if eval_res[met] == eval_res[met]:
                save_dict[loss_name] = {}
                save_dict[loss_name][prefix] = eval_res[met]
    img_filepath = f"{prefix}_epoch{epoch:04d}_eval.png"
    save_img_path = os.path.join(img_folder, img_filepath)
    # Filter out Nan pck curves
    evaluator_results = {
        eval_name: res
        for eval_name, res in evaluator_results.items()
        if res["epe_mean"] == res["epe_mean"]
    }
    evalvis.eval_vis(evaluator_results, save_img_path, fig=fig)
    pickle_path = save_img_path.replace(".png", ".pkl")
    with open(pickle_path, "wb") as p_f:
        pickle.dump(evaluator_results, p_f)
    return save_dict, avg_meters, evaluator_results