import gradio as gr
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing import image

# 모델 로드 (파일명 수정)
model = tf.keras.models.load_model("cell_organelle_classifier.keras")

# 모델의 class_indices 순서에 맞춘 리스트 (오타 수정 포함)
class_names = ['class_ER', 'class_chloroplast', 'class_golgi', 'class_mitochondria']

# 한글 라벨 매핑
kor_label_map = {
    'class_ER': '소포체',
    'class_chloroplast': '엽록체',
    'class_golgi': '골지체',
    'class_mitochondria': '미토콘드리아'
}

# 관련 기사 링크 (클래스명 키도 수정)
related_links = {
    'class_mitochondria': [
        "https://biz.chosun.com/science-chosun/medicine-health/2025/06/26/LWAWZ4KCDFB4NHXRPNSXZO4BZ4/",
        "https://www.joongang.co.kr/article/25326557"
    ],
    'class_ER': [
        "https://www.snu.ac.kr/research/highlights?md=v&bbsidx=155485",
        "https://m.dongascience.com/news.php?idx=68985"
    ],
    'class_golgi': [
        "https://www.chosun.com/economy/science/2024/07/16/YS7Q7BVSM7G5H7SSPUPFSTWH2I/",
        "https://www.snu.ac.kr/snunow/press?md=v&bbsidx=149578"
    ],
    'class_chloroplast': [
        "https://www.fnewstv.com/news/newsview.php?ncode=1065602102008791&dt=m",
        "https://www.hani.co.kr/arti/opinion/column/1205729.html"
    ]
}

def classify_image(img):
    img = img.resize((224, 224))
    img_array = image.img_to_array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    pred = model.predict(img_array)
    pred_class = np.argmax(pred[0])
    pred_label = class_names[pred_class]
    pred_label_kor = kor_label_map.get(pred_label, pred_label)
    confidence = pred[0][pred_class] * 100
    links = related_links[pred_label]

    result = f"🔬 예측 결과: {pred_label_kor} ({confidence:.2f}%)\n\n🔗 관련 기사:\n" + "\n".join(links)
    return result

demo = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="세포소기관 분류기🔬 | 엽록체, 미토콘드리아, 소포체, 골지체 분류",
    description="이미지를 업로드하면 어떤 세포소기관인지 분류하고 관련 기사를 보여줍니다."
)

demo.launch()