yanranxiaoxi commited on
Commit
7d68b4d
1 Parent(s): f73d8e1

chore: end try

Browse files
Files changed (1) hide show
  1. app.py +13 -25
app.py CHANGED
@@ -7,7 +7,6 @@ import trimesh
7
  import mcubes
8
  import imageio
9
  from torchvision.utils import save_image
10
- from torchvision.transforms import ToPILImage
11
  from PIL import Image
12
  from transformers import AutoModel, AutoConfig
13
  from rembg import remove, new_session
@@ -184,10 +183,6 @@ def generate_mesh(image, source_size=512, render_size=384, mesh_size=512, export
184
 
185
  with torch.no_grad():
186
  planes = model_wrapper.forward(image, source_camera)
187
- planes_pil_image = np.concatenate([
188
- np.concatenate([planes[0][1], planes[0][2]], axis=1),
189
- np.concatenate([planes[0][0], planes[0][0]], axis=1),
190
- ], axis=0)
191
 
192
  if export_mesh:
193
  grid_out = model_wrapper.model.synthesizer.forward_grid(planes=planes, grid_size=mesh_size)
@@ -201,7 +196,7 @@ def generate_mesh(image, source_size=512, render_size=384, mesh_size=512, export
201
  mesh_path = "xiaoxis_mesh.obj"
202
  mesh.export(mesh_path, 'obj')
203
 
204
- return planes_pil_image, mesh_path
205
 
206
  if export_video:
207
  render_cameras = _default_render_cameras(batch_size=1).to(model_wrapper.device)
@@ -225,23 +220,23 @@ def generate_mesh(image, source_size=512, render_size=384, mesh_size=512, export
225
  video_path = "xiaoxis_video.mp4"
226
  imageio.mimwrite(video_path, frames, fps=fps)
227
 
228
- return planes_pil_image, video_path
229
 
230
- return planes_pil_image, None
231
 
232
  return None, None
233
 
234
  def step_1_generate_planes(image):
235
- planes_pil_image, _ = generate_mesh(image)
236
- return planes_pil_image
237
 
238
  def step_2_generate_obj(image):
239
- planes_pil_image, mesh_path = generate_mesh(image, export_mesh=True)
240
- return planes_pil_image, mesh_path, mesh_path
241
 
242
  def step_3_generate_video(image):
243
- planes_pil_image, video_path = generate_mesh(image, export_video=True)
244
- return planes_pil_image, video_path
245
 
246
  # 从 assets 文件夹中设置示例文件,并限制最多读取 10 个文件
247
  example_folder = "assets"
@@ -261,15 +256,9 @@ with gr.Blocks() as demo:
261
  with gr.Column():
262
  img_input = gr.Image(type="pil", label="输入图像")
263
  examples_component = gr.Examples(examples=examples, inputs=img_input, outputs=None, examples_per_page=5)
264
- generate_planes_button = gr.Button("生成平面图")
265
  generate_mesh_button = gr.Button("生成模型")
266
  generate_video_button = gr.Button("生成视频")
267
  with gr.Column():
268
- planes_output = gr.Image(
269
- label="平面图",
270
- type="pil",
271
- interactive=False
272
- )
273
  model_output = LitModel3D(
274
  clear_color=[0, 0, 0, 0], # 可调整背景颜色,以获得更好的对比度
275
  label="模型可视化",
@@ -294,15 +283,14 @@ with gr.Blocks() as demo:
294
  def clear_model_viewer():
295
  """在加载新模型前重置 Gradio。"""
296
  update_output = gr.update(value=None)
297
- return update_output, update_output
298
 
299
  # 首先清除输出的数据
300
- img_input.change(clear_model_viewer, inputs=None, outputs=[planes_output, model_output])
301
 
302
  # 然后生成模型和视频
303
- generate_planes_button.click(step_1_generate_planes, inputs=img_input, outputs=planes_output)
304
- generate_mesh_button.click(step_2_generate_obj, inputs=img_input, outputs=[planes_output, obj_file_output, model_output])
305
- generate_video_button.click(step_3_generate_video, inputs=img_input, outputs=[planes_output, video_file_output])
306
 
307
  demo.launch(
308
  auth=(os.environ.get('AUTH_USERNAME'), os.environ.get('AUTH_PASSWORD'))
 
7
  import mcubes
8
  import imageio
9
  from torchvision.utils import save_image
 
10
  from PIL import Image
11
  from transformers import AutoModel, AutoConfig
12
  from rembg import remove, new_session
 
183
 
184
  with torch.no_grad():
185
  planes = model_wrapper.forward(image, source_camera)
 
 
 
 
186
 
187
  if export_mesh:
188
  grid_out = model_wrapper.model.synthesizer.forward_grid(planes=planes, grid_size=mesh_size)
 
196
  mesh_path = "xiaoxis_mesh.obj"
197
  mesh.export(mesh_path, 'obj')
198
 
199
+ return None, mesh_path
200
 
201
  if export_video:
202
  render_cameras = _default_render_cameras(batch_size=1).to(model_wrapper.device)
 
220
  video_path = "xiaoxis_video.mp4"
221
  imageio.mimwrite(video_path, frames, fps=fps)
222
 
223
+ return None, video_path
224
 
225
+ return planes, None
226
 
227
  return None, None
228
 
229
  def step_1_generate_planes(image):
230
+ planes, _ = generate_mesh(image)
231
+ return planes
232
 
233
  def step_2_generate_obj(image):
234
+ _, mesh_path = generate_mesh(image, export_mesh=True)
235
+ return mesh_path, mesh_path
236
 
237
  def step_3_generate_video(image):
238
+ _, video_path = generate_mesh(image, export_video=True)
239
+ return video_path
240
 
241
  # 从 assets 文件夹中设置示例文件,并限制最多读取 10 个文件
242
  example_folder = "assets"
 
256
  with gr.Column():
257
  img_input = gr.Image(type="pil", label="输入图像")
258
  examples_component = gr.Examples(examples=examples, inputs=img_input, outputs=None, examples_per_page=5)
 
259
  generate_mesh_button = gr.Button("生成模型")
260
  generate_video_button = gr.Button("生成视频")
261
  with gr.Column():
 
 
 
 
 
262
  model_output = LitModel3D(
263
  clear_color=[0, 0, 0, 0], # 可调整背景颜色,以获得更好的对比度
264
  label="模型可视化",
 
283
  def clear_model_viewer():
284
  """在加载新模型前重置 Gradio。"""
285
  update_output = gr.update(value=None)
286
+ return update_output
287
 
288
  # 首先清除输出的数据
289
+ img_input.change(clear_model_viewer, inputs=None, outputs=model_output)
290
 
291
  # 然后生成模型和视频
292
+ generate_mesh_button.click(step_2_generate_obj, inputs=img_input, outputs=[obj_file_output, model_output])
293
+ generate_video_button.click(step_3_generate_video, inputs=img_input, outputs=video_file_output)
 
294
 
295
  demo.launch(
296
  auth=(os.environ.get('AUTH_USERNAME'), os.environ.get('AUTH_PASSWORD'))