jens commited on
Commit
025dcd6
·
1 Parent(s): 371a984

Basic layout

Browse files
Files changed (1) hide show
  1. app.py +13 -39
app.py CHANGED
@@ -5,55 +5,29 @@ from inference import DepthPredictor, SegmentPredictor
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, video):
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 else None
32
 
33
- if video:
34
- # Add video processing here if needed
35
- pass
36
-
37
- return [image, depth_result, sam_result, rgb_gltf_path, point_cloud_fig]
38
 
39
- # Interface inputs
40
- image_input = gr.Image(source="webcam", tool=None, label="Input Image", type="pil")
41
- depth_map_button = gr.Button(label="Produce Depth Map", value=False)
42
- segmentation_map_button = gr.Button(label="Produce Segmentation Map", value=False)
43
- video_input = gr.Video(source="webcam")
44
 
45
- # Interface outputs
46
- output_image = gr.Image(label="RGB")
47
- output_depth_map = gr.Image(label="Predicted Depth")
48
- output_segmentation_map = gr.Image(label="Predicted Segmentation")
49
- output_3d_reconstruction = gr.Model3D(label="3D mesh reconstruction - RGB", clear_color=[1.0, 1.0, 1.0, 1.0])
50
- output_point_cloud = gr.Plot(label="Point Cloud")
51
 
52
- # Interface
53
  demo = gr.Interface(
54
  snap,
55
- inputs=[image_input, depth_map_button, segmentation_map_button, video_input],
56
- outputs=[output_image, output_depth_map, output_segmentation_map, output_3d_reconstruction, output_point_cloud]
 
 
 
 
 
 
57
  )
58
 
59
  if __name__ == "__main__":
 
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"),
28
+ gr.Model3D(label="3D mesh reconstruction - RGB",
29
+ clear_color=[1.0, 1.0, 1.0, 1.0]),
30
+ gr.Plot()]
31
  )
32
 
33
  if __name__ == "__main__":