yanranxiaoxi commited on
Commit
cc398fb
·
verified ·
1 Parent(s): 66937e7

chore: remove try

Browse files
Files changed (1) hide show
  1. app.py +47 -52
app.py CHANGED
@@ -178,58 +178,53 @@ def _default_render_cameras(batch_size: int = 1):
178
 
179
  @spaces.GPU
180
  def generate_mesh(image, source_size=512, render_size=384, mesh_size=512, export_mesh=False, export_video=False, fps=30):
181
- try:
182
- image = preprocess_image(image, source_size).to(model_wrapper.device)
183
- source_camera = _default_source_camera(batch_size=1).to(model_wrapper.device)
184
-
185
- with torch.no_grad():
186
- planes = model_wrapper.forward(image, source_camera)
187
-
188
- if export_mesh:
189
- grid_out = model_wrapper.model.synthesizer.forward_grid(planes=planes, grid_size=mesh_size)
190
- vtx, faces = mcubes.marching_cubes(grid_out['sigma'].float().squeeze(0).squeeze(-1).cpu().numpy(), 1.0)
191
- vtx = vtx / (mesh_size - 1) * 2 - 1
192
- vtx_tensor = torch.tensor(vtx, dtype=torch.float32, device=model_wrapper.device).unsqueeze(0)
193
- vtx_colors = model_wrapper.model.synthesizer.forward_points(planes, vtx_tensor)['rgb'].float().squeeze(0).cpu().numpy()
194
- vtx_colors = (vtx_colors * 255).astype(np.uint8)
195
- mesh = trimesh.Trimesh(vertices=vtx, faces=faces, vertex_colors=vtx_colors)
196
-
197
- mesh_path = "xiaoxis_mesh.obj"
198
- mesh.export(mesh_path, 'obj')
199
-
200
- return None, mesh_path
201
-
202
- if export_video:
203
- render_cameras = _default_render_cameras(batch_size=1).to(model_wrapper.device)
204
- frames = []
205
- chunk_size = 1
206
- for i in range(0, render_cameras.shape[1], chunk_size):
207
- frame_chunk = model_wrapper.model.synthesizer(
208
- planes,
209
- render_cameras[:, i:i + chunk_size],
210
- render_size,
211
- render_size,
212
- 0,
213
- 0
214
- )
215
- frames.append(frame_chunk['images_rgb'])
216
-
217
- frames = torch.cat(frames, dim=1)
218
- frames = frames.squeeze(0)
219
- frames = (frames.permute(0, 2, 3, 1).cpu().numpy() * 255).astype(np.uint8)
220
-
221
- video_path = "xiaoxis_video.mp4"
222
- imageio.mimwrite(video_path, frames, fps=fps)
223
-
224
- return None, video_path
225
-
226
- return planes, None
227
-
228
- return None, None
229
-
230
- except Exception as e:
231
- print(f"Error: {str(e)}")
232
- return None, None
233
 
234
  def step_1_generate_planes(image):
235
  planes, _ = generate_mesh(image)
 
178
 
179
  @spaces.GPU
180
  def generate_mesh(image, source_size=512, render_size=384, mesh_size=512, export_mesh=False, export_video=False, fps=30):
181
+ image = preprocess_image(image, source_size).to(model_wrapper.device)
182
+ source_camera = _default_source_camera(batch_size=1).to(model_wrapper.device)
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)
189
+ vtx, faces = mcubes.marching_cubes(grid_out['sigma'].float().squeeze(0).squeeze(-1).cpu().numpy(), 1.0)
190
+ vtx = vtx / (mesh_size - 1) * 2 - 1
191
+ vtx_tensor = torch.tensor(vtx, dtype=torch.float32, device=model_wrapper.device).unsqueeze(0)
192
+ vtx_colors = model_wrapper.model.synthesizer.forward_points(planes, vtx_tensor)['rgb'].float().squeeze(0).cpu().numpy()
193
+ vtx_colors = (vtx_colors * 255).astype(np.uint8)
194
+ mesh = trimesh.Trimesh(vertices=vtx, faces=faces, vertex_colors=vtx_colors)
195
+
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)
203
+ frames = []
204
+ chunk_size = 1
205
+ for i in range(0, render_cameras.shape[1], chunk_size):
206
+ frame_chunk = model_wrapper.model.synthesizer(
207
+ planes,
208
+ render_cameras[:, i:i + chunk_size],
209
+ render_size,
210
+ render_size,
211
+ 0,
212
+ 0
213
+ )
214
+ frames.append(frame_chunk['images_rgb'])
215
+
216
+ frames = torch.cat(frames, dim=1)
217
+ frames = frames.squeeze(0)
218
+ frames = (frames.permute(0, 2, 3, 1).cpu().numpy() * 255).astype(np.uint8)
219
+
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)