Spaces:

toshas
/

repainting_3d_assets

Running on A10G

App Files Files Community

repainting_3d_assets / main.py

toshas

drop ply format to enable frontend upload of glb

61a0789 8 months ago

raw history blame contribute delete

No virus

6.89 kB

	from __future__ import annotations

	import math
	import os
	import subprocess

	import gradio as gr
	import pygltflib
	import trimesh


	def convert_formats(path_input, target_ext):
	"""
	Converts an input 3D model under input path to the target extensions format and returns a path to that file.
	:param path_input: path to user input
	:param target_ext: target extension
	:return: path to the input 3D model stored in target format.
	"""
	path_input_base, ext = os.path.splitext(path_input)
	if ext == "." + target_ext:
	return path_input
	path_output = path_input_base + "." + target_ext
	if not os.path.exists(path_output):
	trimesh.load_mesh(path_input).export(path_output)
	return path_output


	def add_lights(path_input, path_output):
	"""
	Adds directional lights in the horizontal plane to the glb file.
	:param path_input: path to input glb
	:param path_output: path to output glb
	:return: None
	"""
	glb = pygltflib.GLTF2().load(path_input)

	N = 3 # default max num lights in Babylon.js is 4
	angle_step = 2 * math.pi / N

	lights_extension = {
	"lights": [
	{"type": "directional", "color": [1.0, 1.0, 1.0], "intensity": 2.0}
	for _ in range(N)
	]
	}

	if "KHR_lights_punctual" not in glb.extensionsUsed:
	glb.extensionsUsed.append("KHR_lights_punctual")
	glb.extensions["KHR_lights_punctual"] = lights_extension

	light_nodes = []
	for i in range(N):
	angle = i * angle_step
	rotation = [0.0, math.sin(angle / 2), 0.0, math.cos(angle / 2)]
	node = {
	"rotation": rotation,
	"extensions": {"KHR_lights_punctual": {"light": i}},
	}
	light_nodes.append(node)

	light_node_indices = list(range(len(glb.nodes), len(glb.nodes) + N))
	glb.nodes.extend(light_nodes)

	root_node_index = glb.scenes[glb.scene].nodes[0]
	root_node = glb.nodes[root_node_index]
	if hasattr(root_node, "children"):
	root_node.children.extend(light_node_indices)
	else:
	root_node.children = light_node_indices

	glb.save(path_output)


	def breathe_new_life_into_3d_model(path_input, prompt, path_output=None):
	"""
	Paints the input 3D model using the following method:
	@inproceedings{wang2023breathing,
	title={Breathing New Life into 3D Assets with Generative Repainting},
	author={Wang, Tianfu and Kanakis, Menelaos and Schindler, Konrad and Van Gool, Luc and Obukhov, Anton},
	booktitle={Proceedings of the British Machine Vision Conference (BMVC)},
	year={2023},
	publisher={BMVA Press}
	}
	:param path_input: path to input 3D model
	:param prompt: text description of the expected output
	:param path_output: path to precomputed (cached) output if available
	:return: path to the painted 3D model in glb format with lights for Gradio web component
	"""
	if path_output is not None:
	path_output = path_output.name
	path_output_glb_vis = path_output[:-4] + "_vis.glb"
	add_lights(path_output, path_output_glb_vis)
	return path_output_glb_vis

	path_output_dir = path_input + ".output"
	os.makedirs(path_output_dir, exist_ok=True)

	path_input_ply = convert_formats(path_input, "ply")

	cmd = [
	"bash",
	"/repainting_3d_assets/code/scripts/conda_run.sh",
	"/repainting_3d_assets",
	path_input_ply,
	path_output_dir,
	prompt,
	]
	result = subprocess.run(cmd, env=os.environ, text=True)

	if result.returncode != 0:
	print(f"Output: {result.stdout}")
	print(f"Stderr: {result.stderr}")
	raise RuntimeError("Processing failed")

	path_output_glb = os.path.join(path_output_dir, "model_draco.glb")
	path_output_glb_vis = path_output_glb[:-4] + "_vis.glb"
	add_lights(path_output_glb, path_output_glb_vis)

	return path_output_glb_vis


	def run():
	"""
	Gradio entry point
	"""
	desc = """
	<p align="center">
	<a title="Website" href="https://www.obukhov.ai/repainting_3d_assets" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
	<img src="https://www.obukhov.ai/img/badges/badge-website.svg">
	</a>
	<a title="arXiv" href="https://arxiv.org/abs/2309.08523" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
	<img src="https://www.obukhov.ai/img/badges/badge-pdf.svg">
	</a>
	<a title="Github" href="https://github.com/kongdai123/repainting_3d_assets" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
	<img src="https://img.shields.io/github/stars/kongdai123/repainting_3d_assets?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="badge-github-stars">
	</a>
	<a title="Social" href="https://twitter.com/antonobukhov1" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
	<img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
	</a>
	</p>
	<p align="justify">
	Repaint your 3D models with a text prompt, guided by a method from our BMVC'2023 Oral paper 'Breathing New Life
	into 3D Assets with Generative Repainting'. Simply drop a model into the left pane, specify your repainting
	preferences, and wait for the outcome (~20 min). Explore precomputed examples at the bottom, or follow the
	Project Website badge for additional precomputed models and comparison with other repainting techniques.
	</p>
	"""
	demo = gr.Interface(
	title="Repainting 3D Assets",
	description=desc,
	thumbnail="thumbnail.jpg",
	fn=breathe_new_life_into_3d_model,
	inputs=[
	gr.Model3D(
	camera_position=(30.0, 90.0, 3.0),
	elem_classes="viewport",
	label="Input Model",
	),
	gr.Textbox(
	label="Text Prompt",
	),
	gr.File(visible=False, label="Cached Output"),
	],
	outputs=[
	gr.Model3D(
	camera_position=(30.0, 90.0, 3.0),
	elem_classes="viewport",
	label="Repainted Model",
	),
	],
	examples=[
	[
	os.path.join(os.path.dirname(__file__), "files/horse.glb"),
	"pastel superhero unicorn",
	os.path.join(
	os.path.dirname(__file__), "files/horse.glb.output/output.glb"
	),
	],
	],
	cache_examples=True,
	css="""
	.viewport {
	aspect-ratio: 16/9;
	}
	""",
	allow_flagging="never",
	)

	demo.queue().launch(server_name="0.0.0.0", server_port=7860)


	if __name__ == "__main__":
	run()