Tensorflow2加载图片数据集并训练

import os
os.environ['TF_CPP_MIN_LOG_LEVEL']='2'
import tensorflow as tf
import random
import pathlib
import numpy as np
def load_image_label(dirpath):
    data_path = pathlib.Path(dirpath)
    all_image_paths  = list(data_path.glob('*/*'))
    all_image_paths = [str(path) for path in all_image_paths]
    random.shuffle(all_image_paths)
    label_names = sorted(item.name for item in data_path.glob('*/') if item.is_dir())
    label_count = len(label_names)
    label_to_index = dict((name, index) for index, name in enumerate(label_names))
    labels = [label_to_index[pathlib.Path(path).parent.name] for path in all_image_paths]
    return all_image_paths, labels, label_count
path = './Train'
all_image_paths, labels, label_count = load_image_label(path)
all_iamge_labels = np.zeros((len(all_image_paths),label_count))
for num, each_label in enumerate(labels):
    all_iamge_labels[num, each_label] = 1
ds = tf.data.Dataset.from_tensor_slices((all_image_paths, all_iamge_labels))
def load_and_preprocess_from_path_label(path, label):
    image = tf.io.read_file(path)  # 读取图片
    image = tf.image.decode_jpeg(image, channels=3)
    image = tf.image.resize(image, [60, 60])  
    image /= 255.0  # 归一化到[0,1]范围
    return image, label
image_label_ds = ds.map(load_and_preprocess_from_path_label)
train_image = []
train_label = []
for image, label in zip(all_image_paths, all_iamge_labels):
    r_image, r_label = load_and_preprocess_from_path_label(image, label)
    train_image.append(r_image)
    train_label.append(r_label)
train_images = np.array(train_image)
train_labels = np.array(train_label)
print(train_images.shape)
print(train_labels.shape)

from tensorflow import keras
from tensorflow.keras import layers
model = keras.Sequential(
[
    layers.Conv2D(5,3,activation = 'relu',padding = 'same',input_shape=[60, 60, 3]),
    layers.Flatten(),
    layers.Dense(128, activation='relu'),
    layers.Dense(3, activation='softmax')
])
model.compile(optimizer='adam',
             loss='categorical_crossentropy',
             metrics=['accuracy'])
model.fit(train_images, train_labels, epochs=5,batch_size=3)

import cv2 as cv
test_img = cv.imread("test.jpg")
test_img = cv.resize(test_img,[60,60])
print(test_img.shape)
test_img = np.expand_dims(np.array(test_img),axis=0)
result = model.predict(test_img)
print(result)

参考文档

https://blog.csdn.net/qq_36758914/article/details/106883812
https://blog.csdn.net/Black_Friend/article/details/104529859