Spaces:
Runtime error
Runtime error
File size: 1,873 Bytes
077fc91 1689431 077fc91 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
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() |