|
|
|
|
|
|
|
|
|
|
|
|
|
import os |
|
import sys |
|
sys.path.append(os.path.abspath('./modules')) |
|
|
|
import math |
|
import tempfile |
|
import gradio |
|
import torch |
|
import spaces |
|
import numpy as np |
|
import functools |
|
import trimesh |
|
import copy |
|
from PIL import Image |
|
from scipy.spatial.transform import Rotation |
|
|
|
from modules.pe3r.images import Images |
|
|
|
from modules.dust3r.inference import inference |
|
from modules.dust3r.image_pairs import make_pairs |
|
from modules.dust3r.utils.image import load_images, rgb |
|
from modules.dust3r.utils.device import to_numpy |
|
from modules.dust3r.viz import add_scene_cam, CAM_COLORS, OPENGL, pts3d_to_trimesh, cat_meshes |
|
from modules.dust3r.cloud_opt import global_aligner, GlobalAlignerMode |
|
from copy import deepcopy |
|
import cv2 |
|
from typing import Any, Dict, Generator,List |
|
import matplotlib.pyplot as pl |
|
|
|
from modules.mobilesamv2.utils.transforms import ResizeLongestSide |
|
from modules.pe3r.models import Models |
|
import torchvision.transforms as tvf |
|
|
|
from modules.mast3r.model import AsymmetricMASt3R |
|
|
|
silent = False |
|
device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
|
|
MAST3R_CKP = 'naver/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric' |
|
mast3r = AsymmetricMASt3R.from_pretrained(MAST3R_CKP).to(device) |
|
|
|
|
|
|
|
def _convert_scene_output_to_glb(outdir, imgs, pts3d, mask, focals, cams2world, cam_size=0.05, |
|
cam_color=None, as_pointcloud=False, |
|
transparent_cams=False): |
|
assert len(pts3d) == len(mask) <= len(imgs) <= len(cams2world) == len(focals) |
|
pts3d = to_numpy(pts3d) |
|
imgs = to_numpy(imgs) |
|
focals = to_numpy(focals) |
|
cams2world = to_numpy(cams2world) |
|
|
|
scene = trimesh.Scene() |
|
|
|
|
|
if as_pointcloud: |
|
pts = np.concatenate([p[m] for p, m in zip(pts3d, mask)]) |
|
col = np.concatenate([p[m] for p, m in zip(imgs, mask)]) |
|
pct = trimesh.PointCloud(pts.reshape(-1, 3), colors=col.reshape(-1, 3)) |
|
scene.add_geometry(pct) |
|
else: |
|
meshes = [] |
|
for i in range(len(imgs)): |
|
meshes.append(pts3d_to_trimesh(imgs[i], pts3d[i], mask[i])) |
|
mesh = trimesh.Trimesh(**cat_meshes(meshes)) |
|
scene.add_geometry(mesh) |
|
|
|
|
|
for i, pose_c2w in enumerate(cams2world): |
|
if isinstance(cam_color, list): |
|
camera_edge_color = cam_color[i] |
|
else: |
|
camera_edge_color = cam_color or CAM_COLORS[i % len(CAM_COLORS)] |
|
add_scene_cam(scene, pose_c2w, camera_edge_color, |
|
None if transparent_cams else imgs[i], focals[i], |
|
imsize=imgs[i].shape[1::-1], screen_width=cam_size) |
|
|
|
rot = np.eye(4) |
|
rot[:3, :3] = Rotation.from_euler('y', np.deg2rad(180)).as_matrix() |
|
scene.apply_transform(np.linalg.inv(cams2world[0] @ OPENGL @ rot)) |
|
outfile = os.path.join(outdir, 'scene.glb') |
|
if not silent: |
|
print('(exporting 3D scene to', outfile, ')') |
|
scene.export(file_obj=outfile) |
|
return outfile |
|
|
|
def get_3D_model_from_scene(outdir, scene, min_conf_thr=3, as_pointcloud=False, mask_sky=False, |
|
clean_depth=False, transparent_cams=False, cam_size=0.05): |
|
""" |
|
extract 3D_model (glb file) from a reconstructed scene |
|
""" |
|
if scene is None: |
|
return None |
|
|
|
if clean_depth: |
|
scene = scene.clean_pointcloud() |
|
if mask_sky: |
|
scene = scene.mask_sky() |
|
|
|
|
|
rgbimg = scene.ori_imgs |
|
focals = scene.get_focals().cpu() |
|
cams2world = scene.get_im_poses().cpu() |
|
|
|
pts3d = to_numpy(scene.get_pts3d()) |
|
scene.min_conf_thr = float(scene.conf_trf(torch.tensor(min_conf_thr))) |
|
msk = to_numpy(scene.get_masks()) |
|
return _convert_scene_output_to_glb(outdir, rgbimg, pts3d, msk, focals, cams2world, as_pointcloud=as_pointcloud, |
|
transparent_cams=transparent_cams, cam_size=cam_size) |
|
|
|
def mask_nms(masks, threshold=0.8): |
|
keep = [] |
|
mask_num = len(masks) |
|
suppressed = np.zeros((mask_num), dtype=np.int64) |
|
for i in range(mask_num): |
|
if suppressed[i] == 1: |
|
continue |
|
keep.append(i) |
|
for j in range(i + 1, mask_num): |
|
if suppressed[j] == 1: |
|
continue |
|
intersection = (masks[i] & masks[j]).sum() |
|
if min(intersection / masks[i].sum(), intersection / masks[j].sum()) > threshold: |
|
suppressed[j] = 1 |
|
return keep |
|
|
|
def filter(masks, keep): |
|
ret = [] |
|
for i, m in enumerate(masks): |
|
if i in keep: ret.append(m) |
|
return ret |
|
|
|
def mask_to_box(mask): |
|
if mask.sum() == 0: |
|
return np.array([0, 0, 0, 0]) |
|
|
|
|
|
rows = np.any(mask, axis=1) |
|
cols = np.any(mask, axis=0) |
|
|
|
|
|
top = np.argmax(rows) |
|
bottom = len(rows) - 1 - np.argmax(np.flip(rows)) |
|
left = np.argmax(cols) |
|
right = len(cols) - 1 - np.argmax(np.flip(cols)) |
|
|
|
return np.array([left, top, right, bottom]) |
|
|
|
def box_xyxy_to_xywh(box_xyxy): |
|
box_xywh = deepcopy(box_xyxy) |
|
box_xywh[2] = box_xywh[2] - box_xywh[0] |
|
box_xywh[3] = box_xywh[3] - box_xywh[1] |
|
return box_xywh |
|
|
|
def get_seg_img(mask, box, image): |
|
image = image.copy() |
|
x, y, w, h = box |
|
|
|
box_area = w * h |
|
mask_area = mask.sum() |
|
if 1 - (mask_area / box_area) < 0.2: |
|
image[mask == 0] = np.array([0, 0, 0], dtype=np.uint8) |
|
else: |
|
random_values = np.random.randint(0, 255, size=image.shape, dtype=np.uint8) |
|
image[mask == 0] = random_values[mask == 0] |
|
seg_img = image[y:y+h, x:x+w, ...] |
|
return seg_img |
|
|
|
def pad_img(img): |
|
h, w, _ = img.shape |
|
l = max(w,h) |
|
pad = np.zeros((l,l,3), dtype=np.uint8) |
|
if h > w: |
|
pad[:,(h-w)//2:(h-w)//2 + w, :] = img |
|
else: |
|
pad[(w-h)//2:(w-h)//2 + h, :, :] = img |
|
return pad |
|
|
|
def batch_iterator(batch_size: int, *args) -> Generator[List[Any], None, None]: |
|
assert len(args) > 0 and all( |
|
len(a) == len(args[0]) for a in args |
|
), "Batched iteration must have inputs of all the same size." |
|
n_batches = len(args[0]) // batch_size + int(len(args[0]) % batch_size != 0) |
|
for b in range(n_batches): |
|
yield [arg[b * batch_size : (b + 1) * batch_size] for arg in args] |
|
|
|
def slerp(u1, u2, t): |
|
""" |
|
Perform spherical linear interpolation (Slerp) between two unit vectors. |
|
|
|
Args: |
|
- u1 (torch.Tensor): First unit vector, shape (1024,) |
|
- u2 (torch.Tensor): Second unit vector, shape (1024,) |
|
- t (float): Interpolation parameter |
|
|
|
Returns: |
|
- torch.Tensor: Interpolated vector, shape (1024,) |
|
""" |
|
|
|
dot_product = torch.sum(u1 * u2) |
|
|
|
|
|
dot_product = torch.clamp(dot_product, -1.0, 1.0) |
|
|
|
|
|
theta = torch.acos(dot_product) |
|
|
|
|
|
sin_theta = torch.sin(theta) |
|
if sin_theta == 0: |
|
|
|
return u1 + t * (u2 - u1) |
|
|
|
s1 = torch.sin((1 - t) * theta) / sin_theta |
|
s2 = torch.sin(t * theta) / sin_theta |
|
|
|
|
|
return s1 * u1 + s2 * u2 |
|
|
|
def slerp_multiple(vectors, t_values): |
|
""" |
|
Perform spherical linear interpolation (Slerp) for multiple vectors. |
|
|
|
Args: |
|
- vectors (torch.Tensor): Tensor of vectors, shape (n, 1024) |
|
- a_values (torch.Tensor): Tensor of values corresponding to each vector, shape (n,) |
|
|
|
Returns: |
|
- torch.Tensor: Interpolated vector, shape (1024,) |
|
""" |
|
n = vectors.shape[0] |
|
|
|
|
|
interpolated_vector = vectors[0] |
|
|
|
|
|
for i in range(1, n): |
|
|
|
t = t_values[i] / (t_values[i] + t_values[i-1]) |
|
interpolated_vector = slerp(interpolated_vector, vectors[i], t) |
|
|
|
return interpolated_vector |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@spaces.GPU(duration=180) |
|
def get_reconstructed_scene(outdir, filelist, schedule, niter, min_conf_thr, |
|
as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size, |
|
scenegraph_type, winsize, refid): |
|
""" |
|
from a list of images, run dust3r inference, global aligner. |
|
then run get_3D_model_from_scene |
|
""" |
|
if len(filelist) < 2: |
|
raise gradio.Error("Please input at least 2 images.") |
|
|
|
images = Images(filelist=filelist, device=device) |
|
|
|
|
|
|
|
|
|
|
|
rev_cog_seg_maps = [] |
|
for tmp_img in images.np_images: |
|
rev_seg_map = -np.ones(tmp_img.shape[:2], dtype=np.int64) |
|
rev_cog_seg_maps.append(rev_seg_map) |
|
cog_seg_maps = rev_cog_seg_maps |
|
cog_feats = torch.zeros((1, 1024)) |
|
imgs = load_images(images, rev_cog_seg_maps, size=512, verbose=not silent) |
|
|
|
if len(imgs) == 1: |
|
imgs = [imgs[0], copy.deepcopy(imgs[0])] |
|
imgs[1]['idx'] = 1 |
|
|
|
if scenegraph_type == "swin": |
|
scenegraph_type = scenegraph_type + "-" + str(winsize) |
|
elif scenegraph_type == "oneref": |
|
scenegraph_type = scenegraph_type + "-" + str(refid) |
|
|
|
pairs = make_pairs(imgs, scene_graph=scenegraph_type, prefilter=None, symmetrize=True) |
|
output = inference(pairs, mast3r, device, batch_size=1, verbose=not silent) |
|
mode = GlobalAlignerMode.PointCloudOptimizer if len(imgs) > 2 else GlobalAlignerMode.PairViewer |
|
scene_1 = global_aligner(output, cog_seg_maps, rev_cog_seg_maps, cog_feats, device=device, mode=mode, verbose=not silent) |
|
lr = 0.01 |
|
|
|
loss = scene_1.compute_global_alignment(tune_flg=True, init='mst', niter=niter, schedule=schedule, lr=lr) |
|
|
|
try: |
|
ImgNorm = tvf.Compose([tvf.ToTensor(), tvf.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) |
|
for i in range(len(imgs)): |
|
|
|
imgs[i]['img'] = ImgNorm(scene_1.imgs[i])[None] |
|
pairs = make_pairs(imgs, scene_graph=scenegraph_type, prefilter=None, symmetrize=True) |
|
output = inference(pairs, mast3r, device, batch_size=1, verbose=not silent) |
|
mode = GlobalAlignerMode.PointCloudOptimizer if len(imgs) > 2 else GlobalAlignerMode.PairViewer |
|
scene = global_aligner(output, cog_seg_maps, rev_cog_seg_maps, cog_feats, device=device, mode=mode, verbose=not silent) |
|
ori_imgs = scene.ori_imgs |
|
lr = 0.01 |
|
|
|
loss = scene.compute_global_alignment(tune_flg=False, init='mst', niter=niter, schedule=schedule, lr=lr) |
|
except Exception as e: |
|
scene = scene_1 |
|
scene.imgs = ori_imgs |
|
scene.ori_imgs = ori_imgs |
|
print(e) |
|
|
|
outfile = get_3D_model_from_scene(outdir, scene, min_conf_thr, as_pointcloud, mask_sky, |
|
clean_depth, transparent_cams, cam_size) |
|
torch.cuda.empty_cache() |
|
|
|
return scene, outfile |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
with tempfile.TemporaryDirectory(suffix='pe3r_gradio_demo') as tmpdirname: |
|
recon_fun = functools.partial(get_reconstructed_scene, tmpdirname) |
|
|
|
|
|
|
|
with gradio.Blocks(css=""".gradio-container {margin: 0 !important; min-width: 100%};""", title="PE3R Demo") as demo: |
|
|
|
scene = gradio.State(None) |
|
gradio.HTML('<h2 style="text-align: center;">PE3R Demo</h2>') |
|
with gradio.Column(): |
|
inputfiles = gradio.File(file_count="multiple") |
|
with gradio.Row(): |
|
schedule = gradio.Dropdown(["linear", "cosine"], |
|
value='linear', label="schedule", info="For global alignment!", |
|
visible=False) |
|
niter = gradio.Number(value=300, precision=0, minimum=0, maximum=5000, |
|
label="num_iterations", info="For global alignment!", |
|
visible=False) |
|
scenegraph_type = gradio.Dropdown([("complete: all possible image pairs", "complete"), |
|
("swin: sliding window", "swin"), |
|
("oneref: match one image with all", "oneref")], |
|
value='complete', label="Scenegraph", |
|
info="Define how to make pairs", |
|
interactive=True, |
|
visible=False) |
|
winsize = gradio.Slider(label="Scene Graph: Window Size", value=1, |
|
minimum=1, maximum=1, step=1, visible=False) |
|
refid = gradio.Slider(label="Scene Graph: Id", value=0, minimum=0, maximum=0, step=1, visible=False) |
|
|
|
run_btn = gradio.Button("Reconstruct") |
|
|
|
with gradio.Row(): |
|
|
|
min_conf_thr = gradio.Slider(label="min_conf_thr", value=3.0, minimum=1.0, maximum=20, step=0.1, visible=False) |
|
|
|
cam_size = gradio.Slider(label="cam_size", value=0.05, minimum=0.001, maximum=0.1, step=0.001, visible=False) |
|
with gradio.Row(): |
|
as_pointcloud = gradio.Checkbox(value=True, label="As pointcloud", visible=False) |
|
|
|
mask_sky = gradio.Checkbox(value=False, label="Mask sky", visible=False) |
|
clean_depth = gradio.Checkbox(value=True, label="Clean-up depthmaps", visible=False) |
|
transparent_cams = gradio.Checkbox(value=True, label="Transparent cameras", visible=False) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
outmodel = gradio.Model3D() |
|
|
|
|
|
run_btn.click(fn=recon_fun, |
|
inputs=[inputfiles, schedule, niter, min_conf_thr, as_pointcloud, |
|
mask_sky, clean_depth, transparent_cams, cam_size, |
|
scenegraph_type, winsize, refid], |
|
outputs=[scene, outmodel]) |
|
|
|
|
|
|
|
|
|
|
|
demo.launch(show_error=True, share=None, server_name=None, server_port=None) |
|
|