| | import argparse |
| | import os |
| | from contextlib import nullcontext |
| |
|
| | import torch |
| | from PIL import Image |
| | from tqdm import tqdm |
| | from transparent_background import Remover |
| |
|
| | from spar3d.models.mesh import QUAD_REMESH_AVAILABLE, TRIANGLE_REMESH_AVAILABLE |
| | from spar3d.system import SPAR3D |
| | from spar3d.utils import foreground_crop, get_device, remove_background |
| |
|
| |
|
| | def check_positive(value): |
| | ivalue = int(value) |
| | if ivalue <= 0: |
| | raise argparse.ArgumentTypeError("%s is an invalid positive int value" % value) |
| | return ivalue |
| |
|
| |
|
| | if __name__ == "__main__": |
| | parser = argparse.ArgumentParser() |
| | parser.add_argument( |
| | "image", type=str, nargs="+", help="Path to input image(s) or folder." |
| | ) |
| | parser.add_argument( |
| | "--device", |
| | default=get_device(), |
| | type=str, |
| | help=f"Device to use. If no CUDA/MPS-compatible device is found, the baking will fail. Default: '{get_device()}'", |
| | ) |
| | parser.add_argument( |
| | "--pretrained-model", |
| | default="stabilityai/stable-point-aware-3d", |
| | type=str, |
| | help="Path to the pretrained model. Could be either a huggingface model id is or a local path. Default: 'stabilityai/stable-point-aware-3d'", |
| | ) |
| | parser.add_argument( |
| | "--foreground-ratio", |
| | default=1.3, |
| | type=float, |
| | help="Ratio of the foreground size to the image size. Only used when --no-remove-bg is not specified. Default: 0.85", |
| | ) |
| | parser.add_argument( |
| | "--output-dir", |
| | default="output/", |
| | type=str, |
| | help="Output directory to save the results. Default: 'output/'", |
| | ) |
| | parser.add_argument( |
| | "--texture-resolution", |
| | default=1024, |
| | type=int, |
| | help="Texture atlas resolution. Default: 1024", |
| | ) |
| | parser.add_argument( |
| | "--low-vram-mode", |
| | action="store_true", |
| | help=( |
| | "Use low VRAM mode. SPAR3D consumes 10.5GB of VRAM by default. " |
| | "This mode will reduce the VRAM consumption to roughly 7GB but in exchange " |
| | "the model will be slower. Default: False" |
| | ), |
| | ) |
| |
|
| | remesh_choices = ["none"] |
| | if TRIANGLE_REMESH_AVAILABLE: |
| | remesh_choices.append("triangle") |
| | if QUAD_REMESH_AVAILABLE: |
| | remesh_choices.append("quad") |
| | parser.add_argument( |
| | "--remesh_option", |
| | choices=remesh_choices, |
| | default="none", |
| | help="Remeshing option", |
| | ) |
| | if TRIANGLE_REMESH_AVAILABLE or QUAD_REMESH_AVAILABLE: |
| | parser.add_argument( |
| | "--reduction_count_type", |
| | choices=["keep", "vertex", "faces"], |
| | default="keep", |
| | help="Vertex count type", |
| | ) |
| | parser.add_argument( |
| | "--target_count", |
| | type=check_positive, |
| | help="Selected target count.", |
| | default=2000, |
| | ) |
| | parser.add_argument( |
| | "--batch_size", default=1, type=int, help="Batch size for inference" |
| | ) |
| | args = parser.parse_args() |
| |
|
| | |
| | devices = ["cuda", "mps", "cpu"] |
| | if not any(args.device in device for device in devices): |
| | raise ValueError("Invalid device. Use cuda, mps or cpu") |
| |
|
| | output_dir = args.output_dir |
| | os.makedirs(output_dir, exist_ok=True) |
| |
|
| | device = args.device |
| | if not (torch.cuda.is_available() or torch.backends.mps.is_available()): |
| | device = "cpu" |
| |
|
| | print("Device used: ", device) |
| |
|
| | model = SPAR3D.from_pretrained( |
| | args.pretrained_model, |
| | config_name="config.yaml", |
| | weight_name="model.safetensors", |
| | low_vram_mode=args.low_vram_mode, |
| | ) |
| | model.to(device) |
| | model.eval() |
| |
|
| | bg_remover = Remover(device=device) |
| | images = [] |
| | idx = 0 |
| | for image_path in args.image: |
| |
|
| | def handle_image(image_path, idx): |
| | image = remove_background( |
| | Image.open(image_path).convert("RGBA"), bg_remover |
| | ) |
| | image = foreground_crop(image, args.foreground_ratio) |
| | os.makedirs(os.path.join(output_dir, str(idx)), exist_ok=True) |
| | image.save(os.path.join(output_dir, str(idx), "input.png")) |
| | images.append(image) |
| |
|
| | if os.path.isdir(image_path): |
| | image_paths = [ |
| | os.path.join(image_path, f) |
| | for f in os.listdir(image_path) |
| | if f.endswith((".png", ".jpg", ".jpeg")) |
| | ] |
| | for image_path in image_paths: |
| | handle_image(image_path, idx) |
| | idx += 1 |
| | else: |
| | handle_image(image_path, idx) |
| | idx += 1 |
| |
|
| | vertex_count = ( |
| | -1 |
| | if args.reduction_count_type == "keep" |
| | else ( |
| | args.target_count |
| | if args.reduction_count_type == "vertex" |
| | else args.target_count // 2 |
| | ) |
| | ) |
| |
|
| | for i in tqdm(range(0, len(images), args.batch_size)): |
| | image = images[i : i + args.batch_size] |
| | if torch.cuda.is_available(): |
| | torch.cuda.reset_peak_memory_stats() |
| | with torch.no_grad(): |
| | with ( |
| | torch.autocast(device_type=device, dtype=torch.bfloat16) |
| | if "cuda" in device |
| | else nullcontext() |
| | ): |
| | mesh, glob_dict = model.run_image( |
| | image, |
| | bake_resolution=args.texture_resolution, |
| | remesh=args.remesh_option, |
| | vertex_count=vertex_count, |
| | return_points=True, |
| | ) |
| | if torch.cuda.is_available(): |
| | print("Peak Memory:", torch.cuda.max_memory_allocated() / 1024 / 1024, "MB") |
| | elif torch.backends.mps.is_available(): |
| | print( |
| | "Peak Memory:", torch.mps.driver_allocated_memory() / 1024 / 1024, "MB" |
| | ) |
| |
|
| | if len(image) == 1: |
| | out_mesh_path = os.path.join(output_dir, str(i), "mesh.glb") |
| | mesh.export(out_mesh_path, include_normals=True) |
| | out_points_path = os.path.join(output_dir, str(i), "points.ply") |
| | glob_dict["point_clouds"][0].export(out_points_path) |
| | else: |
| | for j in range(len(mesh)): |
| | out_mesh_path = os.path.join(output_dir, str(i + j), "mesh.glb") |
| | mesh[j].export(out_mesh_path, include_normals=True) |
| | out_points_path = os.path.join(output_dir, str(i + j), "points.ply") |
| | glob_dict["point_clouds"][j].export(out_points_path) |
| |
|
