|
|
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)
|
|
|
|
