import gradio as gr
from tensorflow.keras.models import load_model
import numpy as np
import cv2
from huggingface_hub import hf_hub_download

# Load the model
model_path = hf_hub_download(repo_id="SalmanAboAraj/Tooth1", filename="unet_model.h5")
model = load_model(model_path)

# Define prediction function
def predict(image):
    original_height, original_width, _ = image.shape
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    image = cv2.resize(image, (128, 128))
    image = np.expand_dims(image, axis=0)
    image = np.expand_dims(image, axis=-1)
    image = image / 255.0
    mask = model.predict(image)
    mask = (mask[0] > 0.5).astype(np.uint8) * 255
    mask = cv2.resize(mask, (original_width, original_height))
    return mask

# Gradio Blocks interface
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="numpy", label="Input X-ray Image")
            mask_output = gr.Image(type="numpy", label="Annotation Mask")
        with gr.Column():
            gr.Markdown("# Tooth Segmentation Model")
            gr.Markdown("Upload a dental X-ray image to generate the annotation mask.")
    
    # Linking inputs and outputs with the prediction function
    image_input.change(predict, inputs=image_input, outputs=mask_output)

# Launch the app
if __name__ == "__main__":
    demo.launch()