A19grey commited on
Commit
3f2f727
1 Parent(s): 7ce37c4

serialize depth then 3D so timeouts for 3D dont break whole system if user not want 3D

Browse files
Files changed (1) hide show
  1. app.py +37 -17
app.py CHANGED
@@ -20,7 +20,7 @@ print(f"Timm version: {timm.__version__}")
20
 
21
  subprocess.run(["bash", "get_pretrained_models.sh"])
22
 
23
- @spaces.GPU(duration=20)
24
  def load_model_and_predict(image_path):
25
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
26
  model, transform = depth_pro.create_model_and_transforms()
@@ -195,7 +195,7 @@ def remove_thin_features(mesh, thickness_threshold=0.01):
195
 
196
  return mesh
197
 
198
- @spaces.GPU(duration=20)
199
  def regenerate_3d_model(depth_csv, image_path, focallength_px, simplification_factor, smoothing_iterations, thin_threshold):
200
  # Load depth from CSV
201
  depth = np.loadtxt(depth_csv, delimiter=',')
@@ -208,7 +208,7 @@ def regenerate_3d_model(depth_csv, image_path, focallength_px, simplification_fa
208
  print("regenerated!")
209
  return view_model_path, download_model_path
210
 
211
- @spaces.GPU(duration=20)
212
  def predict_depth(input_image):
213
  temp_file = None
214
  try:
@@ -238,18 +238,32 @@ def predict_depth(input_image):
238
  raw_depth_path = "raw_depth_map.csv"
239
  np.savetxt(raw_depth_path, depth, delimiter=',')
240
 
241
- view_model_path, download_model_path = generate_3d_model(depth, temp_file, focallength_px)
242
- print("Depth created - 3D model generated!")
243
- return output_path, f"Focal length: {focallength_px:.2f} pixels", raw_depth_path, view_model_path, download_model_path, temp_file, focallength_px
244
  except Exception as e:
245
  import traceback
246
  error_message = f"An error occurred: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
247
  print(error_message)
248
- return None, error_message, None, None, None, None, None
249
  finally:
250
  if temp_file and os.path.exists(temp_file):
251
  os.remove(temp_file)
252
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
253
  def get_last_commit_timestamp():
254
  try:
255
  timestamp = subprocess.check_output(['git', 'log', '-1', '--format=%cd', '--date=iso']).decode('utf-8').strip()
@@ -267,12 +281,10 @@ with gr.Blocks() as iface:
267
  "An enhanced demo that creates a textured 3D model from the input image and depth map.\n\n"
268
  "Forked from https://huggingface.co/spaces/akhaliq/depth-pro and model from https://huggingface.co/apple/DepthPro\n"
269
  "**Instructions:**\n"
270
- "1. Upload an image.\n"
271
- "2. The app will predict the depth map, display it, and provide the focal length.\n"
272
- "3. Download the raw depth data as a CSV file.\n"
273
- "4. View the generated 3D model textured with the original image.\n"
274
- "5. Adjust parameters and click 'Regenerate 3D Model' to update the model.\n"
275
- "6. Download the 3D model as an OBJ file if desired.\n\n"
276
  f"Last updated: {last_updated}"
277
  )
278
 
@@ -283,6 +295,8 @@ with gr.Blocks() as iface:
283
  focal_length = gr.Textbox(label="Focal Length")
284
  raw_depth_csv = gr.File(label="Download Raw Depth Map (CSV)")
285
 
 
 
286
  with gr.Row():
287
  view_3d_model = gr.Model3D(label="View 3D Model")
288
  download_3d_model = gr.File(label="Download 3D Model (OBJ)")
@@ -293,22 +307,28 @@ with gr.Blocks() as iface:
293
  thin_threshold = gr.Slider(minimum=0.001, maximum=0.1, value=0.01, step=0.001, label="Thin Feature Threshold")
294
 
295
  regenerate_button = gr.Button("Regenerate 3D Model")
 
296
 
297
  # Hidden components to store intermediate results
298
- hidden_depth_csv = gr.State()
299
  hidden_image_path = gr.State()
300
  hidden_focal_length = gr.State()
301
 
302
  input_image.change(
303
  predict_depth,
304
  inputs=[input_image],
305
- outputs=[depth_map, focal_length, raw_depth_csv, view_3d_model, download_3d_model, hidden_image_path, hidden_focal_length]
 
 
 
 
 
 
306
  )
307
 
308
  regenerate_button.click(
309
- regenerate_3d_model,
310
  inputs=[raw_depth_csv, hidden_image_path, hidden_focal_length, simplification_factor, smoothing_iterations, thin_threshold],
311
- outputs=[view_3d_model, download_3d_model]
312
  )
313
 
314
  # Launch the Gradio interface with sharing enabled
 
20
 
