# -*- coding: utf-8 -*-
'''
训练摔倒检测模型
'''

# import the necessary packages
from oldcare.datasets import SimpleDatasetLoader
from oldcare.preprocessing import AspectAwarePreprocessor_fall
from oldcare.preprocessing import AspectAwarePreprocessor
from oldcare.preprocessing import ImageToArrayPreprocessor
from oldcare.conv import vgg16
from oldcare.callbacks import TrainingMonitor
from imutils import paths
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
from keras.utils import to_categorical
from sklearn.preprocessing import LabelEncoder
from keras.preprocessing.image import ImageDataGenerator
import tensorflow as tf
import cv2
import numpy as np
from skimage import img_as_ubyte

# 全局变量
dataset_path = 'dataset'
output_model_path = 'models/fall_detection.hdf5'
output_plot_path = 'plots/fall_detection.png'

# 全局常量
TARGET_WIDTH = 64
TARGET_HEIGHT = 64
BATCH_SIZE = 128
EPOCHS = 50
LR_INIT = 0.001
DECAY = LR_INIT / EPOCHS
MOMENTUM = 0.9

# 加载图片
aap = AspectAwarePreprocessor(TARGET_WIDTH, TARGET_HEIGHT)
iap = ImageToArrayPreprocessor()
print(type(aap))

#image_opencv = np.array(aap)
#image_opencv = np.uint8(image_opencv)
#image_opencv = img_as_ubyte(image_opencv)



print("[INFO] loading images...")
imagePaths = list(paths.list_images(dataset_path))
sdl = SimpleDatasetLoader(preprocessors=[aap, iap])
#sdl = SimpleDatasetLoader(preprocessors=[aap, fgmask])
(data, labels) = sdl.load(imagePaths, 500, False)
data = data.astype("float") / 255.0

#cv2.imshow("data", data)

#fg = cv2.createBackgroundSubtractorMOG2()
#fgmask = fg.apply(data)
#cv2.imshow("fgmask", fgmask)

# convert the labels from integers to vectors
le = LabelEncoder().fit(labels)
labels = to_categorical(le.transform(labels), 2)

# partition the data into training and testing splits using 80% of
# the data for training and the remaining 20% for testing
(trainX, testX, trainY, testY) = train_test_split(data,labels, test_size=0.20, stratify=labels, random_state=42)

# construct the image generator for data augmentation
aug = ImageDataGenerator(rotation_range=30, width_shift_range=0.1,
                         height_shift_range=0.1, shear_range=0.2,
                         zoom_range=0.2, horizontal_flip=True,
                         fill_mode="nearest")

# initialize the model
print("[INFO] compiling model...")
model = vgg16.build(width=TARGET_WIDTH, height=TARGET_HEIGHT,depth=3,classes=2)
model.build(input_shape=(32,32,3,2))
model.summary()

opt = tf.keras.optimizers.Adamax(learning_rate=LR_INIT, beta_1=0.9, beta_2=0.999, epsilon=1e-07, name='Adamax')
#opt = tf.keras.optimizers.Adagrad(learning_rate=LR_INIT, initial_accumulator_value=0.1, epsilon=1e-07,name='Adagrad')
model.compile(loss="categorical_crossentropy", optimizer=opt,metrics=["accuracy"])

# construct the set of callbacks
callbacks = [TrainingMonitor(output_plot_path)]

# train the network
print("[INFO] training network...")
H = model.fit(trainX, trainY, validation_data=(testX, testY), batch_size=BATCH_SIZE, epochs=EPOCHS, callbacks=callbacks, verbose=1)

# evaluate the network
print("[INFO] evaluating network...")
predictions = model.predict(testX, batch_size=BATCH_SIZE)
print(classification_report(testY.argmax(axis=1),
                            predictions.argmax(axis=1), target_names=le.classes_))

# save the model to disk
print("[INFO] serializing network...")
model.save(output_model_path)
