| import gradio as gr |
| from segment_anything import SamAutomaticMaskGenerator, sam_model_registry |
| import supervision as sv |
| from inference import DepthPredictor, SegmentPredictor |
| from utils import create_3d_obj, create_3d_pc, point_cloud, generate_PCL |
| import numpy as np |
|
|
| def produce_depth_map(image): |
| depth_predictor = DepthPredictor() |
| depth_result = depth_predictor.predict(image) |
| return depth_result |
|
|
| def produce_segmentation_map(image): |
| segment_predictor = SegmentPredictor() |
| sam_result = segment_predictor.predict(image) |
| return sam_result |
|
|
| def produce_3d_reconstruction(image): |
| depth_predictor = DepthPredictor() |
| depth_result = depth_predictor.predict(image) |
| rgb_gltf_path = create_3d_obj(np.array(image), depth_result, path='./rgb.gltf') |
| return rgb_gltf_path |
|
|
| def produce_point_cloud(depth_map, segmentation_map): |
| return point_cloud(np.array(segmentation_map), depth_map) |
|
|
| def snap(image, depth_map, segmentation_map): |
| depth_result = produce_depth_map(image) if depth_map else None |
| sam_result = produce_segmentation_map(image) if segmentation_map else None |
| rgb_gltf_path = produce_3d_reconstruction(image) if depth_map else None |
| point_cloud_fig = produce_point_cloud(depth_result, sam_result) if (segmentation_map and depth_map) else None |
|
|
| return [image, depth_result, sam_result, rgb_gltf_path, point_cloud_fig] |
| demo = gr.Interface( |
| snap, |
| inputs=[gr.Image(source="webcam", tool=None, label="Input Image", type="pil"), |
| "checkbox", |
| "checkbox"], |
| outputs=[gr.Image(label="RGB"), |
| gr.Image(label="predicted depth"), |
| gr.Image(label="predicted segmentation"), |
| gr.Model3D(label="3D mesh reconstruction - RGB", |
| clear_color=[1.0, 1.0, 1.0, 1.0]), |
| gr.Plot()] |
| ) |
|
|
| if __name__ == "__main__": |
| demo.launch() |