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
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, point_cloud, video):
    if depth_map:
        depth_result = produce_depth_map(image)
    else:
        depth_result = None
        
    if segmentation_map:
        sam_result = produce_segmentation_map(image)
    else:
        sam_result = None
        
    if point_cloud:
        point_cloud_fig = produce_point_cloud(depth_result, sam_result)
    else:
        point_cloud_fig = None
        
    if video:
        # Add video processing here if needed
        pass

    return [image, depth_result, sam_result, point_cloud_fig]

# Interface inputs
image_input = gr.Image(source="webcam", tool=None, label="Input Image", type="pil")
depth_map_button = gr.Button(label="Produce Depth Map", value=False)
segmentation_map_button = gr.Button(label="Produce Segmentation Map", value=False)
point_cloud_button = gr.Button(label="Produce Point Cloud", value=False)
video_input = gr.Video(source="webcam")

# Interface outputs
output_image = gr.Image(label="RGB")
output_depth_map = gr.Image(label="Predicted Depth")
output_segmentation_map = gr.Image(label="Predicted Segmentation")
output_point_cloud = gr.Plot(label="Point Cloud")

# Interface
demo = gr.Interface(
    snap,
    inputs=[image_input, depth_map_button, segmentation_map_button, point_cloud_button, video_input],
    outputs=[output_image, output_depth_map, output_segmentation_map, None, output_point_cloud]
)

if __name__ == "__main__":
    demo.launch()