forked from VAST-AI-Research/TripoSR
-
Notifications
You must be signed in to change notification settings - Fork 0
/
run.py
162 lines (142 loc) · 4.97 KB
/
run.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
import argparse
import logging
import os
import time
import numpy as np
import rembg
import torch
from PIL import Image
from tsr.system import TSR
from tsr.utils import remove_background, resize_foreground, save_video
class Timer:
def __init__(self):
self.items = {}
self.time_scale = 1000.0 # ms
self.time_unit = "ms"
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) -> float:
if name not in self.items:
return
if torch.cuda.is_available():
torch.cuda.synchronize()
start_time = self.items.pop(name)
delta = time.time() - start_time
t = delta * self.time_scale
logging.info(f"{name} finished in {t:.2f}{self.time_unit}.")
timer = Timer()
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
)
parser = argparse.ArgumentParser()
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. If no CUDA-compatible device is found, will fallback to 'cpu'. Default: 'cuda:0'",
)
parser.add_argument(
"--pretrained-model-name-or-path",
default="stabilityai/TripoSR",
type=str,
help="Path to the pretrained model. Could be either a huggingface model id is or a local path. Default: 'stabilityai/TripoSR'",
)
parser.add_argument(
"--chunk-size",
default=8192,
type=int,
help="Evaluation chunk size for surface extraction and rendering. Smaller chunk size reduces VRAM usage but increases computation time. 0 for no chunking. Default: 8192",
)
parser.add_argument(
"--mc-resolution",
default=256,
type=int,
help="Marching cubes grid resolution. Default: 256"
)
parser.add_argument(
"--no-remove-bg",
action="store_true",
help="If specified, the background will NOT be automatically removed from the input image, and the input image should be an RGB image with gray background and properly-sized foreground. Default: false",
)
parser.add_argument(
"--foreground-ratio",
default=0.85,
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(
"--model-save-format",
default="obj",
type=str,
choices=["obj", "glb"],
help="Format to save the extracted mesh. Default: 'obj'",
)
parser.add_argument(
"--render",
action="store_true",
help="If specified, save a NeRF-rendered video. Default: false",
)
args = parser.parse_args()
output_dir = args.output_dir
os.makedirs(output_dir, exist_ok=True)
device = args.device
if not torch.cuda.is_available():
device = "cpu"
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")
timer.start("Processing images")
images = []
if args.no_remove_bg:
rembg_session = None
else:
rembg_session = rembg.new_session()
for i, image_path in enumerate(args.image):
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), f"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([image], device=device)
timer.end("Running model")
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), f"render.mp4"), fps=30
)
timer.end("Rendering")
timer.start("Exporting mesh")
meshes = model.extract_mesh(scene_codes, resolution=args.mc_resolution)
meshes[0].export(os.path.join(output_dir, str(i), f"mesh.{args.model_save_format}"))
timer.end("Exporting mesh")