app.py

import gradio as gr
import numpy as np
from PIL import Image
import os

# ------------------------------
# 1) 모델 파일 경로 지정
# ------------------------------
MODEL_PATH = "crack_detection.h5"

IS_TF = MODEL_PATH.endswith(".h5") or MODEL_PATH.endswith(".keras")

# ------------------------------
# 2) TensorFlow 모델 로드
# ------------------------------
if IS_TF:
    import tensorflow as tf
    model = tf.keras.models.load_model(MODEL_PATH)
    print("🔥 Loaded TensorFlow crack classifier")

# ------------------------------
# 3) PyTorch 모델 로드
# ------------------------------
else:
    import torch
    from torch import nn
    
    class CNN(nn.Module):
        # 네가 가진 모델 구조 맞게 조정 필요
        def __init__(self):
            super().__init__()
            self.net = nn.Sequential(
                nn.Conv2d(3, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
                nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
                nn.Flatten(),
                nn.Linear(32 * 56 * 56, 2)  # 이미지 크기 맞게 조정 필요
            )
        def forward(self, x):
            return self.net(x)

    model = CNN()
    model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu"))
    model.eval()
    print("🔥 Loaded PyTorch crack classifier")


# ------------------------------
# 4) 전체 예측 함수
# ------------------------------
def predict(img: Image.Image):
    # 입력 변환
    img_resized = img.resize((224, 224))
    arr = np.array(img_resized) / 255.0

    if IS_TF:
        # TensorFlow
        X = arr.reshape(1, 224, 224, 3)
        probs = model.predict(X)[0]  # [p_normal, p_crack]

    else:
        # PyTorch
        import torch
        X = (
            torch.tensor(arr)
            .permute(2, 0, 1)
            .unsqueeze(0)
            .float()
        )
        probs = torch.softmax(model(X), dim=1).detach().numpy()[0]

    p_normal = float(probs[0])
    p_crack  = float(probs[1])

    if p_crack > p_normal:
        label = "crack"
        conf  = p_crack
    else:
        label = "normal"
        conf  = p_normal

    # ------------------------------
    # 프론트가 요구하는 JSON 구조
    # ------------------------------
    return {
        "data": [
            {
                "label": label,
                "confidence": conf
            }
        ]
    }


# ------------------------------
# 5) Gradio API Interface
# ------------------------------
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.JSON(label="Detection Result"),
    title="Crack Detection Classifier",
    description="사진을 업로드하면 균열/정상 여부와 확률(%)을 반환합니다.",
    flagging_mode="never"
)

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