diff --git a/3d_generation/TripoSR/README.md b/3d_generation/TripoSR/README.md index 8dac32b..4e247ea 100644 --- a/3d_generation/TripoSR/README.md +++ b/3d_generation/TripoSR/README.md @@ -1,4 +1,4 @@ -# TripoSR +# TripoSR
Teaser Video @@ -38,14 +38,9 @@ python run.py examples/chair.png --output-dir output/ ``` This will save the reconstructed 3D model to `output/`. You can also specify more than one image path separated by spaces. The default options takes about **6GB VRAM** for a single image input. -For defailed usage of this script, use `python run.py --help`. +For detailed usage of this script, use `python run.py --help`. ### Local Gradio App -Install Gradio: -```sh -pip install gradio -``` -Start the Gradio App: ```sh python gradio_app.py ``` diff --git a/3d_generation/TripoSR/examples/chair.png b/3d_generation/TripoSR/examples/chair.png index c691634..fee2a21 100644 Binary files a/3d_generation/TripoSR/examples/chair.png and b/3d_generation/TripoSR/examples/chair.png differ diff --git a/3d_generation/TripoSR/examples/flamingo.png b/3d_generation/TripoSR/examples/flamingo.png index 2f2be77..c01ce3d 100644 Binary files a/3d_generation/TripoSR/examples/flamingo.png and b/3d_generation/TripoSR/examples/flamingo.png differ diff --git a/3d_generation/TripoSR/examples/hamburger.png b/3d_generation/TripoSR/examples/hamburger.png index c07b42d..166ef03 100644 Binary files a/3d_generation/TripoSR/examples/hamburger.png and b/3d_generation/TripoSR/examples/hamburger.png differ diff --git a/3d_generation/TripoSR/examples/horse.png b/3d_generation/TripoSR/examples/horse.png index 691a27c..97facd4 100644 Binary files a/3d_generation/TripoSR/examples/horse.png and b/3d_generation/TripoSR/examples/horse.png differ diff --git a/3d_generation/TripoSR/examples/iso_house.png b/3d_generation/TripoSR/examples/iso_house.png index f54335c..717170f 100644 Binary files a/3d_generation/TripoSR/examples/iso_house.png and b/3d_generation/TripoSR/examples/iso_house.png differ diff --git a/3d_generation/TripoSR/examples/marble.png b/3d_generation/TripoSR/examples/marble.png index cd426b0..0d72674 100644 Binary files a/3d_generation/TripoSR/examples/marble.png and b/3d_generation/TripoSR/examples/marble.png differ diff --git a/3d_generation/TripoSR/examples/mario.png b/3d_generation/TripoSR/examples/mario.png new file mode 100644 index 0000000..0919e91 Binary files /dev/null and b/3d_generation/TripoSR/examples/mario.png differ diff --git a/3d_generation/TripoSR/examples/police_woman.png b/3d_generation/TripoSR/examples/police_woman.png index 4bc99ab..adda32a 100644 Binary files a/3d_generation/TripoSR/examples/police_woman.png and b/3d_generation/TripoSR/examples/police_woman.png differ diff --git a/3d_generation/TripoSR/examples/poly_fox.png b/3d_generation/TripoSR/examples/poly_fox.png index 83cee5d..3a137fc 100644 Binary files a/3d_generation/TripoSR/examples/poly_fox.png and b/3d_generation/TripoSR/examples/poly_fox.png differ diff --git a/3d_generation/TripoSR/examples/robot.png b/3d_generation/TripoSR/examples/robot.png index 9f12ee4..31ce968 100644 Binary files a/3d_generation/TripoSR/examples/robot.png and b/3d_generation/TripoSR/examples/robot.png differ diff --git a/3d_generation/TripoSR/examples/stripes.png b/3d_generation/TripoSR/examples/stripes.png index 36bc9eb..125cac6 100644 Binary files a/3d_generation/TripoSR/examples/stripes.png and b/3d_generation/TripoSR/examples/stripes.png differ diff --git a/3d_generation/TripoSR/examples/teapot.png b/3d_generation/TripoSR/examples/teapot.png index 5f8514c..864694f 100644 Binary files a/3d_generation/TripoSR/examples/teapot.png and b/3d_generation/TripoSR/examples/teapot.png differ diff --git a/3d_generation/TripoSR/examples/tiger_girl.png b/3d_generation/TripoSR/examples/tiger_girl.png index 84e7e1a..4288fc7 100644 Binary files a/3d_generation/TripoSR/examples/tiger_girl.png and b/3d_generation/TripoSR/examples/tiger_girl.png differ diff --git a/3d_generation/TripoSR/examples/unicorn.png b/3d_generation/TripoSR/examples/unicorn.png index 2b6386c..d3be7f6 100644 Binary files a/3d_generation/TripoSR/examples/unicorn.png and b/3d_generation/TripoSR/examples/unicorn.png differ diff --git a/3d_generation/TripoSR/gradio_app.py b/3d_generation/TripoSR/gradio_app.py index b7b9616..9f67db7 100644 --- a/3d_generation/TripoSR/gradio_app.py +++ b/3d_generation/TripoSR/gradio_app.py @@ -84,7 +84,7 @@ with gr.Blocks(title="TripoSR") as interface: **Tips:** 1. If you find the result is unsatisfied, please try to change the foreground ratio. It might improve the results. - 2. You can disable "Remove Background" for the provided examples since they have been already preprocessed. + 2. It's better to disable "Remove Background" for the provided examples (except fot the last one) since they have been already preprocessed. 3. Otherwise, please disable "Remove Background" option only if your input image is RGBA with transparent background, image contents are centered and occupy more than 70% of image width or height. """ ) @@ -114,7 +114,7 @@ with gr.Blocks(title="TripoSR") as interface: mc_resolution = gr.Slider( label="Marching Cubes Resolution", minimum=32, - maximum=320, + maximum=1024, value=256, step=32 ) @@ -148,7 +148,7 @@ with gr.Blocks(title="TripoSR") as interface: "examples/iso_house.png", "examples/marble.png", "examples/police_woman.png", - "examples/captured_p.png", + "examples/captured.jpeg", ], inputs=[input_image], outputs=[processed_image, output_model_obj, output_model_glb], diff --git a/3d_generation/TripoSR/requirements.txt b/3d_generation/TripoSR/requirements.txt index 5fe390d..d368a9b 100644 --- a/3d_generation/TripoSR/requirements.txt +++ b/3d_generation/TripoSR/requirements.txt @@ -7,3 +7,4 @@ trimesh==4.0.5 rembg huggingface-hub imageio[ffmpeg] +gradio diff --git a/3d_generation/TripoSR/run.py b/3d_generation/TripoSR/run.py index 19ce20d..f27a8ce 100644 --- a/3d_generation/TripoSR/run.py +++ b/3d_generation/TripoSR/run.py @@ -3,100 +3,120 @@ import logging import os import sys import time -from PIL import Image import numpy as np -import rembg import torch +from PIL import Image +import rembg from tsr.system import TSR -from tsr.utils import remove_background, resize_foreground, save_video +from tsr.utils import remove_background, resize_foreground, save_video, to_gradio_3d_orientation -# Setup logging configuration -logging.basicConfig(format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO) - -# Print environment and GPU information -def print_environment_info(): - logging.info("Environment variables:") - for key, value in os.environ.items(): - logging.info(f"{key}: {value}") - - logging.info('========================================') - logging.info(f'Python interpreter: {sys.executable}') - logging.info(f'PyTorch version: {torch.__version__}') - logging.info(f'CUDA Available: {torch.cuda.is_available()}') - logging.info(f'CUDA Device count: {torch.cuda.device_count()}') - logging.info(f'CUDA threads for parallelizing CPU operations: {torch.get_num_threads()}') - logging.info(f'CUDA architectures library was compiled for: {torch.cuda.get_arch_list()}') - logging.info('========================================') - -print_environment_info() +def print_gpu_info(): + print('========================================') + print('Python interpreter:', sys.executable) + print('PyTorch version:', torch.__version__) + print('CUDA Available:', torch.cuda.is_available()) + print('CUDA Device count:', torch.cuda.device_count()) + print('CUDA threads for parallelizing CPU operations:', torch.get_num_threads()) + print('CUDA architectures library was compiled for:', torch.cuda.get_arch_list()) + print('========================================', flush=True) class Timer: def __init__(self): - self.start_times = {} + self.items = {} + self.time_scale = 1000.0 # ms + self.time_unit = "ms" - def start(self, name: str): - self.start_times[name] = time.time() + def start(self, name: str) -> None: + if torch.cuda.is_available(): + torch.cuda.synchronize() + self.items[name] = time.time() + logging.info(f"{name} ...") - def end(self, name: str): - if name in self.start_times: - elapsed_time = time.time() - self.start_times[name] - logging.info(f"{name} finished in {elapsed_time * 1000:.2f}ms.") - del self.start_times[name] + def end(self, name: str) -> None: + if name not in self.items: + return + if torch.cuda.is_available(): + torch.cuda.synchronize() + elapsed = (time.time() - self.items.pop(name)) * self.time_scale + logging.info(f"{name} finished in {elapsed:.2f}{self.time_unit}.") -timer = Timer() +def main(): + logging.basicConfig(format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO) -def process_image(image_path, remove_bg=True, foreground_ratio=0.85): - image = Image.open(image_path).convert("RGB") - if remove_bg: - rembg_session = rembg.new_session() - image = remove_background(image, rembg_session) - image = resize_foreground(image, foreground_ratio) - image = np.array(image).astype(np.float32) / 255.0 - image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5 # Blend with grey background - return Image.fromarray((image * 255.0).astype(np.uint8)) + print_gpu_info() + + parser = argparse.ArgumentParser() + parser.add_argument("--input", type=str, required=True, nargs="+", help="Path to input image(s).") + parser.add_argument("--device", default="cuda:0", type=str, help="Device to use. Defaults to 'cuda:0' if available.") + parser.add_argument("--pretrained-model-name-or-path", default="stabilityai/TripoSR", type=str, help="Pretrained model path or Hugging Face model ID.") + parser.add_argument("--chunk-size", default=8192, type=int, help="Chunk size for processing. Smaller values use less VRAM.") + parser.add_argument("--mc-resolution", default=256, type=int, help="Resolution for the marching cubes algorithm.") + parser.add_argument("--no-remove-bg", action="store_true", help="Disable automatic background removal.") + parser.add_argument("--foreground-ratio", default=0.85, type=float, help="Foreground to image size ratio.") + parser.add_argument("--output-dir", default="output/", type=str, help="Directory to save the output.") + parser.add_argument("--model-save-format", default="glb", choices=["obj", "glb"], help="Mesh save format.") + parser.add_argument("--render", action="store_true", help="Render a video from the model output.") + args = parser.parse_args() -def main(args): output_dir = args.output_dir os.makedirs(output_dir, exist_ok=True) device = args.device if torch.cuda.is_available() else "cpu" + timer = Timer() timer.start("Initializing model") model = TSR.from_pretrained(args.pretrained_model_name_or_path, config_name="config.yaml", weight_name="model.ckpt") model.renderer.set_chunk_size(args.chunk_size) model.to(device) timer.end("Initializing model") - for i, image_path in enumerate(args.image): - timer.start(f"Processing image {i+1}") - processed_image = process_image(image_path, not args.no_remove_bg, args.foreground_ratio) - processed_image_path = os.path.join(output_dir, f"processed_{i}.png") - processed_image.save(processed_image_path) - timer.end(f"Processing image {i+1}") + timer.start("Processing images") + images = [] + if args.no_remove_bg: + rembg_session = None + else: + rembg_session = rembg.new_session() - timer.start(f"Running model on image {i+1}") + for i, image_path in enumerate(args.input): + if args.no_remove_bg: + image = np.array(Image.open(image_path).convert("RGB")) + else: + image = remove_background(Image.open(image_path), rembg_session) + image = resize_foreground(image, args.foreground_ratio) + image = np.array(image).astype(np.float32) / 255.0 + image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5 + image = Image.fromarray((image * 255.0).astype(np.uint8)) + if not os.path.exists(os.path.join(output_dir, str(i))): + os.makedirs(os.path.join(output_dir, str(i))) + image.save(os.path.join(output_dir, str(i), "input.png")) + images.append(image) + timer.end("Processing images") + + for i, image in enumerate(images): + logging.info(f"Running image {i + 1}/{len(images)} ...") + + timer.start("Running model") with torch.no_grad(): - scene_codes = model([np.array(processed_image).astype(np.float32) / 255.0], device=device) - timer.end(f"Running model on image {i+1}") + scene_codes = model([image], device=device) + timer.end("Running model") - timer.start(f"Exporting mesh for image {i+1}") + if args.render: + timer.start("Rendering") + render_images = model.render(scene_codes, n_views=30, return_type="pil") + for ri, render_image in enumerate(render_images[0]): + render_image.save(os.path.join(output_dir, str(i), f"render_{ri:03d}.png")) + save_video(render_images[0], os.path.join(output_dir, str(i), "render.mp4"), fps=30) + timer.end("Rendering") + + timer.start("Exporting mesh") meshes = model.extract_mesh(scene_codes, resolution=args.mc_resolution) - mesh_export_path = os.path.join(output_dir, f"mesh_{i}.{args.model_save_format}") - meshes[0].export(mesh_export_path, file_type=args.model_save_format) - timer.end(f"Exporting mesh for image {i+1}") + meshes[0] = to_gradio_3d_orientation(meshes[0]) + mesh_path = os.path.join(output_dir, str(i), f"mesh.{args.model_save_format}") + meshes[0].export(mesh_path) + logging.info(f"Mesh saved to {mesh_path}") + timer.end("Exporting mesh") if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Run TSR model on input images.") - parser.add_argument("image", type=str, nargs="+", help="Path to input image(s).") - parser.add_argument("--device", default="cuda:0", type=str, help="Device to use. Defaults to 'cuda:0', falls back to 'cpu' if CUDA is not available.") - parser.add_argument("--pretrained-model-name-or-path", default="stabilityai/TripoSR", type=str, help="Pretrained model identifier or file path.") - parser.add_argument("--chunk-size", default=8192, type=int, help="Evaluation chunk size for rendering. Defaults to 8192.") - parser.add_argument("--mc-resolution", default=256, type=int, help="Resolution for marching cubes algorithm. Defaults to 256.") - parser.add_argument("--no-remove-bg", action="store_true", help="Don't remove background from input images.") - parser.add_argument("--foreground-ratio", default=0.85, type=float, help="Foreground ratio for resizing foreground. Defaults to 0.85.") - parser.add_argument("--output-dir", default="output/", type=str, help="Directory to save the output. Defaults to 'output/'.") - parser.add_argument("--model-save-format", default="obj", choices=["obj", "glb"], type=str, help="Format for saving the model. Supports 'obj' and 'glb'. Defaults to 'obj'.") - args = parser.parse_args() - - main(args) \ No newline at end of file + main() + diff --git a/3d_generation/TripoSR/tsr/models/nerf_renderer.py b/3d_generation/TripoSR/tsr/models/nerf_renderer.py index 0cf3ab6..b00224f 100644 --- a/3d_generation/TripoSR/tsr/models/nerf_renderer.py +++ b/3d_generation/TripoSR/tsr/models/nerf_renderer.py @@ -1,9 +1,10 @@ from dataclasses import dataclass -from typing import Dict +from typing import Dict, Optional import torch import torch.nn.functional as F from einops import rearrange, reduce +from torchmcubes import marching_cubes from ..utils import ( BaseModule, @@ -38,20 +39,78 @@ class TriplaneNeRFRenderer(BaseModule): ), "chunk_size must be a non-negative integer (0 for no chunking)." self.chunk_size = chunk_size + def interpolate_triplane(self, triplane: torch.Tensor, resolution: int): + coords = torch.linspace(-1.0, 1.0, resolution, device = triplane.device) + x, y = torch.meshgrid(coords, coords, indexing="ij") + verts2D = torch.cat([x.view(resolution, resolution,1), y.view(resolution, resolution,1)], dim = -1) + verts2D = verts2D.expand(3, -1, -1, -1) + return F.grid_sample(triplane, verts2D, align_corners=False,mode="bilinear") # [3 40 H W] xy xz yz + + def block_based_marchingcube(self, decoder: torch.nn.Module, triplane: torch.Tensor, resolution: int, threshold, block_resolution = 128) -> torch.Tensor: + resolution += 1 # sample 1 more line of density, 1024 + 1 == 1025, 0 mapping to -1.0f, 512 mapping to 0.0f, 1025 mapping to 1.0f, for better floating point precision. + block_size = 2.0 * block_resolution / (resolution - 1) + voxel_size = block_size / block_resolution + interpolated = self.interpolate_triplane(triplane, resolution) + + pos_list = [] + indices_list = [] + for x in range(0, resolution - 1, block_resolution): + size_x = resolution - x if x + block_resolution >= resolution else block_resolution + 1 # sample 1 more line of density, so marching cubes resolution match block_resolution + for y in range(0, resolution - 1, block_resolution): + size_y = resolution - y if y + block_resolution >= resolution else block_resolution + 1 + for z in range(0, resolution - 1, block_resolution): + size_z = resolution - z if z + block_resolution >= resolution else block_resolution + 1 + xyplane = interpolated[0:1, :, x:x+size_x, y:y+size_y].expand(size_z, -1, -1, -1, -1).permute(3, 4, 0, 1, 2) + xzplane = interpolated[1:2, :, x:x+size_x, z:z+size_z].expand(size_y, -1, -1, -1, -1).permute(3, 0, 4, 1, 2) + yzplane = interpolated[2:3, :, y:y+size_y, z:z+size_z].expand(size_x, -1, -1, -1, -1).permute(0, 3, 4, 1, 2) + sz = size_x * size_y * size_z + out = torch.cat([xyplane, xzplane, yzplane], dim=3).view(sz, 3, -1) + + if self.cfg.feature_reduction == "concat": + out = out.view(sz, -1) + elif self.cfg.feature_reduction == "mean": + out = reduce(out, "N Np Cp -> N Cp", Np=3, reduction="mean") + else: + raise NotImplementedError + net_out = decoder(out) + out = None # discard samples + density = net_out["density"] + net_out = None # discard colors + density = get_activation(self.cfg.density_activation)(density + self.cfg.density_bias).view(size_x, size_y, size_z) + try: # now do the marching cube + v_pos, indices = marching_cubes(density.detach(), threshold) + except AttributeError: + print("torchmcubes was not compiled with CUDA support, use CPU version instead.") + v_pos, indices = self.mc_func(density.detach().cpu(), 0.0) + offset = torch.tensor([x * voxel_size - 1.0, y * voxel_size - 1.0, z * voxel_size - 1.0], device = triplane.device) + v_pos = v_pos[..., [2, 1, 0]] * voxel_size + offset + + indices_list.append(indices) + pos_list.append(v_pos) + + vertex_count = 0 + for i in range(0, len(pos_list)): + indices_list[i] += vertex_count + vertex_count += pos_list[i].size(0) + + return torch.cat(pos_list), torch.cat(indices_list) + def query_triplane( self, decoder: torch.nn.Module, positions: torch.Tensor, triplane: torch.Tensor, + scale_pos = True ) -> Dict[str, torch.Tensor]: input_shape = positions.shape[:-1] positions = positions.view(-1, 3) # positions in (-radius, radius) # normalized to (-1, 1) for grid sample - positions = scale_tensor( - positions, (-self.cfg.radius, self.cfg.radius), (-1, 1) - ) + if scale_pos: + positions = scale_tensor( + positions, (-self.cfg.radius, self.cfg.radius), (-1, 1) + ) def _query_chunk(x): indices2D: torch.Tensor = torch.stack( diff --git a/3d_generation/TripoSR/tsr/system.py b/3d_generation/TripoSR/tsr/system.py index 86c7ee7..8248bf1 100644 --- a/3d_generation/TripoSR/tsr/system.py +++ b/3d_generation/TripoSR/tsr/system.py @@ -13,7 +13,6 @@ from huggingface_hub import hf_hub_download from omegaconf import OmegaConf from PIL import Image -from .models.isosurface import MarchingCubeHelper from .utils import ( BaseModule, ImagePreprocessor, @@ -160,44 +159,25 @@ class TSR(BaseModule): return images - def set_marching_cubes_resolution(self, resolution: int): - if ( - self.isosurface_helper is not None - and self.isosurface_helper.resolution == resolution - ): - return - self.isosurface_helper = MarchingCubeHelper(resolution) - def extract_mesh(self, scene_codes, resolution: int = 256, threshold: float = 25.0): - self.set_marching_cubes_resolution(resolution) meshes = [] for scene_code in scene_codes: with torch.no_grad(): - density = self.renderer.query_triplane( - self.decoder, - scale_tensor( - self.isosurface_helper.grid_vertices.to(scene_codes.device), - self.isosurface_helper.points_range, - (-self.renderer.cfg.radius, self.renderer.cfg.radius), - ), + v_pos, t_pos_idx = self.renderer.block_based_marchingcube(self.decoder.to(scene_codes.device), scene_code, - )["density_act"] - v_pos, t_pos_idx = self.isosurface_helper(-(density - threshold)) - v_pos = scale_tensor( - v_pos, - self.isosurface_helper.points_range, - (-self.renderer.cfg.radius, self.renderer.cfg.radius), - ) - with torch.no_grad(): - color = self.renderer.query_triplane( - self.decoder, + resolution, + threshold + ) + color = self.renderer.query_triplane(self.decoder.to(scene_codes.device), v_pos.to(scene_codes.device), scene_code, False)["color"] + v_pos = scale_tensor( v_pos, - scene_code, - )["color"] - mesh = trimesh.Trimesh( - vertices=v_pos.cpu().numpy(), - faces=t_pos_idx.cpu().numpy(), - vertex_colors=color.cpu().numpy(), - ) - meshes.append(mesh) + (-1.0, 1.0), + (-self.renderer.cfg.radius, self.renderer.cfg.radius) + ) + mesh = trimesh.Trimesh( + vertices=v_pos.cpu().numpy(), + faces=t_pos_idx.cpu().numpy(), + vertex_colors=color.cpu().numpy(), + ) + meshes.append(mesh) return meshes