demo.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
from spann3r.tools.vis import render_frames, find_render_cam, vis_pred_and_imgs
from pose_utils import solve_cemara
from backend_utils import improved_multiway_registration, pts2normal, point2mesh, combine_and_clean_point_clouds

def get_args_parser():
    parser = argparse.ArgumentParser('Spann3R demo', add_help=False)
    parser.add_argument('--save_path', type=str, default='./output/demo/', help='Path to experiment folder')
    parser.add_argument('--demo_path', type=str, default='./examples/s00567', help='Path to experiment folder')
    parser.add_argument('--ckpt_path', type=str, default='./checkpoints/spann3r.pth', help='ckpt path')
    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=1e-3, help='confidence threshold')
    parser.add_argument('--kf_every', type=int, default=10, help='map every kf_every frames')
    parser.add_argument('--vis', action='store_true', help='visualize')
    parser.add_argument('--voxel_size', type=float, default=0.004, help='voxel size for multiway registration')
    return parser

import tempfile
import subprocess
def extract_frames(video_path: str, duration: float = 20.0, fps: float = 3.0) -> str:
    temp_dir = tempfile.mkdtemp()
    output_path = os.path.join(temp_dir, "%03d.jpg")
    
    filter_complex = f"select='if(lt(t,{duration}),1,0)',fps={fps}"

    command = [
        "ffmpeg",
        "-i", video_path,
        "-vf", filter_complex,
        "-vsync", "0",
        output_path
    ]
    
    subprocess.run(command, check=True)
    return temp_dir

@torch.no_grad()
def main(args):

    workspace = args.save_path
    os.makedirs(workspace, exist_ok=True)

    ##### Load model
    model = Spann3R(dus3r_name='./checkpoints/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth', 
                use_feat=False).to(args.device)
    
    model.load_state_dict(torch.load(args.ckpt_path)['model'])
    model.eval()
    
    if args.demo_path.endswith('.mp4') or args.demo_path.endswith('.avi') or args.demo_path.endswith('.webm'):
        args.demo_path = extract_frames(args.demo_path)
        args.kf_every = 1

    ##### Load dataset
    dataset = Demo(ROOT=args.demo_path, resolution=224, full_video=True, kf_every=args.kf_every)
    dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0)

    batch = dataloader.__iter__().__next__()

    ##### Inference
    for view in batch:
         view['img'] = view['img'].to(args.device, non_blocking=True)
           

    demo_name = args.demo_path.split("/")[-1]

    print(f'Started reconstruction for {demo_name}')

    if args.offline:
        imgs_all = []
        for j, view in enumerate(batch):
            img = view['img']
            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 {demo_name}, FPS: {fps:.2f}')

    ##### Save results

    save_demo_path = osp.join(workspace, demo_name)
    os.makedirs(save_demo_path, exist_ok=True)

    pts_all = []
    pts_normal_all = []
    pts_gt_all = []
    images_all = []
    masks_all = []
    conf_sig_all = []
    cameras_all = []

    last_focal = None
    for j, view in enumerate(ordered_batch):
        
        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_normal = pts2normal(preds[j]['pts3d' if j==0 else 'pts3d_in_other_view'][0]).cpu().numpy()
        conf = preds[j]['conf'][0].cpu().data.numpy()
        conf_sig = (conf - 1) / conf
        pts_gt = view['pts3d'].cpu().numpy()[0]

        camera, last_focal, depth_map = solve_cemara(torch.tensor(pts), torch.tensor(conf_sig) > args.conf_thresh, 
                                          args.device, focal=last_focal)
        pts_scale = depth_map / last_focal

        images_all.append((image[None, ...] + 1.0)/2.0)
        pts_all.append(pts[None, ...])
        pts_normal_all.append(pts_normal[None, ...])
        pts_gt_all.append(pts_gt[None, ...])
        pts_scale_all.append(pts_scale[None, ...])
        masks_all.append(mask[None, ...])
        conf_sig_all.append(conf_sig[None, ...])

    images_all = np.concatenate(images_all, axis=0)
    pts_all = np.concatenate(pts_all, axis=0)
    pts_normal_all = np.concatenate(pts_normal_all, axis=0)
    pts_gt_all = np.concatenate(pts_gt_all, axis=0)
    masks_all = np.concatenate(masks_all, axis=0)
    conf_sig_all = np.concatenate(conf_sig_all, axis=0)

    # Create point clouds for multiway registration
    pcds = []
    for j in range(len(pts_all)):
        pcd = o3d.geometry.PointCloud()
        mask = conf_sig_all[j] > args.conf_thresh
        pcd.points = o3d.utility.Vector3dVector(pts_all[j][mask])
        pcd.colors = o3d.utility.Vector3dVector(images_all[j][mask])
        pcd.normals = o3d.utility.Vector3dVector(pts_normal_all[j][mask])
        pcds.append(pcd)

    print("Performing global registration...")
    pcd_combined, _, _ = improved_multiway_registration(pcds, voxel_size=0.001)
    # pcd_combined = combine_and_clean_point_clouds(transformed_pcds, voxel_size=args.voxel_size * 0.1)
    mesh_recon = point2mesh(pcd_combined)


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