jens commited on
Commit
2eca80e
·
1 Parent(s): 025dcd6
Files changed (1) hide show
  1. app.py +21 -7
app.py CHANGED
@@ -5,23 +5,37 @@ from inference import DepthPredictor, SegmentPredictor
5
  from utils import create_3d_obj, create_3d_pc, point_cloud
6
  import numpy as np
7
 
8
-
9
- def snap(image, video):
10
  depth_predictor = DepthPredictor()
11
  depth_result = depth_predictor.predict(image)
12
- rgb_gltf_path = create_3d_obj(np.array(image), depth_result, path='./rgb.gltf')
13
-
 
14
  segment_predictor = SegmentPredictor()
15
  sam_result = segment_predictor.predict(image)
16
- fig = point_cloud(np.array(sam_result), depth_result)
 
 
 
 
 
 
17
 
18
- return [image, depth_result, sam_result, rgb_gltf_path, fig]#[depth_result, gltf_path, gltf_path]
 
19
 
 
 
 
 
 
20
 
 
21
  demo = gr.Interface(
22
  snap,
23
  inputs=[gr.Image(source="webcam", tool=None, label="Input Image", type="pil"),
24
- gr.Video(source="webcam")],
 
25
  outputs=[gr.Image(label="RGB"),
26
  gr.Image(label="predicted depth"),
27
  gr.Image(label="predicted segmentation"),
 
5
  from utils import create_3d_obj, create_3d_pc, point_cloud
6
  import numpy as np
7
 
8
+ def produce_depth_map(image):
 
9
  depth_predictor = DepthPredictor()
10
  depth_result = depth_predictor.predict(image)
11
+ return depth_result
12
+
13
+ def produce_segmentation_map(image):
14
  segment_predictor = SegmentPredictor()
15
  sam_result = segment_predictor.predict(image)
16
+ return sam_result
17
+
18
+ def produce_3d_reconstruction(image):
19
+ depth_predictor = DepthPredictor()
20
+ depth_result = depth_predictor.predict(image)
21
+ rgb_gltf_path = create_3d_obj(np.array(image), depth_result, path='./rgb.gltf')
22
+ return rgb_gltf_path
23
 
24
+ def produce_point_cloud(depth_map, segmentation_map):
25
+ return point_cloud(np.array(segmentation_map), depth_map)
26
 
27
+ def snap(image, depth_map, segmentation_map):
28
+ depth_result = produce_depth_map(image) if depth_map else None
29
+ sam_result = produce_segmentation_map(image) if segmentation_map else None
30
+ rgb_gltf_path = produce_3d_reconstruction(image) if depth_map else None
31
+ point_cloud_fig = produce_point_cloud(depth_result, sam_result) if (segmentation_map and depth_map) else None
32
 
33
+ return [image, depth_result, sam_result, rgb_gltf_path, point_cloud_fig]
34
  demo = gr.Interface(
35
  snap,
36
  inputs=[gr.Image(source="webcam", tool=None, label="Input Image", type="pil"),
37
+ "checkbox",
38
+ "checkbox"],
39
  outputs=[gr.Image(label="RGB"),
40
  gr.Image(label="predicted depth"),
41
  gr.Image(label="predicted segmentation"),