train.py

import os, json
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.optimizers import RMSprop

print("TF:", tf.__version__)

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.join(BASE_DIR, "sample")
OUT_DIR = os.path.join(BASE_DIR, "artifacts")
os.makedirs(OUT_DIR, exist_ok=True)

MODEL_PATH = os.path.join(OUT_DIR, "model.keras")
CLASSES_PATH = os.path.join(OUT_DIR, "class_names.json")

IMG_SIZE = (200, 200)
BATCH_SIZE = 2
SEED = 42

# ✅ 데이터가 너무 적으니 일단 train에 전부 사용 (validation 없음)
train_gen = ImageDataGenerator(rescale=1.0 / 255.0)

train_ds = train_gen.flow_from_directory(
    DATA_DIR,
    target_size=IMG_SIZE,
    batch_size=BATCH_SIZE,
    class_mode="binary",
    shuffle=True,
    seed=SEED
)

# 클래스명 저장 (Streamlit에서 동일 매핑을 쓰기 위해 필수)
# 예: {'NG':0, 'OK':1} -> index 순서에 맞게 ["NG","OK"]
class_indices = train_ds.class_indices
class_names = [None] * len(class_indices)
for name, idx in class_indices.items():
    class_names[idx] = name

print("class_indices:", class_indices)
print("class_names:", class_names)

with open(CLASSES_PATH, "w", encoding="utf-8") as f:
    json.dump(class_names, f, ensure_ascii=False, indent=2)

# ✅ 모델
model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(IMG_SIZE[0], IMG_SIZE[1], 3)),
    tf.keras.layers.Conv2D(16, (3,3), activation="relu"),
    tf.keras.layers.MaxPool2D(2,2),
    tf.keras.layers.Conv2D(32, (3,3), activation="relu"),
    tf.keras.layers.MaxPool2D(2,2),
    tf.keras.layers.Conv2D(64, (3,3), activation="relu"),
    tf.keras.layers.MaxPool2D(2,2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation="relu"),
    tf.keras.layers.Dropout(0.3),
    tf.keras.layers.Dense(1, activation="sigmoid")
])

model.compile(
    loss="binary_crossentropy",
    optimizer=RMSprop(learning_rate=0.001),
    metrics=["accuracy"]
)

history = model.fit(train_ds, epochs=30)

# ✅ 저장
model.save(MODEL_PATH)
print("\nSaved model to:", MODEL_PATH)
print("Saved classes to:", CLASSES_PATH)