chore: planes to PIL image

app.py

@@ -7,6 +7,7 @@ import trimesh
 import mcubes
 import imageio
 from torchvision.utils import save_image
+import torchvision.transforms import ToPILImage
 from PIL import Image
 from transformers import AutoModel, AutoConfig
 from rembg import remove, new_session
@@ -183,6 +184,7 @@ def generate_mesh(image, source_size=512, render_size=384, mesh_size=512, export
 
     with torch.no_grad():
         planes = model_wrapper.forward(image, source_camera)
+        planes_pil_image = ToPILImage(planes)
 
         if export_mesh:
             grid_out = model_wrapper.model.synthesizer.forward_grid(planes=planes, grid_size=mesh_size)
@@ -196,7 +198,7 @@ def generate_mesh(image, source_size=512, render_size=384, mesh_size=512, export
             mesh_path = "xiaoxis_mesh.obj"
             mesh.export(mesh_path, 'obj')
 
-            return planes, mesh_path
+            return planes_pil_image, mesh_path
 
         if export_video:
             render_cameras = _default_render_cameras(batch_size=1).to(model_wrapper.device)
@@ -220,23 +222,23 @@ def generate_mesh(image, source_size=512, render_size=384, mesh_size=512, export
             video_path = "xiaoxis_video.mp4"
             imageio.mimwrite(video_path, frames, fps=fps)
 
-            return planes, video_path
+            return planes_pil_image, video_path
 
-        return planes, None
+        return planes_pil_image, None
 
     return None, None
 
 def step_1_generate_planes(image):
-    planes, _ = generate_mesh(image)
-    return planes
+    planes_pil_image, _ = generate_mesh(image)
+    return planes_pil_image
 
 def step_2_generate_obj(image):
-    planes, mesh_path = generate_mesh(image, export_mesh=True)
-    return planes, mesh_path, mesh_path
+    planes_pil_image, mesh_path = generate_mesh(image, export_mesh=True)
+    return planes_pil_image, mesh_path, mesh_path
 
 def step_3_generate_video(image):
-    planes, video_path = generate_mesh(image, export_video=True)
-    return planes, video_path
+    planes_pil_image, video_path = generate_mesh(image, export_video=True)
+    return planes_pil_image, video_path
 
 # 从 assets 文件夹中设置示例文件，并限制最多读取 10 个文件
 example_folder = "assets"