# Copyright (C) 2024-present Naver Corporation. All rights reserved. # Licensed under CC BY-NC-SA 4.0 (non-commercial use only). # # -------------------------------------------------------- # gradio demo # -------------------------------------------------------- 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 silent = True device = 'cuda' if torch.cuda.is_available() else 'cpu' pe3r = Models(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, silent=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() # # full pointcloud # 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) # # add each camera # 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 # # @spaces.GPU(duration=180) # def get_3D_model_from_scene(outdir, silent, 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 # # post processes # if clean_depth: # scene = scene.clean_pointcloud() # if mask_sky: # scene = scene.mask_sky() # # get optimized values from scene # rgbimg = scene.ori_imgs # focals = scene.get_focals().cpu() # cams2world = scene.get_im_poses().cpu() # # 3D pointcloud from depthmap, poses and intrinsics # 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, silent=silent) # 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]) # # Get the rows and columns where the mask is 1 # rows = np.any(mask, axis=1) # cols = np.any(mask, axis=0) # # Get top, bottom, left, right edges # 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 # # image[mask == 0] = np.array([0, 0, 0], dtype=np.uint8) # 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,) # """ # # Compute the dot product # dot_product = torch.sum(u1 * u2) # # Ensure the dot product is within the valid range [-1, 1] # dot_product = torch.clamp(dot_product, -1.0, 1.0) # # Compute the angle between the vectors # theta = torch.acos(dot_product) # # Compute the coefficients for the interpolation # sin_theta = torch.sin(theta) # if sin_theta == 0: # # Vectors are parallel, return a linear interpolation # return u1 + t * (u2 - u1) # s1 = torch.sin((1 - t) * theta) / sin_theta # s2 = torch.sin(t * theta) / sin_theta # # Perform the interpolation # 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] # # Initialize the interpolated vector with the first vector # interpolated_vector = vectors[0] # # Perform Slerp iteratively # for i in range(1, n): # # Perform Slerp between the current interpolated vector and the next vector # t = t_values[i] / (t_values[i] + t_values[i-1]) # interpolated_vector = slerp(interpolated_vector, vectors[i], t) # return interpolated_vector # @torch.no_grad # def get_mask_from_img_sam1(mobilesamv2, yolov8, sam1_image, yolov8_image, original_size, input_size, transform, device): # sam_mask=[] # img_area = original_size[0] * original_size[1] # obj_results = yolov8(yolov8_image,device=device,retina_masks=False,imgsz=1024,conf=0.25,iou=0.95,verbose=False) # input_boxes1 = obj_results[0].boxes.xyxy # input_boxes1 = input_boxes1.cpu().numpy() # input_boxes1 = transform.apply_boxes(input_boxes1, original_size) # input_boxes = torch.from_numpy(input_boxes1).to(device) # # obj_results = yolov8(yolov8_image,device=device,retina_masks=False,imgsz=512,conf=0.25,iou=0.9,verbose=False) # # input_boxes2 = obj_results[0].boxes.xyxy # # input_boxes2 = input_boxes2.cpu().numpy() # # input_boxes2 = transform.apply_boxes(input_boxes2, original_size) # # input_boxes2 = torch.from_numpy(input_boxes2).to(device) # # input_boxes = torch.cat((input_boxes1, input_boxes2), dim=0) # input_image = mobilesamv2.preprocess(sam1_image) # image_embedding = mobilesamv2.image_encoder(input_image)['last_hidden_state'] # image_embedding=torch.repeat_interleave(image_embedding, 320, dim=0) # prompt_embedding=mobilesamv2.prompt_encoder.get_dense_pe() # prompt_embedding=torch.repeat_interleave(prompt_embedding, 320, dim=0) # for (boxes,) in batch_iterator(320, input_boxes): # with torch.no_grad(): # image_embedding=image_embedding[0:boxes.shape[0],:,:,:] # prompt_embedding=prompt_embedding[0:boxes.shape[0],:,:,:] # sparse_embeddings, dense_embeddings = mobilesamv2.prompt_encoder( # points=None, # boxes=boxes, # masks=None,) # low_res_masks, _ = mobilesamv2.mask_decoder( # image_embeddings=image_embedding, # image_pe=prompt_embedding, # sparse_prompt_embeddings=sparse_embeddings, # dense_prompt_embeddings=dense_embeddings, # multimask_output=False, # simple_type=True, # ) # low_res_masks=mobilesamv2.postprocess_masks(low_res_masks, input_size, original_size) # sam_mask_pre = (low_res_masks > mobilesamv2.mask_threshold) # for mask in sam_mask_pre: # if mask.sum() / img_area > 0.002: # sam_mask.append(mask.squeeze(1)) # sam_mask=torch.cat(sam_mask) # sorted_sam_mask = sorted(sam_mask, key=(lambda x: x.sum()), reverse=True) # keep = mask_nms(sorted_sam_mask) # ret_mask = filter(sorted_sam_mask, keep) # return ret_mask # @torch.no_grad # def get_cog_feats(images, device): # cog_seg_maps = [] # rev_cog_seg_maps = [] # inference_state = pe3r.sam2.init_state(images=images.sam2_images, video_height=images.sam2_video_size[0], video_width=images.sam2_video_size[1]) # mask_num = 0 # sam1_images = images.sam1_images # sam1_images_size = images.sam1_images_size # np_images = images.np_images # np_images_size = images.np_images_size # sam1_masks = get_mask_from_img_sam1(pe3r.mobilesamv2, pe3r.yolov8, sam1_images[0], np_images[0], np_images_size[0], sam1_images_size[0], images.sam1_transform, device) # for mask in sam1_masks: # _, _, _ = pe3r.sam2.add_new_mask( # inference_state=inference_state, # frame_idx=0, # obj_id=mask_num, # mask=mask, # ) # mask_num += 1 # video_segments = {} # video_segments contains the per-frame segmentation results # for out_frame_idx, out_obj_ids, out_mask_logits in pe3r.sam2.propagate_in_video(inference_state): # sam2_masks = (out_mask_logits > 0.0).squeeze(1) # video_segments[out_frame_idx] = { # out_obj_id: sam2_masks[i].cpu().numpy() # for i, out_obj_id in enumerate(out_obj_ids) # } # if out_frame_idx == 0: # continue # sam1_masks = get_mask_from_img_sam1(pe3r.mobilesamv2, pe3r.yolov8, sam1_images[out_frame_idx], np_images[out_frame_idx], np_images_size[out_frame_idx], sam1_images_size[out_frame_idx], images.sam1_transform, device) # for sam1_mask in sam1_masks: # flg = 1 # for sam2_mask in sam2_masks: # # print(sam1_mask.shape, sam2_mask.shape) # area1 = sam1_mask.sum() # area2 = sam2_mask.sum() # intersection = (sam1_mask & sam2_mask).sum() # if min(intersection / area1, intersection / area2) > 0.25: # flg = 0 # break # if flg: # video_segments[out_frame_idx][mask_num] = sam1_mask.cpu().numpy() # mask_num += 1 # multi_view_clip_feats = torch.zeros((mask_num+1, 1024)) # multi_view_clip_feats_map = {} # multi_view_clip_area_map = {} # for now_frame in range(0, len(video_segments), 1): # image = np_images[now_frame] # seg_img_list = [] # out_obj_id_list = [] # out_obj_mask_list = [] # out_obj_area_list = [] # # NOTE: background: -1 # rev_seg_map = -np.ones(image.shape[:2], dtype=np.int64) # sorted_dict_items = sorted(video_segments[now_frame].items(), key=lambda x: np.count_nonzero(x[1]), reverse=False) # for out_obj_id, mask in sorted_dict_items: # if mask.sum() == 0: # continue # rev_seg_map[mask] = out_obj_id # rev_cog_seg_maps.append(rev_seg_map) # seg_map = -np.ones(image.shape[:2], dtype=np.int64) # sorted_dict_items = sorted(video_segments[now_frame].items(), key=lambda x: np.count_nonzero(x[1]), reverse=True) # for out_obj_id, mask in sorted_dict_items: # if mask.sum() == 0: # continue # box = np.int32(box_xyxy_to_xywh(mask_to_box(mask))) # if box[2] == 0 and box[3] == 0: # continue # # print(box) # seg_img = get_seg_img(mask, box, image) # pad_seg_img = cv2.resize(pad_img(seg_img), (256,256)) # seg_img_list.append(pad_seg_img) # seg_map[mask] = out_obj_id # out_obj_id_list.append(out_obj_id) # out_obj_area_list.append(np.count_nonzero(mask)) # out_obj_mask_list.append(mask) # if len(seg_img_list) == 0: # cog_seg_maps.append(seg_map) # continue # seg_imgs = np.stack(seg_img_list, axis=0) # b,H,W,3 # seg_imgs = torch.from_numpy(seg_imgs).permute(0,3,1,2) # / 255.0 # inputs = pe3r.siglip_processor(images=seg_imgs, return_tensors="pt") # inputs = {key: value.to(device) for key, value in inputs.items()} # image_features = pe3r.siglip.get_image_features(**inputs) # image_features = image_features / image_features.norm(dim=-1, keepdim=True) # image_features = image_features.detach().cpu() # for i in range(len(out_obj_mask_list)): # for j in range(i + 1, len(out_obj_mask_list)): # mask1 = out_obj_mask_list[i] # mask2 = out_obj_mask_list[j] # intersection = np.logical_and(mask1, mask2).sum() # area1 = out_obj_area_list[i] # area2 = out_obj_area_list[j] # if min(intersection / area1, intersection / area2) > 0.025: # conf1 = area1 / (area1 + area2) # # conf2 = area2 / (area1 + area2) # image_features[j] = slerp(image_features[j], image_features[i], conf1) # for i, clip_feat in enumerate(image_features): # id = out_obj_id_list[i] # if id in multi_view_clip_feats_map.keys(): # multi_view_clip_feats_map[id].append(clip_feat) # multi_view_clip_area_map[id].append(out_obj_area_list[i]) # else: # multi_view_clip_feats_map[id] = [clip_feat] # multi_view_clip_area_map[id] = [out_obj_area_list[i]] # cog_seg_maps.append(seg_map) # del image_features # for i in range(mask_num): # if i in multi_view_clip_feats_map.keys(): # clip_feats = multi_view_clip_feats_map[i] # mask_area = multi_view_clip_area_map[i] # multi_view_clip_feats[i] = slerp_multiple(torch.stack(clip_feats), np.stack(mask_area)) # else: # multi_view_clip_feats[i] = torch.zeros((1024)) # multi_view_clip_feats[mask_num] = torch.zeros((1024)) # return cog_seg_maps, rev_cog_seg_maps, multi_view_clip_feats @spaces.GPU(duration=180) def get_reconstructed_scene(outdir, device, silent, 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) # try: # cog_seg_maps, rev_cog_seg_maps, cog_feats = get_cog_feats(images, device) # imgs = load_images(images, rev_cog_seg_maps, size=512, verbose=not silent) # except Exception as e: # 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, pe3r.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 # # if mode == GlobalAlignerMode.PointCloudOptimizer: # 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)): # # print(imgs[i]['img'].shape, scene.imgs[i].shape, ImgNorm(scene.imgs[i])[None]) # imgs[i]['img'] = ImgNorm(scene_1.imgs[i])[None] # pairs = make_pairs(imgs, scene_graph=scenegraph_type, prefilter=None, symmetrize=True) # output = inference(pairs, pe3r.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 # # if mode == GlobalAlignerMode.PointCloudOptimizer: # 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, silent, scene, min_conf_thr, as_pointcloud, mask_sky, # clean_depth, transparent_cams, cam_size) # # also return rgb, depth and confidence imgs # # depth is normalized with the max value for all images # # we apply the jet colormap on the confidence maps # rgbimg = scene.imgs # depths = to_numpy(scene.get_depthmaps()) # confs = to_numpy([c for c in scene.im_conf]) # # confs = to_numpy([c for c in scene.conf_2]) # cmap = pl.get_cmap('jet') # depths_max = max([d.max() for d in depths]) # depths = [d / depths_max for d in depths] # confs_max = max([d.max() for d in confs]) # confs = [cmap(d / confs_max) for d in confs] # imgs = [] # for i in range(len(rgbimg)): # imgs.append(rgbimg[i]) # imgs.append(rgb(depths[i])) # imgs.append(rgb(confs[i])) # return scene, outfile, imgs # @spaces.GPU(duration=180) # def get_3D_object_from_scene(outdir, pe3r, silent, device, text, threshold, scene, min_conf_thr, as_pointcloud, # mask_sky, clean_depth, transparent_cams, cam_size): # texts = [text] # inputs = pe3r.siglip_tokenizer(text=texts, padding="max_length", return_tensors="pt") # inputs = {key: value.to(device) for key, value in inputs.items()} # with torch.no_grad(): # text_feats =pe3r.siglip.get_text_features(**inputs) # text_feats = text_feats / text_feats.norm(dim=-1, keepdim=True) # scene.render_image(text_feats, threshold) # scene.ori_imgs = scene.rendered_imgs # outfile = get_3D_model_from_scene(outdir, silent, scene, min_conf_thr, as_pointcloud, mask_sky, # clean_depth, transparent_cams, cam_size) # return outfile def set_scenegraph_options(inputfiles, winsize, refid, scenegraph_type): num_files = len(inputfiles) if inputfiles is not None else 1 max_winsize = max(1, math.ceil((num_files - 1) / 2)) if scenegraph_type == "swin": winsize = gradio.Slider(label="Scene Graph: Window Size", value=max_winsize, minimum=1, maximum=max_winsize, step=1, visible=True) refid = gradio.Slider(label="Scene Graph: Id", value=0, minimum=0, maximum=num_files - 1, step=1, visible=False) elif scenegraph_type == "oneref": winsize = gradio.Slider(label="Scene Graph: Window Size", value=max_winsize, minimum=1, maximum=max_winsize, step=1, visible=False) refid = gradio.Slider(label="Scene Graph: Id", value=0, minimum=0, maximum=num_files - 1, step=1, visible=True) else: winsize = gradio.Slider(label="Scene Graph: Window Size", value=max_winsize, minimum=1, maximum=max_winsize, step=1, visible=False) refid = gradio.Slider(label="Scene Graph: Id", value=0, minimum=0, maximum=num_files - 1, step=1, visible=False) return winsize, refid with tempfile.TemporaryDirectory(suffix='pe3r_gradio_demo') as tmpdirname: recon_fun = functools.partial(get_reconstructed_scene, tmpdirname, pe3r, device, silent) # model_from_scene_fun = functools.partial(get_3D_model_from_scene, tmpdirname, silent) # get_3D_object_from_scene_fun = functools.partial(get_3D_object_from_scene, tmpdirname, pe3r, silent, device) with gradio.Blocks(css=""".gradio-container {margin: 0 !important; min-width: 100%};""", title="PE3R Demo") as demo: # scene state is save so that you can change conf_thr, cam_size... without rerunning the inference scene = gradio.State(None) gradio.HTML('