update the code of gradio

gradio_app.py

@@ -4,18 +4,16 @@ import os
 import json
 import torch
 import sys
-import time
-import importlib
+import base64
 import numpy as np
-from omegaconf import OmegaConf
 from huggingface_hub import hf_hub_download
-from diffusers import DiffusionPipeline
 
 import PIL
 from PIL import Image
 from collections import OrderedDict
 import trimesh
 import rembg
+import requests
 import gradio as gr
 from typing import Any
 
@@ -50,7 +48,7 @@ _CITE_ = r"""
 If you find our work useful for your research or applications, please cite using this bibtex:
 ```bibtex
 @article{li2024craftsman,
-author    = {Weiyu Li and Jiarui Liu and Rui Chen and Yixun Liang and Xuelin Chen and Ping Tan and Xiaoxiao Long},
+author    = {Weiyu Li and Jiarui Liu and Hongyu Yan and Rui Chen and Yixun Liang and Xuelin Chen and Ping Tan and Xiaoxiao Long},
 title     = {CraftsMan: High-fidelity Mesh Generation with 3D Native Generation and Interactive Geometry Refiner},
 journal   = {arXiv preprint arXiv:2405.14979},
 year      = {2024},
@@ -67,7 +65,6 @@ If you have any questions, feel free to open a discussion or contact us at <b>we
 model = None
 cached_dir = None
 
-
 def check_input_image(input_image):
     if input_image is None:
         raise gr.Error("No image uploaded!")
@@ -87,7 +84,7 @@ class RMBG(object):
             if image.mode == "RGBA" and image.getextrema()[3][0] < 255:
                 # explain why current do not rm bg
                 print("alhpa channl not enpty, skip remove background, using alpha channel as mask")
-                background = Image.new("RGBA", image.size, background_color)
+                background = Image.new("RGBA", image.size, (*background_color, 0))
                 image = Image.alpha_composite(background, image)
                 do_remove = False
             do_remove = do_remove or force
@@ -97,11 +94,11 @@ class RMBG(object):
             # calculate the min bbox of the image
             alpha = image.split()[-1]
             image = image.crop(alpha.getbbox())
-
             return image
         return _rembg_remove(input_image, None, force_remove=True)
 
-    def run(self, rm_type, image, foreground_ratio, background_choice, background_color=(0, 0, 0, 0)):
+    def run(self, rm_type, image_file, foreground_ratio, background_choice, background_color=(255, 255, 255)):
+        image = Image.open(image_file)
         if "Original" in background_choice:
             return image
         else:
@@ -120,7 +117,7 @@ class RMBG(object):
             # Resize the image while maintaining the aspect ratio
             resized_image = image.resize(new_size)
             # Create a new image with the original size and white background
-            padded_image = PIL.Image.new("RGBA", image.size, (0, 0, 0, 0))
+            padded_image = PIL.Image.new("RGBA", image.size, (*background_color, 0))
             paste_position = ((image.width - resized_image.width) // 2, (image.height - resized_image.height) // 2)
             padded_image.paste(resized_image, paste_position)
 
@@ -129,24 +126,29 @@ class RMBG(object):
             if width == height:
                 return padded_image
             new_size = (max(width, height), max(width, height))
-            image = PIL.Image.new("RGBA", new_size, (0, 0, 0, 0))
+            image = PIL.Image.new("RGBA", new_size, (*background_color, 1))
             paste_position = ((new_size[0] - width) // 2, (new_size[1] - height) // 2)
             image.paste(padded_image, paste_position)
-            return image
+            filepath = tempfile.NamedTemporaryFile(suffix=f".png", delete=False).name
+            image.save(filepath)
+            return filepath
 
-@spaces.GPU
+# @spaces.GPU
 def image2mesh(image: Any, 
                more: bool = False,
                scheluder_name: str ="DDIMScheduler",
                guidance_scale: int = 7.5,
-               steps: int = 30,
+               steps: int = 50,
                seed: int = 4,
                target_face_count: int = 2000,
                octree_depth: int = 7):
     
+    # global rmbg
+    # processed_image = rmbg.run(rm_type, image, foreground_ratio, background_choice)
+    processed_image = Image.open(image)
     sample_inputs = {
         "image": [
-            image
+            processed_image
         ]
     }
 
@@ -169,7 +171,7 @@ def image2mesh(image: Any,
     assert len(mesh_outputs) == 1, "Only support single mesh output for gradio demo"
     mesh = trimesh.Trimesh(mesh_outputs[0][0], mesh_outputs[0][1])
     # filepath = f"{cached_dir}/{time.time()}.obj"
-    filepath = tempfile.NamedTemporaryFile(suffix=f".obj", delete=False).name
+    filepath = tempfile.NamedTemporaryFile(suffix=f".glb", delete=False).name
     mesh.export(filepath, include_normals=True)
 
     if 'Remesh' in more:
@@ -181,6 +183,30 @@ def image2mesh(image: Any,
     
     return filepath
 
+def mesh2texture(mesh_file, image_file):
+    headers = {'Content-Type': 'application/json'}
+    server_url = "114.249.238.184:34119"
+    with open(image_file, 'rb') as f:
+        image_bytes = f.read()
+    with open(mesh_file, 'rb') as f:
+        mesh_bytes = f.read()
+    request = {
+        'png_base64_image': base64.b64encode(image_bytes).decode('utf-8'),
+        'glb_base64_mesh': base64.b64encode(mesh_bytes).decode('utf-8'),
+    }
+    response = requests.post(
+        url=f"http://{server_url}/generate_texture", 
+        headers=headers, 
+        data=json.dumps(request),
+    ).json()
+    mesh_bytes = base64.b64decode(response['glb_base64_mesh'])
+    filepath = tempfile.NamedTemporaryFile(suffix=f".glb", delete=False).name
+    with open(filepath, 'wb') as f:
+        f.write(mesh_bytes)
+
+    return filepath
+
+
 if __name__=="__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--model_path", type=str, default="", help="Path to the object file",)
@@ -238,39 +264,23 @@ if __name__=="__main__":
                             label="Image Input",
                             image_mode="RGBA",
                             sources="upload",
-                            type="pil",
+                            type="filepath",
                         )
                 run_btn = gr.Button('Generate', variant='primary', interactive=True)
 
                 with gr.Row():
                     gr.Markdown('''Try a different <b>seed and MV Model</b> for better results. Good Luck :)''')
                 with gr.Row():
-                    seed = gr.Number(0, label='Seed', show_label=True)
+                    seed = gr.Number(42, label='Seed', show_label=True)
                     more = gr.CheckboxGroup(["Remesh"], label="More", show_label=False)
                     target_face_count = gr.Number(2000, label='Target Face Count', show_label=True)
 
-                with gr.Row():
-                    gr.Examples(
-                        examples=[os.path.join("./examples", i) for i in os.listdir("./examples")],
-                        inputs=[image_input],
-                        examples_per_page=8
-                    )
 
-            with gr.Column(scale=4):
-                with gr.Row():
-                    output_model_obj = gr.Model3D(
-                        label="Output Model (OBJ Format)",
-                        camera_position=(90.0, 90.0, 3.5),
-                        interactive=False,
-                    )
-                with gr.Row():
-                    gr.Markdown('''*please note that the model is fliped due to the gradio viewer, please download the obj file and you will get the correct orientation.''')
-                    
                 with gr.Accordion('Advanced options', open=False):
                     with gr.Row():
                         background_choice = gr.Dropdown(label="Backgroud Choice", value="Auto Remove Background",choices=list(background_choice.keys()))
                         rmbg_type = gr.Dropdown(label="Backgroud Remove Type", value="rembg",choices=['sam', "rembg"])
-                        foreground_ratio = gr.Slider(label="Foreground Ratio", value=1.0, minimum=0.5, maximum=1.0, step=0.01)
+                        foreground_ratio = gr.Slider(label="Foreground Ratio", value=0.95, minimum=0.5, maximum=1.0, step=0.01)
 
                     with gr.Row():
                         guidance_scale = gr.Number(label="3D Guidance Scale", value=7.5, minimum=3.0, maximum=10.0)
@@ -279,10 +289,34 @@ if __name__=="__main__":
                     with gr.Row():
                         scheduler = gr.Dropdown(label="scheluder", value="DDIMScheduler",choices=list(scheluder_dict.keys()))
                         octree_depth = gr.Slider(label="Octree Depth", value=7, minimum=4, maximum=8, step=1)
+
+                with gr.Row():
+                    gr.Examples(
+                        examples=[os.path.join("./assets/examples", i) for i in os.listdir("./assets/examples")],
+                        inputs=[image_input],
+                        examples_per_page=8
+                    )
+
+            with gr.Column(scale=4):
+                with gr.Row():
+                    output_model_obj = gr.Model3D(
+                        label="Output Model (GLB Format)",
+                        camera_position=(90.0, 90.0, 3.5),
+                        interactive=False,
+                    )
+                with gr.Row():
+                    output_model_tex = gr.Model3D(
+                        label="Output Textured Model (GLB Format)",
+                        camera_position=(90.0, 90.0, 3.5),
+                        interactive=False,
+                    )
+                # with gr.Row():
+                #     gr.Markdown('''*please note that the model is fliped due to the gradio viewer, please download the obj file and you will get the correct orientation.''')
                     
         gr.Markdown(_CITE_)
 
-        outputs = [output_model_obj]
+        # outputs = [output_model_obj]
+        # outputs_tex = [output_model_tex]
         rmbg = RMBG()
         
         # model = load_model(ckpt_path, config_path, device)
@@ -304,7 +338,13 @@ if __name__=="__main__":
                     ).success(
                             fn=image2mesh, 
                             inputs=[image_input, more, scheduler, guidance_scale, steps, seed, target_face_count, octree_depth],
-                            outputs=outputs, 
-                            api_name="generate_img2obj")
+                            outputs=[output_model_obj], 
+                            api_name="generate_img2obj"
+                    ).success(
+                            fn=mesh2texture, 
+                            inputs=[output_model_obj, image_input],
+                            outputs=[output_model_tex], 
+                            api_name="generate_obj2tex"
+                    )
         
         demo.queue().launch(share=True, allowed_paths=[args.cached_dir])