Spaces:
Runtime error
Runtime error
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() |