440 lines
14 KiB
Python
440 lines
14 KiB
Python
|
import mimetypes
|
|||
|
|
import os
|
|||
|
|
import secrets
|
|||
|
|
import subprocess
|
|||
|
|
import tempfile
|
|||
|
|
import uuid
|
|||
|
|
from datetime import datetime
|
|||
|
|
from io import BytesIO
|
|||
|
|
|
|||
|
|
import imageio
|
|||
|
|
import numpy as np
|
|||
|
|
import trimesh
|
|||
|
|
|
|||
|
|
import litserve as ls
|
|||
|
|
from minio import Minio
|
|||
|
|
|
|||
|
|
from glb2svg import glb_to_obj, obj_to_step, step_to_svg
|
|||
|
|
from trellis.pipelines import TrellisImageTo3DPipeline
|
|||
|
|
from trellis.utils import render_utils, postprocessing_utils
|
|||
|
|
from utils.new_oss_client import MINIO_URL, MINIO_ACCESS, MINIO_SECRET, MINIO_SECURE, minio_get_image, MINIO_BUCKET, upload_bytes, upload_local_file, download_from_minio
|
|||
|
|
|
|||
|
|
minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def generate_unique_name(original_name: str) -> str:
|
|||
|
|
stem, ext = os.path.splitext(original_name)
|
|||
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|||
|
|
random_part = secrets.token_hex(4) # 8位随机十六进制
|
|||
|
|
return f"{stem}_{timestamp}_{random_part}{ext}"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def load_mesh(file_path):
|
|||
|
|
"""
|
|||
|
|
加载.obj或.glb文件，返回顶点数据
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
file_path: 模型文件路径（支持.obj和.glb格式）
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
numpy.ndarray: 顶点数组，shape为(N, 3)
|
|||
|
|
"""
|
|||
|
|
file_ext = os.path.splitext(file_path)[1].lower()
|
|||
|
|
|
|||
|
|
if file_ext == '.obj':
|
|||
|
|
# 使用trimesh加载obj文件
|
|||
|
|
mesh = trimesh.load(file_path, file_type='obj')
|
|||
|
|
elif file_ext == '.glb' or file_ext == '.gltf':
|
|||
|
|
# 使用trimesh加载glb/gltf文件
|
|||
|
|
mesh = trimesh.load(file_path, file_type='glb')
|
|||
|
|
else:
|
|||
|
|
raise ValueError(f"不支持的文件格式: {file_ext}，仅支持.obj和.glb/.gltf")
|
|||
|
|
|
|||
|
|
# 获取顶点数据
|
|||
|
|
if isinstance(mesh, trimesh.Scene):
|
|||
|
|
# 如果是场景，合并所有几何体的顶点
|
|||
|
|
vertices = []
|
|||
|
|
for geom in mesh.geometry.values():
|
|||
|
|
vertices.append(geom.vertices)
|
|||
|
|
vertices = np.vstack(vertices)
|
|||
|
|
else:
|
|||
|
|
# 如果是单个几何体
|
|||
|
|
vertices = mesh.vertices
|
|||
|
|
|
|||
|
|
if len(vertices) == 0:
|
|||
|
|
raise ValueError("文件中未找到顶点数据")
|
|||
|
|
|
|||
|
|
return vertices
|
|||
|
|
|
|||
|
|
|
|||
|
|
def analyze_mesh(file_path):
|
|||
|
|
"""
|
|||
|
|
分析3D模型文件，计算质心、边界框、尺寸等信息
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
file_path: 模型文件路径（支持.obj和.glb格式）
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
dict: 包含模型分析信息的字典
|
|||
|
|
"""
|
|||
|
|
# 加载模型并获取顶点
|
|||
|
|
vertices = load_mesh(file_path)
|
|||
|
|
|
|||
|
|
# 边界框（每个轴的最小/最大值）
|
|||
|
|
min_coords = vertices.min(axis=0)
|
|||
|
|
max_coords = vertices.max(axis=0)
|
|||
|
|
|
|||
|
|
# 质心
|
|||
|
|
centroid = vertices.mean(axis=0)
|
|||
|
|
|
|||
|
|
# 尺寸 = 边界框维度
|
|||
|
|
size = max_coords - min_coords
|
|||
|
|
|
|||
|
|
# 计算尺寸比例（每个轴占总尺寸的比例）
|
|||
|
|
total_size = np.sum(size)
|
|||
|
|
size_ratio = size / total_size if total_size != 0 else [0, 0, 0]
|
|||
|
|
|
|||
|
|
info = {
|
|||
|
|
# "file_path": file_path,
|
|||
|
|
"file_format": os.path.splitext(file_path)[1].lower(),
|
|||
|
|
"vertex_count": len(vertices),
|
|||
|
|
"centroid": centroid.tolist(),
|
|||
|
|
"bounding_box_min": min_coords.tolist(),
|
|||
|
|
"bounding_box_max": max_coords.tolist(),
|
|||
|
|
"size": size.tolist(),
|
|||
|
|
"size_ratio": size_ratio.tolist(),
|
|||
|
|
"size_ratio_percentage": (size_ratio * 100).tolist()
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return info
|
|||
|
|
|
|||
|
|
|
|||
|
|
def render_glb_preview(glb_path, output_path):
|
|||
|
|
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
|||
|
|
|
|||
|
|
cmd = [
|
|||
|
|
"blender",
|
|||
|
|
"--background",
|
|||
|
|
"--python",
|
|||
|
|
"render_model.py",
|
|||
|
|
"--",
|
|||
|
|
glb_path,
|
|||
|
|
output_path
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
result = subprocess.run(cmd, capture_output=True, text=True)
|
|||
|
|
|
|||
|
|
if result.returncode != 0:
|
|||
|
|
raise RuntimeError(
|
|||
|
|
f"Blender render failed\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return output_path
|
|||
|
|
|
|||
|
|
|
|||
|
|
class TrellisAPI(ls.LitAPI):
|
|||
|
|
|
|||
|
|
def setup(self, device):
|
|||
|
|
os.environ.setdefault("SPCONV_ALGO", "native")
|
|||
|
|
|
|||
|
|
self.pipeline = TrellisImageTo3DPipeline.from_pretrained(
|
|||
|
|
"microsoft/TRELLIS-image-large"
|
|||
|
|
)
|
|||
|
|
self.pipeline.to(device)
|
|||
|
|
|
|||
|
|
def decode_request(self, request):
|
|||
|
|
image_paths = request["image_paths"]
|
|||
|
|
images = []
|
|||
|
|
for path in image_paths:
|
|||
|
|
bucket = path.split('/')[0]
|
|||
|
|
object_name = path[path.find('/') + 1:]
|
|||
|
|
|
|||
|
|
image = minio_get_image(minio_client, bucket, object_name)
|
|||
|
|
images.append(image)
|
|||
|
|
|
|||
|
|
params = {
|
|||
|
|
"file_name": uuid.uuid4().hex,
|
|||
|
|
"model": request.get("model", "single"),
|
|||
|
|
"seed": request.get("seed", 1),
|
|||
|
|
"steps_sparse": request.get("steps_sparse", 12),
|
|||
|
|
"cfg_sparse": request.get("cfg_sparse", 7.5),
|
|||
|
|
"steps_slat": request.get("steps_slat", 12),
|
|||
|
|
"cfg_slat": request.get("cfg_slat", 3.0),
|
|||
|
|
"simplify": request.get("simplify", 0.95),
|
|||
|
|
"texture_size": request.get("texture_size", 1024),
|
|||
|
|
"fps": request.get("fps", 30),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return images, params
|
|||
|
|
|
|||
|
|
def predict(self, inputs):
|
|||
|
|
images, params = inputs
|
|||
|
|
if params["model"] == "single":
|
|||
|
|
outputs = self.pipeline.run(
|
|||
|
|
images[0],
|
|||
|
|
seed=params["seed"],
|
|||
|
|
sparse_structure_sampler_params={
|
|||
|
|
"steps": params["steps_sparse"],
|
|||
|
|
"cfg_strength": params["cfg_sparse"],
|
|||
|
|
},
|
|||
|
|
slat_sampler_params={
|
|||
|
|
"steps": params["steps_slat"],
|
|||
|
|
"cfg_strength": params["cfg_slat"],
|
|||
|
|
},
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
outputs = self.pipeline.run_multi_image(
|
|||
|
|
images,
|
|||
|
|
seed=params['seed'],
|
|||
|
|
sparse_structure_sampler_params={
|
|||
|
|
"steps": params['steps_sparse'],
|
|||
|
|
"cfg_strength": params['cfg_sparse'],
|
|||
|
|
},
|
|||
|
|
slat_sampler_params={
|
|||
|
|
"steps": params['steps_slat'],
|
|||
|
|
"cfg_strength": params['cfg_slat'],
|
|||
|
|
},
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# video_path = self.upload_video(outputs, params)
|
|||
|
|
|
|||
|
|
minio_glb_path, local_glb_path = self.upload_glb(outputs, params)
|
|||
|
|
|
|||
|
|
glb_info = analyze_mesh(local_glb_path)
|
|||
|
|
|
|||
|
|
local_static_model_image_path = os.path.join("glb_output", generate_unique_name("static_model_image.png"))
|
|||
|
|
static_model_image = self.get_static_model_image(model_path=local_glb_path, output_path=local_static_model_image_path)
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"glb_path": minio_glb_path,
|
|||
|
|
"glb_static_img_path": static_model_image,
|
|||
|
|
"glb_info": glb_info,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
def encode_response(self, output):
|
|||
|
|
return output
|
|||
|
|
|
|||
|
|
def upload_video(self, outputs, params):
|
|||
|
|
gaussian_name = f"3d_result/video/{params['file_name']}-gaussian.mp4"
|
|||
|
|
radiance_field_name = f"3d_result/video/{params['file_name']}-radiance_field.mp4"
|
|||
|
|
mesh_name = f"3d_result/video/{params['file_name']}-mesh.mp4"
|
|||
|
|
|
|||
|
|
# gaussian video
|
|||
|
|
video = render_utils.render_video(outputs["gaussian"][0])["color"]
|
|||
|
|
buffer = BytesIO()
|
|||
|
|
imageio.mimsave(buffer, video, format="mp4", fps=params['fps'])
|
|||
|
|
gaussian_video_path = upload_bytes(
|
|||
|
|
buffer.getvalue(),
|
|||
|
|
gaussian_name,
|
|||
|
|
"video/mp4",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# radiance field video
|
|||
|
|
video = render_utils.render_video(outputs["radiance_field"][0])["color"]
|
|||
|
|
buffer = BytesIO()
|
|||
|
|
imageio.mimsave(buffer, video, format="mp4", fps=params['fps'])
|
|||
|
|
radiance_field_video_path = upload_bytes(
|
|||
|
|
buffer.getvalue(),
|
|||
|
|
radiance_field_name,
|
|||
|
|
"video/mp4",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# mesh video
|
|||
|
|
video = render_utils.render_video(outputs["mesh"][0])["normal"]
|
|||
|
|
buffer = BytesIO()
|
|||
|
|
imageio.mimsave(buffer, video, format="mp4", fps=params['fps'])
|
|||
|
|
mesh_path = upload_bytes(
|
|||
|
|
buffer.getvalue(),
|
|||
|
|
mesh_name,
|
|||
|
|
"video/mp4",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"gaussian": gaussian_video_path,
|
|||
|
|
"radiance_field": radiance_field_video_path,
|
|||
|
|
"mesh": mesh_path
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
def upload_glb(self, outputs, params):
|
|||
|
|
file_name = f"3d_result/glb/{params['file_name']}.glb"
|
|||
|
|
local_glb_path = os.path.join("glb_output", generate_unique_name("sample.glb"))
|
|||
|
|
out_dir = os.path.dirname(local_glb_path)
|
|||
|
|
if out_dir:
|
|||
|
|
os.makedirs(out_dir, exist_ok=True)
|
|||
|
|
|
|||
|
|
glb = postprocessing_utils.to_glb(
|
|||
|
|
outputs["gaussian"][0],
|
|||
|
|
outputs["mesh"][0],
|
|||
|
|
simplify=params['simplify'],
|
|||
|
|
texture_size=params['texture_size'],
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
glb.export(
|
|||
|
|
file_obj=local_glb_path,
|
|||
|
|
file_type="glb"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
glb_path = upload_local_file(
|
|||
|
|
local_glb_path,
|
|||
|
|
file_name,
|
|||
|
|
"application/octet-stream",
|
|||
|
|
)
|
|||
|
|
return glb_path, local_glb_path
|
|||
|
|
|
|||
|
|
def upload_ply(self, outputs, params):
|
|||
|
|
file_name = f"3d_result/ply/{params['file_name']}.ply"
|
|||
|
|
|
|||
|
|
with tempfile.NamedTemporaryFile(suffix=".ply") as tmp:
|
|||
|
|
outputs["gaussian"][0].save_ply(tmp.name)
|
|||
|
|
tmp.seek(0)
|
|||
|
|
|
|||
|
|
ply_path = upload_bytes(
|
|||
|
|
tmp.read(),
|
|||
|
|
file_name,
|
|||
|
|
"application/octet-stream",
|
|||
|
|
)
|
|||
|
|
return {"ply": ply_path}
|
|||
|
|
|
|||
|
|
def get_static_model_image(self, model_path, output_path):
|
|||
|
|
local_static_model_image_path = os.path.join(
|
|||
|
|
"glb_output",
|
|||
|
|
generate_unique_name("static_model_image.png")
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
print(f"model_path : {model_path}")
|
|||
|
|
print(f"local_static_model_image_path :{local_static_model_image_path}")
|
|||
|
|
output_path = render_glb_preview(model_path, local_static_model_image_path)
|
|||
|
|
|
|||
|
|
static_model_image = self.upload_local_file(
|
|||
|
|
output_path,
|
|||
|
|
"png"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
print(f"Saved to {static_model_image}")
|
|||
|
|
return static_model_image
|
|||
|
|
|
|||
|
|
def upload_local_file(self, local_path, type):
|
|||
|
|
"""
|
|||
|
|
通用上传函数：支持 SVG, PNG, OBJ 等
|
|||
|
|
"""
|
|||
|
|
object_name = f"3d_result/{type}/{uuid.uuid4().hex}.{type}"
|
|||
|
|
if not os.path.exists(local_path):
|
|||
|
|
print(f"错误: 文件 {local_path} 不存在")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
# 自动根据后缀名识别 Content-Type
|
|||
|
|
# 例如: .svg -> image/svg+xml, .png -> image/png
|
|||
|
|
content_type, _ = mimetypes.guess_type(local_path)
|
|||
|
|
if content_type is None:
|
|||
|
|
content_type = "application/octet-stream"
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
minio_client.fput_object(
|
|||
|
|
bucket_name=MINIO_BUCKET,
|
|||
|
|
object_name=object_name,
|
|||
|
|
file_path=local_path,
|
|||
|
|
content_type=content_type
|
|||
|
|
)
|
|||
|
|
print(f"成功上传 [{content_type}]: {object_name}")
|
|||
|
|
return f"{MINIO_BUCKET}/{object_name}"
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"上传失败: {e}")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ModelToThreeViews(ls.LitAPI):
|
|||
|
|
|
|||
|
|
def setup(self, device):
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
def upload_local_file(self, local_path, type):
|
|||
|
|
"""
|
|||
|
|
通用上传函数：支持 SVG, PNG, OBJ 等
|
|||
|
|
"""
|
|||
|
|
object_name = f"3d_result/{type}/{uuid.uuid4().hex}.{type}"
|
|||
|
|
if not os.path.exists(local_path):
|
|||
|
|
print(f"错误: 文件 {local_path} 不存在")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
# 自动根据后缀名识别 Content-Type
|
|||
|
|
# 例如: .svg -> image/svg+xml, .png -> image/png
|
|||
|
|
content_type, _ = mimetypes.guess_type(local_path)
|
|||
|
|
if content_type is None:
|
|||
|
|
content_type = "application/octet-stream"
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
minio_client.fput_object(
|
|||
|
|
bucket_name=MINIO_BUCKET,
|
|||
|
|
object_name=object_name,
|
|||
|
|
file_path=local_path,
|
|||
|
|
content_type=content_type
|
|||
|
|
)
|
|||
|
|
print(f"成功上传 [{content_type}]: {object_name}")
|
|||
|
|
return f"{MINIO_BUCKET}/{object_name}"
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"上传失败: {e}")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def predict(self, request):
|
|||
|
|
minio_glb_path = request['minio_glb_path']
|
|||
|
|
|
|||
|
|
work_dir = f"glb_to_obj"
|
|||
|
|
os.makedirs(work_dir, exist_ok=True)
|
|||
|
|
|
|||
|
|
glb_path = os.path.join(work_dir, f"model{uuid.uuid4().hex}.glb")
|
|||
|
|
step_dir = os.path.join(work_dir, "step")
|
|||
|
|
svg_dir = os.path.join(work_dir, "svg")
|
|||
|
|
os.makedirs(step_dir, exist_ok=True)
|
|||
|
|
os.makedirs(svg_dir, exist_ok=True)
|
|||
|
|
print(f"""
|
|||
|
|
入参阶段:
|
|||
|
|
input glb-obj minio-path:{minio_glb_path},\n
|
|||
|
|
work_dir : {work_dir},glb_path : {glb_path},step_dir : {step_dir},svg_dir : {svg_dir}\n
|
|||
|
|
""")
|
|||
|
|
print("=" * 10)
|
|||
|
|
|
|||
|
|
print(f" 第一阶段 下载glb文件: ")
|
|||
|
|
# 1 下载
|
|||
|
|
glb_result = download_from_minio(object_path=minio_glb_path, local_path=glb_path)
|
|||
|
|
print(f" 下载结果 : {glb_result} \n")
|
|||
|
|
print("=" * 10)
|
|||
|
|
|
|||
|
|
print(f" 第二阶段 glb -> obj: ")
|
|||
|
|
# 2 glb -> obj
|
|||
|
|
obj_result = glb_to_obj(glb_result)
|
|||
|
|
print(f" glb -> obj 结果 : {obj_result} \n")
|
|||
|
|
print("=" * 10)
|
|||
|
|
|
|||
|
|
print(f" 第三阶段 obj -> step: ")
|
|||
|
|
# 3 obj -> step
|
|||
|
|
step_result = obj_to_step(
|
|||
|
|
input_obj=obj_result,
|
|||
|
|
output_dir=step_dir,
|
|||
|
|
script_path="1_obj_to_step.py"
|
|||
|
|
)
|
|||
|
|
print(f" obj -> step 结果 : {step_result} \n")
|
|||
|
|
print("=" * 10)
|
|||
|
|
|
|||
|
|
print(f" 第四阶段 step -> svg: ")
|
|||
|
|
# 4 step -> svg
|
|||
|
|
combined_svg, combined_png = step_to_svg(
|
|||
|
|
step_path=step_result,
|
|||
|
|
out_dir=svg_dir
|
|||
|
|
)
|
|||
|
|
print(f" step -> svg 结果 : {combined_svg} \n")
|
|||
|
|
print("=" * 10)
|
|||
|
|
|
|||
|
|
# 5 上传
|
|||
|
|
minio_svg_path = self.upload_local_file(combined_png, "svg")
|
|||
|
|
|
|||
|
|
return {"minio_svg_path": minio_svg_path}
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
trellis_api = TrellisAPI(api_path="/canvas/img_to_3D")
|
|||
|
|
model_to_three_api = ModelToThreeViews(api_path="/canvas/3d_to_3views")
|
|||
|
|
server = ls.LitServer([
|
|||
|
|
trellis_api,
|
|||
|
|
model_to_three_api])
|
|||
|
|
server.run(port=8120)
|