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

# إصلاح مشكلة المسارات في Windows
plt = platform.system()
pathlib.WindowsPath = pathlib.PosixPath

# تحميل النموذج من Hugging Face
model_path = hf_hub_download(repo_id="SalmanAboAraj/Tooth1", filename="unet_model_256.h5")
model = load_model(model_path)

# تعريف ألوان الفئات (من بياناتك)
COLOR_MAP = {
    0: (0, 0, 0),        # Background
    1: (215, 179, 255),  # Bone
    2: (246, 51, 81),    # Cavity
    3: (58, 132, 255),   # Crown
    4: (134, 202, 218),  # Dental Implant
    5: (221, 195, 130),  # Dental Implant Crown
    6: (255, 255, 127),  # Dentin
    7: (255, 255, 255),  # Enamel
    8: (1, 13, 27),      # Filling Metal
    9: (0, 133, 255),    # Filling Non-Metal
    10: (24, 250, 143),  # Periapical Radiolucence
    11: (255, 105, 248), # Pulp
    12: (17, 253, 231),  # Root Canal
    13: (255, 146, 119), # Sinus
    14: (131, 224, 112)  # Missing
}

# وظيفة التنبؤ
def predict(image):
    original_height, original_width, _ = image.shape

    # تحويل الصورة إلى حجم يناسب النموذج
    image_resized = cv2.resize(image, (256, 256))
    image_resized = image_resized / 255.0
    image_resized = np.expand_dims(image_resized, axis=0)

    # التنبؤ بالقناع
    mask_pred = model.predict(image_resized)[0]  # الإخراج يكون (256, 256, 15)

    # تحويل القناع إلى صورة بألوان الفئات
    mask_class = np.argmax(mask_pred, axis=-1)  # إخراج التصنيفات (256, 256)
    mask_colored = np.zeros((256, 256, 3), dtype=np.uint8)

    for class_idx, color in COLOR_MAP.items():
        mask_colored[mask_class == class_idx] = color

    # إرجاع القناع بنفس أبعاد الصورة الأصلية
    mask_final = cv2.resize(mask_colored, (original_width, original_height), interpolation=cv2.INTER_NEAREST)
    return mask_final

# إنشاء واجهة Gradio
with gr.Blocks() as iface: 
    gr.Markdown("# Tooth Segmentation Model") 
    gr.Markdown("Upload a dental X-ray image to generate the annotation mask.") 
    with gr.Row(): 
        with gr.Column(): 
            input_image = gr.Image(label="Input X-ray Image", type="numpy") 
            submit_button = gr.Button("Predict") 
        with gr.Column(): 
            output_image = gr.Image(label="Annotation Mask", type="numpy") 
            
    submit_button.click(fn=predict, inputs=input_image, outputs=output_image)
if __name__ == "__main__":
    iface.launch()