3D_Stage/webui.py

import os
import json
import tqdm
import cv2
import numpy as np
import torch, lrm
import torch.nn.functional as F
from lrm.utils.config import load_config
from datetime import datetime
import gradio as gr
from pygltflib import GLTF2
from PIL import Image
from huggingface_hub import hf_hub_download

from refine import refine

device = "cuda"

import trimesh
import pymeshlab
import numpy as np

from huggingface_hub import hf_hub_download, list_repo_files

repo_id = "zjpshadow/CharacterGen"
all_files = list_repo_files(repo_id, revision="main")

for file in all_files:
    if os.path.exists("../" + file):
        continue
    if file.startswith("3D_Stage"):
        hf_hub_download(repo_id, file, local_dir="../")

def traverse(path, back_proj):
    mesh = trimesh.load(f"{path}/model-00.obj")
    mesh.apply_transform(trimesh.transformations.rotation_matrix(np.radians(90.0), [-1, 0, 0]))
    mesh.apply_transform(trimesh.transformations.rotation_matrix(np.radians(180.0), [0, 1, 0]))

    cmesh = pymeshlab.Mesh(mesh.vertices, mesh.faces)
    ms = pymeshlab.MeshSet()
    ms.add_mesh(cmesh)
    ms.apply_coord_laplacian_smoothing(stepsmoothnum=4)
    mesh.vertices = ms.current_mesh().vertex_matrix()

    mesh.export(f'{path}/output.glb', file_type='glb')

    image = Image.open(f"{path}/{'refined_texture_kd.jpg' if back_proj else 'texture_kd.jpg'}")
    texture = np.array(image)
    vertex_colors = np.zeros((mesh.vertices.shape[0], 4), dtype=np.uint8)

    for vertex_index in range(len(mesh.visual.uv)):
        uv = mesh.visual.uv[vertex_index]
        x = int(uv[0] * (texture.shape[1] - 1))
        y = int((1 - uv[1]) * (texture.shape[0] - 1))

        color = texture[y, x, :3]
        vertex_colors[vertex_index] = [color[0], color[1], color[2], 255]
    return trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces, vertex_colors=vertex_colors)

class Inference_API:

    def __init__(self):
        # Load config
        self.cfg = load_config("configs/infer.yaml", makedirs=False)
        # Load system
        print("Loading system")
        self.system = lrm.find(self.cfg.system_cls)(self.cfg.system).to(device)
        self.system.eval()

    def process_images(self, img_input0, img_input1, img_input2, img_input3, back_proj):
        meta = json.load(open("material/meta.json"))
        c2w_cond = [np.array(loc["transform_matrix"]) for loc in meta["locations"]]
        c2w_cond = torch.from_numpy(np.stack(c2w_cond, axis=0)).float()[None].to(device)
        
        # Prepare input data
        rgb_cond = []
        files = [img_input0, img_input1, img_input2, img_input3]
        new_image = []
        for file in files:
            image = np.array(file)
            image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
            new_image.append(Image.fromarray(image.astype(np.uint8)).convert("RGB"))
            rgb = cv2.resize(image, (self.cfg.data.cond_width, 
                                     self.cfg.data.cond_height)).astype(np.float32) / 255.0
            rgb_cond.append(rgb)
        assert len(rgb_cond) == 4, "Please provide 4 images"

        rgb_cond = torch.from_numpy(np.stack(rgb_cond, axis=0)).float()[None].to(device)

        # Run inference
        with torch.no_grad():
            scene_codes = self.system({"rgb_cond": rgb_cond, "c2w_cond": c2w_cond})
            exporter_output = self.system.exporter([f"{i:02d}" for i in range(rgb_cond.shape[0])], scene_codes)

        # Save output
        save_dir = os.path.join("./outputs", datetime.now().strftime("@%Y%m%d-%H%M%S"))
        os.makedirs(save_dir, exist_ok=True)
        self.system.set_save_dir(save_dir)

        for out in exporter_output:
            save_func_name = f"save_{out.save_type}"
            save_func = getattr(self.system, save_func_name)
            save_func(f"{out.save_name}", **out.params)

        if back_proj:
            refine(save_dir, new_image[1], new_image[0], new_image[3], new_image[2])

        new_obj = traverse(save_dir, back_proj)
        new_obj.export(f'{save_dir}/output.obj', file_type='obj')

        gltf = GLTF2().load(f'{save_dir}/output.glb')
        for material in gltf.materials:
            if material.pbrMetallicRoughness:
                material.pbrMetallicRoughness.baseColorFactor = [1.0, 1.0, 1.0, 100.0]
                material.pbrMetallicRoughness.metallicFactor = 0.0
                material.pbrMetallicRoughness.roughnessFactor = 1.0
        gltf.save(f'{save_dir}/output.glb')

        return save_dir, f"{save_dir}/output.obj", f"{save_dir}/output.glb"

inferapi = Inference_API()

# Define the interface
with gr.Blocks() as demo:
    gr.Markdown("# [SIGGRAPH'24] CharacterGen: Efficient 3D Character Generation from Single Images with Multi-View Pose Calibration")
    gr.Markdown("# 3D Stage: Four View Images to 3D Mesh")
    with gr.Row(variant="panel"):
        with gr.Column():
            with gr.Row():
                img_input0 = gr.Image(type="pil", label="Back Image", image_mode="RGBA", width=256, height=384)
                img_input1 = gr.Image(type="pil", label="Front Image", image_mode="RGBA", width=256, height=384)
            with gr.Row():
                img_input2 = gr.Image(type="pil", label="Right Image", image_mode="RGBA", width=256, height=384)
                img_input3 = gr.Image(type="pil", label="Left Image", image_mode="RGBA", width=256, height=384)
            with gr.Row():
                gr.Examples(
                    examples=
                    [["material/examples/1/1.png",
                    "material/examples/1/2.png",
                    "material/examples/1/3.png",
                    "material/examples/1/4.png"]],
                    label="Example Images",
                    inputs=[img_input0, img_input1, img_input2, img_input3]
                )
        with gr.Column():
            with gr.Row():
                back_proj = gr.Checkbox(label="Back Projection")
                submit_button = gr.Button("Process")
            output_dir = gr.Textbox(label="Output Directory")
            with gr.Column():
                with gr.Tab("GLB"):
                    output_model_glb = gr.Model3D( label="Output Model (GLB Format)", height = 768)
                    gr.Markdown("Note: The model shown here has a darker appearance. Download to get correct results.")
                with gr.Tab("OBJ"):
                    output_model_obj = gr.Model3D( label="Output Model (OBJ Format)", height = 768)
                    gr.Markdown("Note: The model shown here is flipped. Download to get correct results.")

    submit_button.click(inferapi.process_images, inputs=[img_input0, img_input1, img_input2, img_input3, back_proj],
                        outputs=[output_dir, output_model_obj, output_model_glb])

# Run the interface
demo.launch()