app.py

import os, json
import numpy as np
import streamlit as st
import tensorflow as tf
from PIL import Image

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ART_DIR = os.path.join(BASE_DIR, "artifacts")
MODEL_PATH = os.path.join(ART_DIR, "model.keras")
CLASSES_PATH = os.path.join(ART_DIR, "class_names.json")

IMG_SIZE = (200, 200)
THRESHOLD = 0.5

@st.cache_resource
def load_artifacts():
    model = tf.keras.models.load_model(MODEL_PATH)
    with open(CLASSES_PATH, "r", encoding="utf-8") as f:
        class_names = json.load(f)
    return model, class_names

st.title("NG / OK Image Classifier")

if not os.path.exists(MODEL_PATH) or not os.path.exists(CLASSES_PATH):
    st.error("먼저 train.py를 실행해서 artifacts/model.keras 와 class_names.json 을 생성하세요.")
    st.stop()

model, class_names = load_artifacts()

uploaded = st.file_uploader("이미지 업로드", type=["png", "jpg", "jpeg", "webp"])

if uploaded is not None:
    img = Image.open(uploaded).convert("RGB")
    st.image(img, caption="Uploaded image", use_container_width=True)

    img_resized = img.resize(IMG_SIZE)
    x = np.array(img_resized, dtype=np.float32) / 255.0
    x = np.expand_dims(x, axis=0)  # (1,200,200,3)

    prob = float(model.predict(x, verbose=0)[0][0])
    pred_idx = 1 if prob >= THRESHOLD else 0

    # binary sigmoid 출력이라 pred_idx(0/1)에 따라 class_names 매핑
    pred_label = class_names[pred_idx] if pred_idx < len(class_names) else str(pred_idx)

    st.subheader("Prediction")
    st.write(f"**Label:** {pred_label}")
    st.write(f"**OK probability (sigmoid):** {prob:.4f}")
    st.write(f"**Threshold:** {THRESHOLD}")