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