import tensorflow as tf
import matplotlib .pyplot as plt
from PIL import Image
import numpy as np
import os
from myResNet import resnet35


#   设置相关底层配置
physical_devices = tf.config.experimental.list_physical_devices(‘GPU‘)
assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
tf.config.experimental.set_memory_growth(physical_devices[0], True)

size = 100
mybatch = 30
# l_rate = 5e-7
load_name = 1e-9
l_rate = 1e-9
ep = 500
img_shape = (1,size,size,3)
each_shape = (size,size)
model_shape = (None,size, size, 3)
def getdata():
    img_path = ‘D:\深度学习\my_imgtest\myimg‘
    x_shape = img_shape
    x_np = np.zeros(shape=x_shape)
    y_shape = (1,1)
    y_np = np.zeros(shape=y_shape)
    y_flag = [‘cui‘,‘ma‘,‘mei‘,‘po‘,‘yi‘]
    for each_person in os.listdir(img_path):
        person_imgs_path = os.path.join(img_path,each_person)
        # print(person_imgs_path)
        #   确定y标签
        flag = person_imgs_path.split(‘\\‘)[-1]
        flag = flag.split(‘_‘)[0]
        # print(flag)
        flag_index = y_flag.index(flag)+1
        # print(flag_index)
        for _,_,imgs_name in os.walk(person_imgs_path):
            for each_img_name in imgs_name:
                each_img = os.path.join(person_imgs_path,each_img_name)
                img1 = Image.open(each_img)
                img1 = img1.resize(each_shape)
                img_np = np.array(img1)
                # print(img_np)
                # print(img_np.shape)
                if np.max(x_np[0]) == 0:
                    # print(‘第一个是0‘)
                    x_np[0] = img_np
                    # print(‘[flag_index]‘, [flag_index])
                    # print(‘y_np:‘, y_np)
                    y_np[0] = [flag_index]
                    continue
                x_np = np.insert(x_np,0,img_np,axis=0)
                y_np = np.insert(y_np,0,flag_index,axis=0)
                # print(‘x_np:‘, x_np)
        # break
    print(‘读取图片成功！！‘)
    # print(‘x_np.shape:‘,x_np.shape)
    # print(‘y_np.shape‘,y_np.shape)
    # x_np.dtype = ‘int32‘

    return x_np,y_np

def preprocess(x,y):
    x = tf.cast(x,dtype=tf.float32) / 255
    y = tf.cast(y,dtype=tf.int32)
    return x,y

if __name__ == ‘__main__‘:
    x,y = getdata()
    # #   画图
    # for i in range(x.shape[0]):
    #     if i != 14 :
    #         continue
    #     print(x[i].shape)
    #     plt.imshow(x[i].astype(‘int32‘))
    #     y_flag = [‘cui‘, ‘ma‘, ‘mei‘, ‘po‘, ‘yi‘]
    #     print(‘y[i]:‘,y[i][0])
    #     name = y_flag[int(y[i][0])-1]
    #     plt.title(name)
    #     plt.show()

    #   识别训练
    y = tf.squeeze(y, axis=1)
    print(x.shape,y.shape)
    train_db = tf.data.Dataset.from_tensor_slices((x, y))
    train_db = train_db.map(preprocess).shuffle(30).batch(mybatch)
    #   打印看下数据的形状
    sample = next(iter(train_db))
    print(‘sample:‘, sample[0].shape, sample[1].shape
          , tf.reduce_min(sample[0]), tf.reduce_max(sample[0]))

    model = resnet35()
    model.build(input_shape=model_shape)
    #   定义优化器
    optimizer = tf.optimizers.Adam(lr=l_rate)

    #   加载参数
    load_path = os.path.join(‘D:\深度学习\my_imgtest\model_save‘,str(load_name))
    model.load_weights(load_path)
    #   训练
    for epoch in range(ep):
        for step, (x, y) in enumerate(train_db):
            with tf.GradientTape() as tape:
                # print(‘x.shape:‘,x.shape)
                logits = model(x)
                #
                y_onehot = tf.one_hot(y, depth=5)
                loss = tf.losses.categorical_crossentropy(y_onehot, logits, from_logits=True)
                loss = tf.reduce_mean(loss)
            grads = tape.gradient(loss, model.trainable_variables)
            if loss > 10:
                print(‘grads:‘,grads)
                break
                break
            optimizer.apply_gradients(zip(grads, model.trainable_variables))
            if step % 100 == 0:
                print(epoch, step, ‘loss:‘, float(loss))
        for x, y in train_db:
            # print()
            logits = model(x)
            prob = tf.nn.softmax(logits, axis=1)
            pred = tf.argmax(prob, axis=1)
            pred = tf.cast(pred, tf.int32)
            correct = tf.cast(tf.equal(pred, y), dtype=tf.int32)
            correct = tf.reduce_mean(tf.cast(correct, dtype=tf.float32))
        print(‘acc:‘, float(correct))

    save_path = ‘D:\深度学习\my_imgtest\model_save‘
    save_path = os.path.join(save_path,str(l_rate))
    model.save_weights(save_path)