Support texture generation

3d_generation/TripoSR/README.md

@@ -1,4 +1,4 @@
-# TripoSR <a href="https://huggingface.co/stabilityai/TripoSR"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Model_Card-Huggingface-orange"></a> <a href="https://huggingface.co/spaces/stabilityai/TripoSR"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Gradio%20Demo-Huggingface-orange"></a> <a href="https://arxiv.org/abs/2403.02151"><img src="https://img.shields.io/badge/Arxiv-2403.02151-B31B1B.svg"></a>
+# TripoSR <a href="https://huggingface.co/stabilityai/TripoSR"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Model_Card-Huggingface-orange"></a> <a href="https://huggingface.co/spaces/stabilityai/TripoSR"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Gradio%20Demo-Huggingface-orange"></a> <a href="https://huggingface.co/papers/2403.02151"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Paper-Huggingface-orange"></a> <a href="https://arxiv.org/abs/2403.02151"><img src="https://img.shields.io/badge/Arxiv-2403.02151-B31B1B.svg"></a> <a href="https://discord.gg/mvS9mCfMnQ"><img src="https://img.shields.io/badge/Discord-%235865F2.svg?logo=discord&logoColor=white"></a>
 
 <div align="center">
   <img src="figures/teaser800.gif" alt="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
 ```

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],

3d_generation/TripoSR/requirements.txt

@@ -7,3 +7,4 @@ trimesh==4.0.5
 rembg
 huggingface-hub
 imageio[ffmpeg]
+gradio

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)
+    main()
+           

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(

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