| 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() |