app2.py

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

# -----------------------
# Config
# -----------------------
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.join(BASE_DIR, "sample")          # sample/NG, sample/OK
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.set_page_config(page_title="NG/OK Inspector", layout="wide")

# -----------------------
# Load model & classes
# -----------------------
@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)  # index order, e.g., ["NG","OK"]
    return model, class_names

def list_dataset_images(root_dir):
    """Return list of dicts: {true_label, path, filename}"""
    exts = (".png", ".jpg", ".jpeg", ".webp", ".bmp")
    items = []
    for label in sorted(os.listdir(root_dir)):
        label_dir = os.path.join(root_dir, label)
        if not os.path.isdir(label_dir):
            continue
        for fname in sorted(os.listdir(label_dir)):
            if fname.lower().endswith(exts):
                items.append({
                    "true_label": label,
                    "path": os.path.join(label_dir, fname),
                    "filename": fname
                })
    return items

def predict_one(model, pil_img):
    img = pil_img.convert("RGB").resize(IMG_SIZE)
    x = np.array(img, dtype=np.float32) / 255.0
    x = np.expand_dims(x, axis=0)
    prob = float(model.predict(x, verbose=0)[0][0])  # sigmoid
    pred_idx = 1 if prob >= THRESHOLD else 0
    return prob, pred_idx

def badge(label, color):
    st.markdown(
        f"""
        <div style="
            padding: 18px 22px;
            border-radius: 18px;
            background: {color};
            color: white;
            font-size: 40px;
            font-weight: 800;
            text-align: center;
            width: fit-content;
            ">
            {label}
        </div>
        """,
        unsafe_allow_html=True
    )

# -----------------------
# Guard
# -----------------------
if not (os.path.exists(MODEL_PATH) and os.path.exists(CLASSES_PATH)):
    st.error("먼저 train.py 실행해서 artifacts/model.keras, artifacts/class_names.json 생성하세요.")
    st.stop()

model, class_names = load_artifacts()
items = list_dataset_images(DATA_DIR)

if not items:
    st.error("sample/NG, sample/OK 폴더 안에 이미지가 없습니다.")
    st.stop()

# -----------------------
# Sidebar: image list + click
# -----------------------
st.sidebar.title("Dataset Images")

# 초기 선택
if "selected_path" not in st.session_state:
    st.session_state.selected_path = items[0]["path"]

# 오른쪽에 파일들 “쫙”
for it in items:
    colA, colB = st.sidebar.columns([1, 3], vertical_alignment="center")

    # 작은 썸네일
    with colA:
        try:
            thumb = Image.open(it["path"]).convert("RGB")
            colA.image(thumb, use_container_width=True)
        except:
            pass

    # 클릭 버튼
    with colB:
        btn_label = f"{it['true_label']}  |  {it['filename']}"
        if st.button(btn_label, key=it["path"]):
            st.session_state.selected_path = it["path"]

# -----------------------
# Main: show selected & result
# -----------------------
selected_path = st.session_state.selected_path
selected_item = next((x for x in items if x["path"] == selected_path), None)

left, right = st.columns([2, 1], vertical_alignment="top")

with left:
    st.header("Selected Image")
    img = Image.open(selected_path).convert("RGB")
    st.image(img, use_container_width=True)

with right:
    st.header("Result")

    prob, pred_idx = predict_one(model, img)
    pred_label = class_names[pred_idx] if pred_idx < len(class_names) else str(pred_idx)

    # 색 크게 표시
    if pred_label.upper() == "OK":
        badge("OK", "#16a34a")  # green
    else:
        badge("NG", "#dc2626")  # red

    true_label = selected_item["true_label"] if selected_item else "Unknown"

    st.markdown("---")
    st.write(f"**True label (folder):** {true_label}")
    st.write(f"**Pred label (model):** {pred_label}")
    st.write(f"**Sigmoid prob:** {prob:.4f}")
    st.write(f"**Threshold:** {THRESHOLD}")

    # 맞/틀 표시
    correct = (true_label == pred_label)
    st.write("**Match:**", "✅ Correct" if correct else "❌ Wrong")