import cv2
import numpy as np
import os
# 數據集路徑
data_path = 'waste-classification-data'
# 類別列表
categories = ['Cardboard', 'Glass', 'Metal', 'Paper', 'Plastic', 'Trash', 'Battery', 'Clothes', 'Organic', 'Shoes']
# 圖像大小
img_size = 224
# 數據預處理
def prepare_data():
    data = []
    for category in categories:
        path = os.path.join(data_path, category)
        label = categories.index(category)
        for img_name in os.listdir(path):
            img_path = os.path.join(path, img_name)
            img = cv2.imread(img_path)
            img = cv2.resize(img, (img_size, img_size))
            img = img.astype('float32') / 255.0
            data.append([img, label])
    return np.array(data)

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout, Flatten
from tensorflow.keras.applications.vgg16 import VGG16
# 模型構建
def build_model():
    # 加載VGG16模型
    base_model = VGG16(weights='imagenet', include_top=False, input_shape=(img_size, img_size, 3))
    # 凍結前15層
    for layer in base_model.layers[:15]:
        layer.trainable = False
    model = Sequential()
    model.add(base_model)
    model.add(Flatten())
    model.add(Dense(256, activation='relu'))
    model.add(Dropout(0.5))
    model.add(Dense(10, activation='softmax'))
    return model

from sklearn.model_selection import train_test_split
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.losses import categorical_crossentropy
from tensorflow.keras.callbacks import ModelCheckpoint
# 數據預處理
data = prepare_data()
# 數據拆分
X = data[:, 0]
y = data[:, 1]
y = np.eye(10)[y.astype('int')]
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# 模型構建
model = build_model()
# 模型編譯
model.compile(optimizer=Adam(lr=0.001), loss=categorical_crossentropy, metrics=['accuracy'])
# 模型訓練
checkpoint = ModelCheckpoint('model.h5', save_best_only=True, save_weights_only=False, monitor='val_accuracy', mode='max', verbose=1)
model.fit(X_train, y_train, batch_size=32, epochs=10, validation_data=(X_test, y_test), callbacks=[checkpoint])

# 模型評估
loss, accuracy = model.evaluate(X_test, y_test)
print('Test Loss: {:.4f}'.format(loss))
print('Test Accuracy: {:.4f}'.format(accuracy))