eval.py

import os
import time
import torch
import argparse
import numpy as np
import open3d as o3d
import os.path as osp
from dust3r.losses import L21
from spann3r.model import Spann3R
from dust3r.inference import inference
from dust3r.utils.geometry import geotrf
from dust3r.image_pairs import make_pairs
from spann3r.loss import Regr3D_t_ScaleShiftInv
from spann3r.datasets import *
from torch.utils.data import DataLoader
from spann3r.tools.eval_recon import accuracy, completion

def get_args_parser():
    parser = argparse.ArgumentParser('Spann3R evaluation', add_help=False)
    parser.add_argument('--exp_path', type=str, help='Path to experiment folder', default='./checkpoints')
    parser.add_argument('--exp_name', type=str, default='ckpt_best', help='Path to experiment folder')
    parser.add_argument('--ckpt', type=str, default='spann3r.pth', help='ckpt name')
    parser.add_argument('--scenegraph_type', type=str, default='complete', help='scenegraph type')
    parser.add_argument('--offline', action='store_true', help='offline reconstruction')
    parser.add_argument('--device', type=str, default='cuda:0', help='device')
    parser.add_argument('--conf_thresh', type=float, default=0.0, help='confidence threshold')

    return parser
    

def main(args):
    workspace = args.exp_path
    ckpt_path = osp.join(workspace, args.ckpt)
    if not osp.exists(workspace):
        raise FileNotFoundError(f"Workspace {workspace} not found")
    
    exp_path = osp.join(workspace, args.exp_name)
    os.makedirs(exp_path, exist_ok=True)

    datasets_all = {
        '7scenes': SevenScenes(split='test', ROOT="./data/7scenes",
                                resolution=224, num_seq=1, full_video=True, kf_every=20),
        'NRGBD': NRGBD(split='test', ROOT="./data/neural_rgbd", 
                           resolution=224, num_seq=1, full_video=True, kf_every=40),
        'DTU': DTU(split='test', ROOT="./data/dtu_test",
                   resolution=224, num_seq=1, full_video=True, kf_every=5),
    }
    model = Spann3R(dus3r_name='./checkpoints/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth', 
                use_feat=False).to(args.device)
    
    model.load_state_dict(torch.load(ckpt_path)['model'])
    model.eval()    


    criterion = Regr3D_t_ScaleShiftInv(L21, norm_mode=False, gt_scale=True)

    with torch.no_grad():
        for name_data, dataset in datasets_all.items():
            save_path = osp.join(exp_path, name_data)
            if args.offline:
                save_path = osp.join(save_path + '_offline')
            os.makedirs(save_path, exist_ok=True)

            log_file = osp.join(save_path, 'logs.txt')
            os.makedirs(save_path, exist_ok=True)

            acc_all = 0
            acc_all_med = 0
            comp_all = 0
            comp_all_med = 0
            nc1_all = 0
            nc1_all_med = 0
            nc2_all = 0
            nc2_all_med = 0

            fps_all = []
            time_all = []

            dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0)

            for i, batch in enumerate(dataloader):

                for view in batch:
                    for name in 'img pts3d valid_mask camera_pose camera_intrinsics F_matrix corres'.split():  # pseudo_focal
                        if name not in view:
                            continue
                        view[name] = view[name].to(args.device, non_blocking=True)


                print(f'Started reconstruction for {name_data} {i+1}/{len(dataloader)}')
                
                if args.offline:
                    imgs_all = []
                    for j, view in enumerate(batch):
                        img = view['img']
                        shape1 = [img.size()[::-1]]

                        imgs_all.append(
                            dict(
                                img=img,
                                true_shape=torch.tensor(img.shape[2:]).unsqueeze(0),
                                idx=j,
                                instance=str(j)
                            )
                        )
                    start = time.time()
                    pairs = make_pairs(imgs_all, scene_graph=args.scenegraph_type, prefilter=None, symmetrize=True)
                    output = inference(pairs, model.dust3r, args.device, batch_size=2, verbose=True)
                    preds, preds_all, idx_used = model.offline_reconstruction(batch, output) 
                    end = time.time()
                    ordered_batch = [batch[i] for i in idx_used]
                else:
                    start = time.time()
                    preds, preds_all = model.forward(batch) 
                    end = time.time()
                    ordered_batch = batch

                fps = len(batch) / (end - start)                

                print(f'Finished reconstruction for {name_data} {i+1}/{len(dataloader)}, FPS: {fps:.2f}')

                fps_all.append(fps)
                time_all.append(end - start)
                

                # Evaluation
                print(f'Evaluation for {name_data} {i+1}/{len(dataloader)}')
                gt_pts, pred_pts, gt_factor, pr_factor, masks, monitoring = criterion.get_all_pts3d_t(ordered_batch, preds_all)
                pred_scale, gt_scale, pred_shift_z, gt_shift_z  = monitoring['pred_scale'], monitoring['gt_scale'], monitoring['pred_shift_z'], monitoring['gt_shift_z']
                
                in_camera1 = None
                pts_all = []
                pts_gt_all = []
                images_all = []
                masks_all = []
                conf_all = []

                for j, view in enumerate(ordered_batch):
                    if in_camera1 is None:
                        in_camera1 = view['camera_pose'][0].cpu()
                    
                    image = view['img'].permute(0, 2, 3, 1).cpu().numpy()[0]
                    mask = view['valid_mask'].cpu().numpy()[0]

                    # pts = preds[j]['pts3d' if j==0 else 'pts3d_in_other_view'].detach().cpu().numpy()[0]
                    pts = pred_pts[0][j].cpu().numpy()[0] if j < len(pred_pts[0]) else pred_pts[1][-1].cpu().numpy()[0]
                    conf = preds[j]['conf'][0].cpu().data.numpy()

                    pts_gt = gt_pts[j].detach().cpu().numpy()[0]

                    #### Align predicted 3D points to the ground truth 
                    pts[..., -1] += gt_shift_z.cpu().numpy().item()
                    pts = geotrf(in_camera1, pts)

                    pts_gt[..., -1] += gt_shift_z.cpu().numpy().item()
                    pts_gt = geotrf(in_camera1, pts_gt)

                    images_all.append((image[None, ...] + 1.0)/2.0)
                    pts_all.append(pts[None, ...])
                    pts_gt_all.append(pts_gt[None, ...])
                    masks_all.append(mask[None, ...])
                    conf_all.append(conf[None, ...])
                
                images_all = np.concatenate(images_all, axis=0)
                pts_all = np.concatenate(pts_all, axis=0)
                pts_gt_all = np.concatenate(pts_gt_all, axis=0)
                masks_all = np.concatenate(masks_all, axis=0)
                conf_all = np.concatenate(conf_all, axis=0)

                scene_id = view['label'][0].rsplit('/', 1)[0]

                save_params = {}

                save_params['images_all'] = images_all
                save_params['pts_all'] = pts_all
                save_params['pts_gt_all'] = pts_gt_all
                save_params['masks_all'] = masks_all
                save_params['conf_all'] = conf_all

                np.save(os.path.join(save_path, f"{scene_id.replace('/', '_')}.npy"), save_params)


                if 'DTU' in name_data:
                    threshold = 100
                else:
                    threshold = 0.1

                
                pts_all_masked = pts_all[masks_all > 0]
                pts_gt_all_masked = pts_gt_all[masks_all > 0]
                images_all_masked = images_all[masks_all > 0]

                pcd = o3d.geometry.PointCloud()
                pcd.points = o3d.utility.Vector3dVector(pts_all_masked.reshape(-1, 3))
                pcd.colors = o3d.utility.Vector3dVector(images_all_masked.reshape(-1, 3))
                o3d.io.write_point_cloud(os.path.join(save_path, f"{scene_id.replace('/', '_')}-mask.ply"), pcd)

                pcd_gt = o3d.geometry.PointCloud()
                pcd_gt.points = o3d.utility.Vector3dVector(pts_gt_all_masked.reshape(-1, 3))
                pcd_gt.colors = o3d.utility.Vector3dVector(images_all_masked.reshape(-1, 3) / 255.0)
                o3d.io.write_point_cloud(os.path.join(save_path, f"{scene_id.replace('/', '_')}-gt.ply"), pcd_gt)

                trans_init = np.eye(4)
                reg_p2p = o3d.pipelines.registration.registration_icp(
                    pcd, pcd_gt, threshold, trans_init,
                    o3d.pipelines.registration.TransformationEstimationPointToPoint())
                
                transformation = reg_p2p.transformation
                            
                pcd = pcd.transform(transformation)
                pcd.estimate_normals()
                pcd_gt.estimate_normals()

                gt_normal = np.asarray(pcd_gt.normals)
                pred_normal = np.asarray(pcd.normals)

                acc, acc_med, nc1, nc1_med = accuracy(pcd_gt.points, pcd.points, gt_normal, pred_normal)
                comp, comp_med, nc2, nc2_med = completion(pcd_gt.points, pcd.points, gt_normal, pred_normal)
                

                print(f"Idx: {scene_id}, Acc: {acc}, Comp: {comp}, NC1: {nc1}, NC2: {nc2} - Acc_med: {acc_med}, Compc_med: {comp_med}, NC1c_med: {nc1_med}, NC2c_med: {nc2_med}", file=open(log_file, "a"))

                acc_all += acc
                comp_all += comp
                nc1_all += nc1
                nc2_all += nc2

                acc_all_med += acc_med
                comp_all_med += comp_med
                nc1_all_med += nc1_med
                nc2_all_med += nc2_med


                # release cuda memory
                torch.cuda.empty_cache()

                print(f"Finished evaluation for {name_data} {i+1}/{len(dataloader)}")



                # Get depth from pcd and run TSDFusion
                
            
            print(f"Dataset: {name_data}, Accuracy: {acc_all/len(dataloader)}, Completion: {comp_all/len(dataloader)}, NC1: {nc1_all/len(dataloader)}, NC2: {nc2_all/len(dataloader)} - Acc_med: {acc_all_med/len(dataloader)}, Comp_med: {comp_all_med/len(dataloader)}, NC1_med: {nc1_all_med/len(dataloader)}, NC2_med: {nc2_all_med/len(dataloader)}", file=open(log_file, "a"))
            print(f"Average fps: {sum(fps) / len(fps)}, Average time: {sum(time_all) / len(time_all)}", file=open(log_file, "a"))
                


if __name__ == '__main__':
    parser = get_args_parser()
    args = parser.parse_args()
    main(args)