import gradio as gr import spaces from gradio_litmodel3d import LitModel3D import os import shutil os.environ['SPCONV_ALGO'] = 'native' from typing import * import torch import numpy as np import imageio from easydict import EasyDict as edict from PIL import Image from trellis.pipelines import TrellisVGGTTo3DPipeline from trellis.representations import Gaussian, MeshExtractResult from trellis.utils import render_utils, postprocessing_utils from wheels.vggt.vggt.utils.load_fn import load_and_preprocess_images from wheels.vggt.vggt.utils.pose_enc import pose_encoding_to_extri_intri import open3d as o3d from torchvision import transforms as TF from PIL import Image import sys sys.path.append("wheels") from wheels.mast3r.model import AsymmetricMASt3R from wheels.mast3r.fast_nn import fast_reciprocal_NNs from wheels.dust3r.dust3r.inference import inference from wheels.dust3r.dust3r.utils.image import load_images_new from trellis.utils.general_utils import * import copy MAX_SEED = np.iinfo(np.int32).max TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp') # TMP_DIR = "tmp/Trellis-demo" # os.environ['GRADIO_TEMP_DIR'] = 'tmp' os.makedirs(TMP_DIR, exist_ok=True) device = "cuda" if torch.cuda.is_available() else "cpu" def start_session(req: gr.Request): user_dir = os.path.join(TMP_DIR, str(req.session_hash)) os.makedirs(user_dir, exist_ok=True) def end_session(req: gr.Request): user_dir = os.path.join(TMP_DIR, str(req.session_hash)) shutil.rmtree(user_dir) @spaces.GPU def preprocess_image(image: Image.Image) -> Image.Image: """ Preprocess the input image for 3D generation. This function is called when a user uploads an image or selects an example. It applies background removal and other preprocessing steps necessary for optimal 3D model generation. Args: image (Image.Image): The input image from the user Returns: Image.Image: The preprocessed image ready for 3D generation """ processed_image = pipeline.preprocess_image(image) return processed_image @spaces.GPU def preprocess_videos(video: str) -> List[Tuple[Image.Image, str]]: """ Preprocess the input video for multi-image 3D generation. This function is called when a user uploads a video. It extracts frames from the video and processes each frame to prepare them for the multi-image 3D generation pipeline. Args: video (str): The path to the input video file Returns: List[Tuple[Image.Image, str]]: The list of preprocessed images ready for 3D generation """ vid = imageio.get_reader(video, 'ffmpeg') fps = vid.get_meta_data()['fps'] images = [] for i, frame in enumerate(vid): if i % max(int(fps * 1), 1) == 0: img = Image.fromarray(frame) W, H = img.size img = img.resize((int(W / H * 512), 512)) images.append(img) vid.close() processed_images = [pipeline.preprocess_image(image) for image in images] return processed_images @spaces.GPU def preprocess_images(images: List[Tuple[Image.Image, str]]) -> List[Image.Image]: """ Preprocess a list of input images for multi-image 3D generation. This function is called when users upload multiple images in the gallery. It processes each image to prepare them for the multi-image 3D generation pipeline. Args: images (List[Tuple[Image.Image, str]]): The input images from the gallery Returns: List[Image.Image]: The preprocessed images ready for 3D generation """ images = [image[0] for image in images] processed_images = [pipeline.preprocess_image(image) for image in images] return processed_images def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict: return { 'gaussian': { **gs.init_params, '_xyz': gs._xyz.cpu().numpy(), '_features_dc': gs._features_dc.cpu().numpy(), '_scaling': gs._scaling.cpu().numpy(), '_rotation': gs._rotation.cpu().numpy(), '_opacity': gs._opacity.cpu().numpy(), }, 'mesh': { 'vertices': mesh.vertices.cpu().numpy(), 'faces': mesh.faces.cpu().numpy(), }, } def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]: gs = Gaussian( aabb=state['gaussian']['aabb'], sh_degree=state['gaussian']['sh_degree'], mininum_kernel_size=state['gaussian']['mininum_kernel_size'], scaling_bias=state['gaussian']['scaling_bias'], opacity_bias=state['gaussian']['opacity_bias'], scaling_activation=state['gaussian']['scaling_activation'], ) gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda') gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda') gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda') gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda') gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda') mesh = edict( vertices=torch.tensor(state['mesh']['vertices'], device='cuda'), faces=torch.tensor(state['mesh']['faces'], device='cuda'), ) return gs, mesh def get_seed(randomize_seed: bool, seed: int) -> int: """ Get the random seed for generation. This function is called by the generate button to determine whether to use a random seed or the user-specified seed value. Args: randomize_seed (bool): Whether to generate a random seed seed (int): The user-specified seed value Returns: int: The seed to use for generation """ return np.random.randint(0, MAX_SEED) if randomize_seed else seed def align_camera(num_frames, extrinsic, intrinsic, rend_extrinsics, rend_intrinsics): extrinsic_tmp = extrinsic.clone() camera_relative = torch.matmul(extrinsic_tmp[:num_frames,:3,:3].permute(0,2,1), extrinsic_tmp[num_frames:,:3,:3]) camera_relative_angle = torch.acos(((camera_relative[:,0,0] + camera_relative[:,1,1] + camera_relative[:,2,2] - 1) / 2).clamp(-1, 1)) idx = torch.argmin(camera_relative_angle) target_extrinsic = rend_extrinsics[idx:idx+1].clone() focal_x = intrinsic[:num_frames,0,0].mean() focal_y = intrinsic[:num_frames,1,1].mean() focal = (focal_x + focal_y) / 2 rend_focal = (rend_intrinsics[0][0,0] + rend_intrinsics[0][1,1]) * 518 / 2 focal_scale = rend_focal / focal target_intrinsic = intrinsic[num_frames:].clone() fxy = (target_intrinsic[:,0,0] + target_intrinsic[:,1,1]) / 2 * focal_scale target_intrinsic[:,0,0] = fxy target_intrinsic[:,1,1] = fxy return target_extrinsic, target_intrinsic def refine_pose_mast3r(rend_image_pil, target_image_pil, original_size, fxy, target_extrinsic, rend_depth): images_mast3r = load_images_new([rend_image_pil, target_image_pil], size=512, square_ok=True) with torch.no_grad(): output = inference([tuple(images_mast3r)], mast3r_model, device, batch_size=1, verbose=False) view1, pred1 = output['view1'], output['pred1'] view2, pred2 = output['view2'], output['pred2'] del output desc1, desc2 = pred1['desc'].squeeze(0).detach(), pred2['desc'].squeeze(0).detach() # find 2D-2D matches between the two images matches_im0, matches_im1 = fast_reciprocal_NNs(desc1, desc2, subsample_or_initxy1=8, device=device, dist='dot', block_size=2**13) # ignore small border around the edge H0, W0 = view1['true_shape'][0] valid_matches_im0 = (matches_im0[:, 0] >= 3) & (matches_im0[:, 0] < int(W0) - 3) & ( matches_im0[:, 1] >= 3) & (matches_im0[:, 1] < int(H0) - 3) H1, W1 = view2['true_shape'][0] valid_matches_im1 = (matches_im1[:, 0] >= 3) & (matches_im1[:, 0] < int(W1) - 3) & ( matches_im1[:, 1] >= 3) & (matches_im1[:, 1] < int(H1) - 3) valid_matches = valid_matches_im0 & valid_matches_im1 matches_im0, matches_im1 = matches_im0[valid_matches], matches_im1[valid_matches] scale_x = original_size[1] / W0.item() scale_y = original_size[0] / H0.item() for pixel in matches_im1: pixel[0] *= scale_x pixel[1] *= scale_y for pixel in matches_im0: pixel[0] *= scale_x pixel[1] *= scale_y depth_map = rend_depth[0] fx, fy, cx, cy = fxy.item(), fxy.item(), original_size[1]/2, original_size[0]/2 # Example values for focal lengths and principal point K = np.array([ [fx, 0, cx], [0, fy, cy], [0, 0, 1] ]) dist_eff = np.array([0,0,0,0], dtype=np.float32) predict_c2w_ini = np.linalg.inv(target_extrinsic[0].cpu().numpy()) predict_w2c_ini = target_extrinsic[0].cpu().numpy() initial_rvec, _ = cv2.Rodrigues(predict_c2w_ini[:3,:3].astype(np.float32)) initial_tvec = predict_c2w_ini[:3,3].astype(np.float32) K_inv = np.linalg.inv(K) height, width = depth_map.shape x_coords, y_coords = np.meshgrid(np.arange(width), np.arange(height)) x_flat = x_coords.flatten() y_flat = y_coords.flatten() depth_flat = depth_map.flatten() x_normalized = (x_flat - K[0, 2]) / K[0, 0] y_normalized = (y_flat - K[1, 2]) / K[1, 1] X_camera = depth_flat * x_normalized Y_camera = depth_flat * y_normalized Z_camera = depth_flat points_camera = np.vstack((X_camera, Y_camera, Z_camera, np.ones_like(X_camera))) points_world = predict_c2w_ini @ points_camera X_world = points_world[0, :] Y_world = points_world[1, :] Z_world = points_world[2, :] points_3D = np.vstack((X_world, Y_world, Z_world)) scene_coordinates_gs = points_3D.reshape(3, original_size[0], original_size[1]) points_3D_at_pixels = np.zeros((matches_im0.shape[0], 3)) for i, (x, y) in enumerate(matches_im0): points_3D_at_pixels[i] = scene_coordinates_gs[:, y, x] success, rvec, tvec, inliers = cv2.solvePnPRansac(points_3D_at_pixels.astype(np.float32), matches_im1.astype(np.float32), K, \ dist_eff,rvec=initial_rvec,tvec=initial_tvec, useExtrinsicGuess=True, reprojectionError=1.0,\ iterationsCount=2000,flags=cv2.SOLVEPNP_EPNP) R = perform_rodrigues_transformation(rvec) trans = -R.T @ np.matrix(tvec) predict_c2w_refine = np.eye(4) predict_c2w_refine[:3,:3] = R.T predict_c2w_refine[:3,3] = trans.reshape(3) target_extrinsic_final = torch.tensor(predict_c2w_refine).inverse().cuda()[None].float() return target_extrinsic_final def pointcloud_registration(rend_image_pil, target_image_pil, original_size, fxy, target_extrinsic, rend_depth, target_pointmap, down_pcd, pcd): images_mast3r = load_images_new([rend_image_pil, target_image_pil], size=512, square_ok=True) with torch.no_grad(): output = inference([tuple(images_mast3r)], mast3r_model, device, batch_size=1, verbose=False) view1, pred1 = output['view1'], output['pred1'] view2, pred2 = output['view2'], output['pred2'] del output desc1, desc2 = pred1['desc'].squeeze(0).detach(), pred2['desc'].squeeze(0).detach() # find 2D-2D matches between the two images matches_im0, matches_im1 = fast_reciprocal_NNs(desc1, desc2, subsample_or_initxy1=8, device=device, dist='dot', block_size=2**13) # ignore small border around the edge H0, W0 = view1['true_shape'][0] valid_matches_im0 = (matches_im0[:, 0] >= 3) & (matches_im0[:, 0] < int(W0) - 3) & ( matches_im0[:, 1] >= 3) & (matches_im0[:, 1] < int(H0) - 3) H1, W1 = view2['true_shape'][0] valid_matches_im1 = (matches_im1[:, 0] >= 3) & (matches_im1[:, 0] < int(W1) - 3) & ( matches_im1[:, 1] >= 3) & (matches_im1[:, 1] < int(H1) - 3) valid_matches = valid_matches_im0 & valid_matches_im1 matches_im0, matches_im1 = matches_im0[valid_matches], matches_im1[valid_matches] scale_x = original_size[1] / W0.item() scale_y = original_size[0] / H0.item() for pixel in matches_im1: pixel[0] *= scale_x pixel[1] *= scale_y for pixel in matches_im0: pixel[0] *= scale_x pixel[1] *= scale_y depth_map = rend_depth[0] fx, fy, cx, cy = fxy.item(), fxy.item(), original_size[1]/2, original_size[0]/2 # Example values for focal lengths and principal point K = np.array([ [fx, 0, cx], [0, fy, cy], [0, 0, 1] ]) dist_eff = np.array([0,0,0,0], dtype=np.float32) predict_c2w_ini = np.linalg.inv(target_extrinsic[0].cpu().numpy()) predict_w2c_ini = target_extrinsic[0].cpu().numpy() initial_rvec, _ = cv2.Rodrigues(predict_c2w_ini[:3,:3].astype(np.float32)) initial_tvec = predict_c2w_ini[:3,3].astype(np.float32) K_inv = np.linalg.inv(K) height, width = depth_map.shape x_coords, y_coords = np.meshgrid(np.arange(width), np.arange(height)) x_flat = x_coords.flatten() y_flat = y_coords.flatten() depth_flat = depth_map.flatten() x_normalized = (x_flat - K[0, 2]) / K[0, 0] y_normalized = (y_flat - K[1, 2]) / K[1, 1] X_camera = depth_flat * x_normalized Y_camera = depth_flat * y_normalized Z_camera = depth_flat points_camera = np.vstack((X_camera, Y_camera, Z_camera, np.ones_like(X_camera))) points_world = predict_c2w_ini @ points_camera X_world = points_world[0, :] Y_world = points_world[1, :] Z_world = points_world[2, :] points_3D = np.vstack((X_world, Y_world, Z_world)) scene_coordinates_gs = points_3D.reshape(3, original_size[0], original_size[1]) points_3D_at_pixels = np.zeros((matches_im0.shape[0], 3)) for i, (x, y) in enumerate(matches_im0): points_3D_at_pixels[i] = scene_coordinates_gs[:, y, x] points_3D_at_pixels_2 = np.zeros((matches_im1.shape[0], 3)) for i, (x, y) in enumerate(matches_im1): points_3D_at_pixels_2[i] = target_pointmap[:, y, x] dist_1 = np.linalg.norm(points_3D_at_pixels - points_3D_at_pixels.mean(axis=0), axis=1) scale_1 = dist_1[dist_1 < np.percentile(dist_1, 99)].mean() dist_2 = np.linalg.norm(points_3D_at_pixels_2 - points_3D_at_pixels_2.mean(axis=0), axis=1) scale_2 = dist_2[dist_2 < np.percentile(dist_2, 99)].mean() # scale_1 = np.linalg.norm(points_3D_at_pixels - points_3D_at_pixels.mean(axis=0), axis=1).mean() # scale_2 = np.linalg.norm(points_3D_at_pixels_2 - points_3D_at_pixels_2.mean(axis=0), axis=1).mean() points_3D_at_pixels_2 = points_3D_at_pixels_2 * (scale_1 / scale_2) pcd_1 = o3d.geometry.PointCloud() pcd_1.points = o3d.utility.Vector3dVector(points_3D_at_pixels) pcd_2 = o3d.geometry.PointCloud() pcd_2.points = o3d.utility.Vector3dVector(points_3D_at_pixels_2) indices = np.arange(points_3D_at_pixels.shape[0]) correspondences = np.stack([indices, indices], axis=1) correspondences = o3d.utility.Vector2iVector(correspondences) result = o3d.pipelines.registration.registration_ransac_based_on_correspondence( pcd_2, pcd_1, correspondences, 0.03, estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPoint(False), ransac_n=5, criteria=o3d.pipelines.registration.RANSACConvergenceCriteria(10000, 10000), ) transformation_matrix = result.transformation.copy() transformation_matrix[:3,:3] = transformation_matrix[:3,:3] * (scale_1 / scale_2) evaluation = o3d.pipelines.registration.evaluate_registration( down_pcd, pcd, 0.02, transformation_matrix ) return transformation_matrix, evaluation.fitness @spaces.GPU(duration=120) def generate_and_extract_glb( multiimages: List[Tuple[Image.Image, str]], seed: int, ss_guidance_strength: float, ss_sampling_steps: int, slat_guidance_strength: float, slat_sampling_steps: int, multiimage_algo: Literal["multidiffusion", "stochastic"], mesh_simplify: float, texture_size: int, refine: Literal["Yes", "No"], ss_refine: Literal["noise", "deltav", "No"], registration_num_frames: int, trellis_stage1_lr: float, trellis_stage1_start_t: float, trellis_stage2_lr: float, trellis_stage2_start_t: float, req: gr.Request, ) -> Tuple[dict, str, str, str]: """ Convert an image to a 3D model and extract GLB file. Args: image (Image.Image): The input image. multiimages (List[Tuple[Image.Image, str]]): The input images in multi-image mode. is_multiimage (bool): Whether is in multi-image mode. seed (int): The random seed. ss_guidance_strength (float): The guidance strength for sparse structure generation. ss_sampling_steps (int): The number of sampling steps for sparse structure generation. slat_guidance_strength (float): The guidance strength for structured latent generation. slat_sampling_steps (int): The number of sampling steps for structured latent generation. multiimage_algo (Literal["multidiffusion", "stochastic"]): The algorithm for multi-image generation. mesh_simplify (float): The mesh simplification factor. texture_size (int): The texture resolution. Returns: dict: The information of the generated 3D model. str: The path to the video of the 3D model. str: The path to the extracted GLB file. str: The path to the extracted GLB file (for download). """ user_dir = os.path.join(TMP_DIR, str(req.session_hash)) image_files = [image[0] for image in multiimages] # Generate 3D model outputs, coords, ss_noise = pipeline.run( image=image_files, seed=seed, formats=["gaussian", "mesh"], preprocess_image=False, sparse_structure_sampler_params={ "steps": ss_sampling_steps, "cfg_strength": ss_guidance_strength, }, slat_sampler_params={ "steps": slat_sampling_steps, "cfg_strength": slat_guidance_strength, }, mode=multiimage_algo, ) if refine == "Yes": try: images, alphas = load_and_preprocess_images(multiimages) images, alphas = images.to(device), alphas.to(device) with torch.no_grad(): with torch.cuda.amp.autocast(dtype=pipeline.VGGT_dtype): images = images[None] aggregated_tokens_list, ps_idx = pipeline.VGGT_model.aggregator(images) # Predict Cameras pose_enc = pipeline.VGGT_model.camera_head(aggregated_tokens_list)[-1] # Extrinsic and intrinsic matrices, following OpenCV convention (camera from world) extrinsic, intrinsic = pose_encoding_to_extri_intri(pose_enc, images.shape[-2:]) # Predict Point Cloud point_map, point_conf = pipeline.VGGT_model.point_head(aggregated_tokens_list, images, ps_idx) del aggregated_tokens_list mask = (alphas[:,0,...][...,None] > 0.8) conf_threshold = np.percentile(point_conf.cpu().numpy(), 50) confidence_mask = (point_conf[0] > conf_threshold) & (point_conf[0] > 1e-5) mask = mask & confidence_mask[...,None] point_map_by_unprojection = point_map[0] point_map_clean = point_map_by_unprojection[mask[...,0]] center_point = point_map_clean.mean(0) scale = np.percentile((point_map_clean - center_point[None]).norm(dim=-1).cpu().numpy(), 98) outlier_mask = (point_map_by_unprojection - center_point[None]).norm(dim=-1) <= scale final_mask = mask & outlier_mask[...,None] point_map_perframe = (point_map_by_unprojection - center_point[None, None, None]) / (2 * scale) point_map_perframe[~final_mask[...,0]] = 127/255 point_map_perframe = point_map_perframe.permute(0,3,1,2) images = images[0].permute(0,2,3,1) images[~(alphas[:,0,...][...,None] > 0.8)[...,0]] = 0. input_images = images.permute(0,3,1,2).clone() vggt_extrinsic = extrinsic[0] vggt_extrinsic = torch.cat([vggt_extrinsic, torch.tensor([[[0,0,0,1]]]).repeat(vggt_extrinsic.shape[0], 1, 1).to(vggt_extrinsic)], dim=1) vggt_intrinsic = intrinsic[0] vggt_intrinsic[:,:2] = vggt_intrinsic[:,:2] / 518 vggt_extrinsic[:,:3,3] = (torch.matmul(vggt_extrinsic[:,:3,:3], center_point[None,:,None].float())[...,0] + vggt_extrinsic[:,:3,3]) / (2 * scale) pointcloud = point_map_perframe.permute(0,2,3,1)[final_mask[...,0]] idxs = torch.randperm(pointcloud.shape[0])[:min(50000, pointcloud.shape[0])] pcd = o3d.geometry.PointCloud() pcd.points = o3d.utility.Vector3dVector(pointcloud[idxs].cpu().numpy()) cl, ind = pcd.remove_statistical_outlier(nb_neighbors=30, std_ratio=3.0) inlier_cloud = pcd.select_by_index(ind) outlier_cloud = pcd.select_by_index(ind, invert=True) distance = np.array(inlier_cloud.points) - np.array(inlier_cloud.points).mean(axis=0)[None] scale = np.percentile(np.linalg.norm(distance, axis=1), 97) voxel_size = 1/64*scale*2 down_pcd = inlier_cloud.voxel_down_sample(voxel_size) torch.cuda.empty_cache() video, rend_extrinsics, rend_intrinsics = render_utils.render_multiview(outputs['gaussian'][0], num_frames=registration_num_frames) rend_extrinsics = torch.stack(rend_extrinsics, dim=0) rend_intrinsics = torch.stack(rend_intrinsics, dim=0) target_extrinsics = [] target_intrinsics = [] target_transforms = [] target_fitnesses = [] pcd = o3d.geometry.PointCloud() mesh = outputs['mesh'][0] idxs = torch.randperm(mesh.vertices.shape[0])[:min(50000, mesh.vertices.shape[0])] pcd.points = o3d.utility.Vector3dVector(mesh.vertices[idxs].cpu().numpy()) distance = np.array(pcd.points) - np.array(pcd.points).mean(axis=0)[None] scale = np.linalg.norm(distance, axis=1).max() voxel_size = 1/64*scale*2 pcd = pcd.voxel_down_sample(voxel_size) # pcd.points = o3d.utility.Vector3dVector((coords[:,1:].cpu().numpy() + 0.5) / 64 - 0.5) for k in range(len(image_files)): images = torch.stack([TF.ToTensor()(render_image) for render_image in video['color']] + [TF.ToTensor()(image_files[k].convert("RGB"))], dim=0) # if len(images) == 0: with torch.no_grad(): with torch.cuda.amp.autocast(dtype=pipeline.VGGT_dtype): # predictions = vggt_model(images.cuda()) aggregated_tokens_list, ps_idx = pipeline.VGGT_model.aggregator(images[None].cuda()) pose_enc = pipeline.VGGT_model.camera_head(aggregated_tokens_list)[-1] extrinsic, intrinsic = pose_encoding_to_extri_intri(pose_enc, images.shape[-2:]) extrinsic, intrinsic = extrinsic[0], intrinsic[0] extrinsic = torch.cat([extrinsic, torch.tensor([0,0,0,1])[None,None].repeat(extrinsic.shape[0], 1, 1).to(extrinsic.device)], dim=1) del aggregated_tokens_list, ps_idx target_extrinsic, target_intrinsic = align_camera(registration_num_frames, extrinsic, intrinsic, rend_extrinsics, rend_intrinsics) fxy = target_intrinsic[:,0,0] target_intrinsic_tmp = target_intrinsic.clone() target_intrinsic_tmp[:,:2] = target_intrinsic_tmp[:,:2] / 518 target_extrinsic_list = [target_extrinsic] iou_list = [] iterations = 3 for i in range(iterations + 1): j = 0 rend = render_utils.render_frames(outputs['gaussian'][0], target_extrinsic, target_intrinsic_tmp, {'resolution': 518, 'bg_color': (0, 0, 0)}, need_depth=True) rend_image = rend['color'][j] # (518, 518, 3) rend_depth = rend['depth'][j] # (3, 518, 518) depth_single = rend_depth[0].astype(np.float32) # (H, W) mask = (depth_single != 0).astype(np.uint8) # kernel = np.ones((3, 3), np.uint8) mask_eroded = cv2.erode(mask, kernel, iterations=3) depth_eroded = depth_single * mask_eroded rend_depth_eroded = np.stack([depth_eroded]*3, axis=0) rend_image = torch.tensor(rend_image).permute(2,0,1) / 255 target_image = images[registration_num_frames:].to(target_extrinsic.device)[j] original_size = (rend_image.shape[1], rend_image.shape[2]) import torchvision torchvision.utils.save_image(rend_image, 'rend_image_{}.png'.format(k)) torchvision.utils.save_image(target_image, 'target_image_{}.png'.format(k)) mask_rend = (rend_image.detach().cpu() > 0).any(dim=0) mask_target = (target_image.detach().cpu() > 0).any(dim=0) intersection = (mask_rend & mask_target).sum().item() union = (mask_rend | mask_target).sum().item() iou = intersection / union if union > 0 else 0.0 iou_list.append(iou) if i == iterations: break rend_image = rend_image * torch.from_numpy(mask_eroded[None]).to(rend_image.device) rend_image_pil = Image.fromarray((rend_image.permute(1,2,0).cpu().numpy() * 255).astype(np.uint8)) target_image_pil = Image.fromarray((target_image.permute(1,2,0).cpu().numpy() * 255).astype(np.uint8)) target_extrinsic[j:j+1] = refine_pose_mast3r(rend_image_pil, target_image_pil, original_size, fxy[j:j+1], target_extrinsic[j:j+1], rend_depth_eroded) target_extrinsic_list.append(target_extrinsic[j:j+1]) idx = iou_list.index(max(iou_list)) target_extrinsic[j:j+1] = target_extrinsic_list[idx] target_transform, fitness = pointcloud_registration(rend_image_pil, target_image_pil, original_size, fxy[j:j+1], target_extrinsic[j:j+1], \ rend_depth_eroded, point_map_perframe[k].cpu().numpy(), down_pcd, pcd) target_transforms.append(target_transform) target_fitnesses.append(fitness) target_extrinsics.append(target_extrinsic[j:j+1]) target_intrinsics.append(target_intrinsic_tmp[j:j+1]) target_extrinsics = torch.cat(target_extrinsics, dim=0) target_intrinsics = torch.cat(target_intrinsics, dim=0) target_fitnesses_filtered = [x for x in target_fitnesses if x < 1] idx = target_fitnesses.index(max(target_fitnesses_filtered)) target_transform = target_transforms[idx] down_pcd_align = copy.deepcopy(down_pcd).transform(target_transform) # pcd = o3d.geometry.PointCloud() # pcd.points = o3d.utility.Vector3dVector(coords[:,1:].cpu().numpy() / 64 - 0.5) reg_p2p = o3d.pipelines.registration.registration_icp( down_pcd_align, pcd, 0.02, np.eye(4), o3d.pipelines.registration.TransformationEstimationPointToPoint(with_scaling=True), o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration = 10000)) down_pcd_align_2 = copy.deepcopy(down_pcd_align).transform(reg_p2p.transformation) input_points = torch.tensor(np.asarray(down_pcd_align_2.points)).to(extrinsic.device).float() input_points = ((input_points + 0.5).clip(0, 1) * 64 - 0.5).to(torch.int32) outputs = pipeline.run_refine( image=image_files, ss_learning_rate=trellis_stage1_lr, ss_start_t=trellis_stage1_start_t, apperance_learning_rate=trellis_stage2_lr, apperance_start_t=trellis_stage2_start_t, extrinsics=target_extrinsics, intrinsics=target_intrinsics, ss_noise=ss_noise, input_points=input_points, ss_refine_type = ss_refine, coords=coords if ss_refine == "No" else None, seed=seed, formats=["mesh", "gaussian"], sparse_structure_sampler_params={ "steps": ss_sampling_steps, "cfg_strength": ss_guidance_strength, }, slat_sampler_params={ "steps": slat_sampling_steps, "cfg_strength": slat_guidance_strength, }, mode=multiimage_algo, ) except Exception as e: print(f"Error during refinement: {e}") # Render video # import uuid # output_id = str(uuid.uuid4()) # os.makedirs(f"{TMP_DIR}/{output_id}", exist_ok=True) # video_path = f"{TMP_DIR}/{output_id}/preview.mp4" # glb_path = f"{TMP_DIR}/{output_id}/mesh.glb" video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color'] video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal'] video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))] video_path = os.path.join(user_dir, 'sample.mp4') imageio.mimsave(video_path, video, fps=15) # Extract GLB gs = outputs['gaussian'][0] mesh = outputs['mesh'][0] glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False) glb_path = os.path.join(user_dir, 'sample.glb') glb.export(glb_path) # Pack state for optional Gaussian extraction state = pack_state(gs, mesh) torch.cuda.empty_cache() return state, video_path, glb_path, glb_path @spaces.GPU def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]: """ Extract a Gaussian splatting file from the generated 3D model. This function is called when the user clicks "Extract Gaussian" button. It converts the 3D model state into a .ply file format containing Gaussian splatting data for advanced 3D applications. Args: state (dict): The state of the generated 3D model containing Gaussian data req (gr.Request): Gradio request object for session management Returns: Tuple[str, str]: Paths to the extracted Gaussian file (for display and download) """ user_dir = os.path.join(TMP_DIR, str(req.session_hash)) gs, _ = unpack_state(state) gaussian_path = os.path.join(user_dir, 'sample.ply') gs.save_ply(gaussian_path) torch.cuda.empty_cache() return gaussian_path, gaussian_path def prepare_multi_example() -> List[Image.Image]: multi_case = list(set([i.split('_')[0] for i in os.listdir("assets/example_multi_image")])) images = [] for case in multi_case: _images = [] for i in range(1, 9): if os.path.exists(f'assets/example_multi_image/{case}_{i}.png'): img = Image.open(f'assets/example_multi_image/{case}_{i}.png') W, H = img.size img = img.resize((int(W / H * 512), 512)) _images.append(np.array(img)) if len(_images) > 0: images.append(Image.fromarray(np.concatenate(_images, axis=1))) return images def split_image(image: Image.Image) -> List[Image.Image]: """ Split a multi-view image into separate view images. This function is called when users select multi-image examples that contain multiple views in a single concatenated image. It automatically splits them based on alpha channel boundaries and preprocesses each view. Args: image (Image.Image): A concatenated image containing multiple views Returns: List[Image.Image]: List of individual preprocessed view images """ image = np.array(image) alpha = image[..., 3] alpha = np.any(alpha>0, axis=0) start_pos = np.where(~alpha[:-1] & alpha[1:])[0].tolist() end_pos = np.where(alpha[:-1] & ~alpha[1:])[0].tolist() images = [] for s, e in zip(start_pos, end_pos): images.append(Image.fromarray(image[:, s:e+1])) return [preprocess_image(image) for image in images] # Create interface demo = gr.Blocks( title="ReconViaGen", css=""" .slider .inner { width: 5px; background: #FFF; } .viewport { aspect-ratio: 4/3; } .tabs button.selected { font-size: 20px !important; color: crimson !important; } h1, h2, h3 { text-align: center; display: block; } .md_feedback li { margin-bottom: 0px !important; } """ ) with demo: gr.Markdown(""" # 💻 ReconViaGen

