mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-17 00:23:14 +00:00
moonshinkiro
213 lines
7.1 KiB
Python
213 lines
7.1 KiB
Python
import argparse
|
|
import logging
|
|
import os
|
|
import time
|
|
import numpy as np
|
|
import torch
|
|
import xatlas
|
|
from PIL import Image, ImageEnhance
|
|
import trimesh
|
|
import rembg
|
|
from tsr.system import TSR
|
|
from tsr.utils import remove_background, resize_foreground, save_video
|
|
from tsr.bake_texture import bake_texture
|
|
|
|
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}.")
|
|
|
|
def create_textured_glb(obj_path, texture_path, output_glb_path):
|
|
obj = trimesh.load(obj_path, file_type='obj')
|
|
texture = Image.open(texture_path)
|
|
|
|
material = trimesh.visual.texture.SimpleMaterial(image=texture)
|
|
obj.visual.material = material
|
|
|
|
transform = trimesh.transformations.rotation_matrix(np.radians(270), [1, 0, 0])
|
|
obj.apply_transform(transform)
|
|
|
|
scene = trimesh.Scene(obj)
|
|
glb_data = trimesh.exchange.gltf.export_glb(scene)
|
|
|
|
with open(output_glb_path, 'wb') as f:
|
|
f.write(glb_data)
|
|
|
|
|
|
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="glb",
|
|
type=str,
|
|
choices=["obj", "glb"],
|
|
help="Format to save the extracted mesh. Default: 'obj'",
|
|
)
|
|
parser.add_argument(
|
|
"--bake-texture",
|
|
action="store_true",
|
|
help="Bake a texture atlas for the extracted mesh, instead of vertex colors",
|
|
)
|
|
parser.add_argument(
|
|
"--texture-resolution",
|
|
default=2048,
|
|
type=int,
|
|
help="Texture atlas resolution, only useful with --bake-texture. Default: 2048"
|
|
)
|
|
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 = Image.open(image_path).convert("RGB")
|
|
else:
|
|
image = remove_background(Image.open(image_path), rembg_session)
|
|
image = resize_foreground(image, args.foreground_ratio)
|
|
image = image.convert("RGBA")
|
|
image_np = np.array(image).astype(np.float32) / 255.0
|
|
if image_np.shape[2] == 4: # Check if image has alpha channel
|
|
image_np = image_np[:, :, :3] * image_np[:, :, 3:4] + (1 - image_np[:, :, 3:4]) * 0.5
|
|
image = Image.fromarray((image_np * 255.0).astype(np.uint8)).convert("RGB")
|
|
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):
|
|
base_name = os.path.splitext(os.path.basename(args.image[i]))[0]
|
|
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"{base_name}_render_{ri:03d}.png"))
|
|
save_video(
|
|
render_images[0], os.path.join(output_dir, str(i), f"{base_name}_render.mp4"), fps=30
|
|
)
|
|
timer.end("Rendering")
|
|
|
|
timer.start("Extracting mesh")
|
|
meshes = model.extract_mesh(scene_codes, not args.bake_texture, resolution=args.mc_resolution)
|
|
timer.end("Extracting mesh")
|
|
|
|
out_mesh_path = os.path.join(output_dir, str(i), f"{base_name}.{args.model_save_format}")
|
|
if args.bake_texture:
|
|
out_texture_path = os.path.join(output_dir, str(i), f"{base_name}_texture.png")
|
|
|
|
timer.start("Baking texture")
|
|
bake_output = bake_texture(meshes[0], model, scene_codes[0], args.texture_resolution)
|
|
timer.end("Baking texture")
|
|
|
|
timer.start("Exporting mesh and texture")
|
|
xatlas.export(out_mesh_path, meshes[0].vertices[bake_output["vmapping"]], bake_output["indices"], bake_output["uvs"], meshes[0].vertex_normals[bake_output["vmapping"]])
|
|
Image.fromarray((bake_output["colors"] * 255.0).astype(np.uint8)).transpose(Image.FLIP_TOP_BOTTOM).save(out_texture_path)
|
|
timer.end("Exporting mesh and texture")
|
|
else:
|
|
timer.start("Exporting mesh")
|
|
meshes[0].export(out_mesh_path)
|
|
timer.end("Exporting mesh")
|
|
|
|
if args.bake_texture:
|
|
glb_path = os.path.join(output_dir, str(i), f"{base_name}.glb")
|
|
create_textured_glb(out_mesh_path, out_texture_path, glb_path)
|
|
|
|
logging.info("Process completed.")
|