Alessio Grancini commited on
Commit
55bd8a1
·
verified ·
1 Parent(s): dfd6eee

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +23 -9
app.py CHANGED
@@ -15,19 +15,33 @@ from point_cloud_generator import display_pcd
15
  # params
16
  CANCEL_PROCESSING = False
17
 
18
- # Initialize models (but actual loading happens in decorated functions)
19
  img_seg = ImageSegmenter(model_type="yolov8s-seg")
20
  depth_estimator = MonocularDepthEstimator(model_type="midas_v21_small_256")
21
 
22
- @spaces.GPU(duration=30) # Adjust duration based on your needs
23
  def process_image(image):
24
- image = utils.resize(image)
25
- image_segmentation, objects_data = img_seg.predict(image)
26
- depthmap, depth_colormap = depth_estimator.make_prediction(image)
27
- dist_image = utils.draw_depth_info(image, depthmap, objects_data)
28
- objs_pcd = utils.generate_obj_pcd(depthmap, objects_data)
29
- plot_fig = display_pcd(objs_pcd)
30
- return image_segmentation, depth_colormap, dist_image, plot_fig
 
 
 
 
 
 
 
 
 
 
 
 
 
 
31
 
32
  @spaces.GPU(duration=30)
33
  def test_process_img(image):
 
15
  # params
16
  CANCEL_PROCESSING = False
17
 
18
+ # Initialize models (but don't load them yet)
19
  img_seg = ImageSegmenter(model_type="yolov8s-seg")
20
  depth_estimator = MonocularDepthEstimator(model_type="midas_v21_small_256")
21
 
22
+ @spaces.GPU
23
  def process_image(image):
24
+ try:
25
+ print("Starting image processing")
26
+ image = utils.resize(image)
27
+ print("Image resized")
28
+
29
+ # Models will be loaded here when needed
30
+ image_segmentation, objects_data = img_seg.predict(image)
31
+ print("Segmentation complete")
32
+
33
+ depthmap, depth_colormap = depth_estimator.make_prediction(image)
34
+ print("Depth estimation complete")
35
+
36
+ dist_image = utils.draw_depth_info(image, depthmap, objects_data)
37
+ objs_pcd = utils.generate_obj_pcd(depthmap, objects_data)
38
+ plot_fig = display_pcd(objs_pcd)
39
+ return image_segmentation, depth_colormap, dist_image, plot_fig
40
+ except Exception as e:
41
+ print(f"Error in process_image: {str(e)}")
42
+ import traceback
43
+ print(traceback.format_exc())
44
+ raise
45
 
46
  @spaces.GPU(duration=30)
47
  def test_process_img(image):