21
  subprocess.run(["bash", "get_pretrained_models.sh"])
22
 
23
+ @spaces.GPU(duration=30)
24
  def load_model_and_predict(image_path):
25
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
26
  model, transform = depth_pro.create_model_and_transforms()
 
195
 
196
  return mesh
197
 
198
+ @spaces.GPU(duration=30) # Increased duration to 30 seconds
199
  def regenerate_3d_model(depth_csv, image_path, focallength_px, simplification_factor, smoothing_iterations, thin_threshold):
200
  # Load depth from CSV
201
  depth = np.loadtxt(depth_csv, delimiter=',')
 
208
  print("regenerated!")
209
  return view_model_path, download_model_path
210
 
211
+ @spaces.GPU(duration=30)
212
  def predict_depth(input_image):
213
  temp_file = None
214
  try:
 
238
  raw_depth_path = "raw_depth_map.csv"
239
  np.savetxt(raw_depth_path, depth, delimiter=',')
240
 
241
+ print("Depth map created!")
242
+ return output_path, f"Focal length: {focallength_px:.2f} pixels", raw_depth_path, temp_file, focallength_px
 
243
  except Exception as e:
244
  import traceback
245
  error_message = f"An error occurred: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
246
  print(error_message)
247
+ return None, error_message, None, None, None
248
  finally:
249
  if temp_file and os.path.exists(temp_file):
250
  os.remove(temp_file)
251
 
252
+ @spaces.GPU(duration=30)
253
+ def create_3d_model(depth_csv, image_path, focallength_px, simplification_factor, smoothing_iterations, thin_threshold):
254
+ try:
255
+ depth = np.loadtxt(depth_csv, delimiter=',')
256
+ view_model_path, download_model_path = generate_3d_model(
257
+ depth, image_path, focallength_px,
258
+ simplification_factor, smoothing_iterations, thin_threshold
259
+ )
260
+ print("3D model generated!")
261
+ return view_model_path, download_model_path, "3D model created successfully!"
262
+ except Exception as e:
263
+ error_message = f"An error occurred during 3D model creation: {str(e)}"
264
+ print(error_message)
265
+ return None, None, error_message
266
+
267
  def get_last_commit_timestamp():
268
  try:
269
  timestamp = subprocess.check_output(['git', 'log', '-1', '--format=%cd', '--date=iso']).decode('utf-8').strip()
 
281
  "An enhanced demo that creates a textured 3D model from the input image and depth map.\n\n"
282
  "Forked from https://huggingface.co/spaces/akhaliq/depth-pro and model from https://huggingface.co/apple/DepthPro\n"
283
  "**Instructions:**\n"
284
+ "1. Upload an image to generate the depth map.\n"
285
+ "2. Click 'Generate 3D Model' to create the 3D visualization.\n"
286
+ "3. Adjust parameters and click 'Regenerate 3D Model' to update the model.\n"
287
+ "4. Download the raw depth data as a CSV file or the 3D model as an OBJ file if desired.\n\n"
 
 
288
  f"Last updated: {last_updated}"
289
  )
290
 
 
295
  focal_length = gr.Textbox(label="Focal Length")
296
  raw_depth_csv = gr.File(label="Download Raw Depth Map (CSV)")
297
 
298
+ generate_3d_button = gr.Button("Generate 3D Model")
299
+
300
  with gr.Row():
301
  view_3d_model = gr.Model3D(label="View 3D Model")
302
  download_3d_model = gr.File(label="Download 3D Model (OBJ)")
 
307
  thin_threshold = gr.Slider(minimum=0.001, maximum=0.1, value=0.01, step=0.001, label="Thin Feature Threshold")
308
 
309
  regenerate_button = gr.Button("Regenerate 3D Model")
310
+ model_status = gr.Textbox(label="3D Model Status")
311
 
312
  # Hidden components to store intermediate results
 
313
  hidden_image_path = gr.State()
314
  hidden_focal_length = gr.State()
315
 
316
  input_image.change(
317
  predict_depth,
318
  inputs=[input_image],
319
+ outputs=[depth_map, focal_length, raw_depth_csv, hidden_image_path, hidden_focal_length]
320
+ )
321
+
322
+ generate_3d_button.click(
323
+ create_3d_model,
324
+ inputs=[raw_depth_csv, hidden_image_path, hidden_focal_length, simplification_factor, smoothing_iterations, thin_threshold],
325
+ outputs=[view_3d_model, download_3d_model, model_status]
326
  )
327
 
328
  regenerate_button.click(
329
+ create_3d_model,
330
  inputs=[raw_depth_csv, hidden_image_path, hidden_focal_length, simplification_factor, smoothing_iterations, thin_threshold],
331
+ outputs=[view_3d_model, download_3d_model, model_status]
332
  )
333
 
334
  # Launch the Gradio interface with sharing enabled