Initial commit

.gitignore

@@ -0,0 +1,5 @@
+.idea
+.DS_Store
+__pycache__
+gradio_cached_examples
+

Dockerfile

@@ -0,0 +1,41 @@
+FROM toshas/repainting_3d_assets:v2
+ENTRYPOINT []
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+ENV REPAINTING3D_DO_NOT_MANAGE_HFCACHE=1
+
+ENV TRANSFORMERS_CACHE=/data/hfcache
+ENV HF_DATASETS_CACHE=/data/hfcache
+ENV HF_HOME=/data/hfcache
+
+RUN mkdir -p ${HF_HOME}
+
+RUN apt-get -y update --fix-missing && \
+    apt-get install -y \
+        python3 \
+        python3-pip \
+        && \
+    apt-get clean
+
+WORKDIR ${ROOT_PROJECT}
+RUN git clone https://github.com/kongdai123/repainting_3d_assets.git && \
+    rm -rf repainting_3d_assets/.git repainting_3d_assets/assets repainting_3d_assets/doc && \
+    rm -rf code && \
+    mv repainting_3d_assets code
+
+COPY ./requirements.txt ./
+RUN pip3 install --no-cache-dir --upgrade -r requirements.txt
+
+RUN useradd -m -u 1000 user
+
+RUN chown -R user ${HF_HOME}
+
+USER user
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+
+ENV PYTHONUNBUFFERED=1
+WORKDIR ${HOME}/app
+COPY --chown=user . ${HOME}/app
+CMD ["python3", "main.py"]

README.md

@@ -1,11 +1,12 @@
 ---
-title: Repainting 3d Assets
-emoji: 🏆
-colorFrom: gray
-colorTo: red
+title: Repainting 3D Assets
+emoji: 🦄
+colorFrom: teal
+colorTo: indigo
 sdk: docker
-pinned: false
+pinned: true
 license: cc-by-nc-sa-4.0
+models: ["stabilityai/stable-diffusion-2-depth"]
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

main.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+import math
+import os
+import subprocess
+from pathlib import Path
+
+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):
+    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)
+
+
+class Model3D(gr.Model3D):
+    """
+    A simple overload of Gradio Model3D that accepts arbitrary 3D formats supported by trimesh.
+    """
+
+    def postprocess(self, y: str | Path | None) -> dict[str, str] | None:
+        if y is not None:
+            y = convert_formats(y, "glb")
+        out = super().postprocess(y)
+        return out
+
+
+def breathe_new_life_into_3d_model(path_input, prompt):
+    """
+    @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}
+    }
+    """
+    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():
+    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=[
+            Model3D(
+                camera_position=(30.0, 90.0, 3.0),
+                elem_classes="viewport",
+                label="Input Model",
+            ),
+            gr.Textbox(label="Text Prompt"),
+        ],
+        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.ply"),
+                "pastel superhero unicorn",
+            ],
+        ],
+        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()

requirements.txt

@@ -0,0 +1,3 @@
+gradio
+pygltflib
+trimesh