badge-github-stars

✨This demo is partial. We will release the whole model later. Stay tuned!✨ """) with gr.Row(): with gr.Column(): with gr.Tabs() as input_tabs: with gr.Tab(label="Input Video or Images", id=0) as multiimage_input_tab: input_video = gr.Video(label="Upload Video", interactive=True, height=300) image_prompt = gr.Image(label="Image Prompt", format="png", visible=False, image_mode="RGBA", type="pil", height=300) multiimage_prompt = gr.Gallery(label="Image Prompt", format="png", type="pil", height=300, columns=3) gr.Markdown(""" Input different views of the object in separate images. """) with gr.Accordion(label="Generation Settings", open=False): seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1) randomize_seed = gr.Checkbox(label="Randomize Seed", value=False) gr.Markdown("Stage 1: Sparse Structure Generation") with gr.Row(): ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=7.5, step=0.1) ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=30, step=1) gr.Markdown("Stage 2: Structured Latent Generation") with gr.Row(): slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1) slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1) multiimage_algo = gr.Radio(["stochastic", "multidiffusion"], label="Multi-image Algorithm", value="multidiffusion") refine = gr.Radio(["Yes", "No"], label="Refinement of Not", value="Yes") ss_refine = gr.Radio(["noise", "deltav", "No"], label="Sparse Structure refinement of not", value="No") registration_num_frames = gr.Slider(20, 50, label="Number of frames in registration", value=30, step=1) trellis_stage1_lr = gr.Slider(1e-4, 1., label="trellis_stage1_lr", value=1e-1, step=5e-4) trellis_stage1_start_t = gr.Slider(0., 1., label="trellis_stage1_start_t", value=0.5, step=0.01) trellis_stage2_lr = gr.Slider(1e-4, 1., label="trellis_stage2_lr", value=1e-1, step=5e-4) trellis_stage2_start_t = gr.Slider(0., 1., label="trellis_stage2_start_t", value=0.5, step=0.01) with gr.Accordion(label="GLB Extraction Settings", open=False): mesh_simplify = gr.Slider(0.9, 0.98, label="Simplify", value=0.95, step=0.01) texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512) generate_btn = gr.Button("Generate & Extract GLB", variant="primary") extract_gs_btn = gr.Button("Extract Gaussian", interactive=False) gr.Markdown(""" *NOTE: Gaussian file can be very large (~50MB), it will take a while to display and download.* """) with gr.Column(): video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300) model_output = LitModel3D(label="Extracted GLB/Gaussian", exposure=10.0, height=300) with gr.Row(): download_glb = gr.DownloadButton(label="Download GLB", interactive=False) download_gs = gr.DownloadButton(label="Download Gaussian", interactive=False) output_buf = gr.State() # Example images at the bottom of the page with gr.Row() as multiimage_example: examples_multi = gr.Examples( examples=prepare_multi_example(), inputs=[image_prompt], fn=split_image, outputs=[multiimage_prompt], run_on_click=True, examples_per_page=8, ) # Handlers demo.load(start_session) demo.unload(end_session) input_video.upload( preprocess_videos, inputs=[input_video], outputs=[multiimage_prompt], ) input_video.clear( lambda: tuple([None, None]), outputs=[input_video, multiimage_prompt], ) multiimage_prompt.upload( preprocess_images, inputs=[multiimage_prompt], outputs=[multiimage_prompt], ) generate_btn.click( get_seed, inputs=[randomize_seed, seed], outputs=[seed], ).then( generate_and_extract_glb, inputs=[multiimage_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps, multiimage_algo, mesh_simplify, texture_size, refine, ss_refine, registration_num_frames, trellis_stage1_lr, trellis_stage1_start_t, trellis_stage2_lr, trellis_stage2_start_t], outputs=[output_buf, video_output, model_output, download_glb], ).then( lambda: tuple([gr.Button(interactive=True), gr.Button(interactive=True)]), outputs=[extract_gs_btn, download_glb], ) video_output.clear( lambda: tuple([gr.Button(interactive=False), gr.Button(interactive=False), gr.Button(interactive=False)]), outputs=[extract_gs_btn, download_glb, download_gs], ) extract_gs_btn.click( extract_gaussian, inputs=[output_buf], outputs=[model_output, download_gs], ).then( lambda: gr.Button(interactive=True), outputs=[download_gs], ) model_output.clear( lambda: tuple([gr.Button(interactive=False), gr.Button(interactive=False)]), outputs=[download_glb, download_gs], ) # Launch the Gradio app if __name__ == "__main__": pipeline = TrellisVGGTTo3DPipeline.from_pretrained("Stable-X/trellis-vggt-v0-2") # pipeline = TrellisVGGTTo3DPipeline.from_pretrained("weights/trellis-vggt-v0-1") pipeline.cuda() pipeline.VGGT_model.cuda() pipeline.birefnet_model.cuda() pipeline.dreamsim_model.cuda() mast3r_model = AsymmetricMASt3R.from_pretrained("naver/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric").cuda().eval() # mast3r_model = AsymmetricMASt3R.from_pretrained("weights/MAST3R/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth").cuda().eval() demo.launch()