File size: 2,424 Bytes
8bc0782 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import rerun as rr
import rerun.blueprint as rrb
from pathlib import Path
from argparse import ArgumentParser
import torch
from mini_dust3r.api import OptimizedResult, inferece_dust3r, log_optimized_result
from mini_dust3r.model import AsymmetricCroCo3DStereo
def create_blueprint(image_name_list: list[str], log_path: Path) -> rrb.Blueprint:
# dont show 2d views if there are more than 4 images as to not clutter the view
if len(image_name_list) > 4:
blueprint = rrb.Blueprint(
rrb.Horizontal(
rrb.Spatial3DView(origin=f"{log_path}"),
),
collapse_panels=True,
)
else:
blueprint = rrb.Blueprint(
rrb.Horizontal(
contents=[
rrb.Spatial3DView(origin=f"{log_path}"),
rrb.Vertical(
contents=[
rrb.Spatial2DView(
origin=f"{log_path}/camera_{i}/pinhole/",
contents=[
"+ $origin/**",
],
)
for i in range(len(image_name_list))
]
),
],
column_shares=[3, 1],
),
collapse_panels=True,
)
return blueprint
def main(image_dir: Path):
if torch.backends.mps.is_available():
device = "mps"
elif torch.cuda.is_available():
device = "cuda"
else:
device = "cpu"
model = AsymmetricCroCo3DStereo.from_pretrained(
"naver/DUSt3R_ViTLarge_BaseDecoder_512_dpt"
).to(device)
optimized_results: OptimizedResult = inferece_dust3r(
image_dir_or_list=image_dir,
model=model,
device=device,
batch_size=1,
)
blueprint = create_blueprint(image_dir, "world")
rr.send_blueprint(blueprint)
log_optimized_result(optimized_results, Path("world"))
if __name__ == "__main__":
parser = ArgumentParser("mini-dust3r rerun demo script")
parser.add_argument(
"--image-dir",
type=Path,
help="Directory containing images to process",
required=True,
)
rr.script_add_args(parser)
args = parser.parse_args()
rr.script_setup(args, "mini-dust3r")
main(args.image_dir)
rr.script_teardown(args) |