基于PaddlePaddle的宝石分类

任务描述：

本次实践是一个多分类任务，需要将照片中的宝石分别进行识别，完成宝石的识别

实践平台：百度AI实训平台-AI Studio、PaddlePaddle2.0.0 动态图

要求：使用CNN方法实现宝石识别！

卷积神经网络

卷积神经网络是提取图像特征的经典网络，其结构一般包含多个卷积层与池化层的交替组合。

数据集介绍

• 数据集文件名为archive_train.zip,archive_test.zip。

• 该数据集包含25个类别不同宝石的图像。

• 这些类别已经分为训练和测试数据。

• 图像大小不一，格式为.jpeg。

# 查看当前挂载的数据集目录, 该目录下的变更重启环境后会自动还原
# View dataset directory. This directory will be recovered automatically after resetting environment. 
!ls /home/aistudio/data
复制代码

data55032  dataset
复制代码

#导入需要的包
import os
import zipfile
import random
import json
import cv2
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import paddle
from paddle.io import Dataset
import paddle.nn as nn 

复制代码

1、数据准备

'''
参数配置
'''
train_parameters = {
    "input_size": [3, 224, 224],                     # 输入图片的shape
    "class_dim": 25,                                 # 分类数
    "src_path":"data/data55032/archive_train.zip",   # 原始数据集路径
    "target_path":"/home/aistudio/data/dataset",     # 要解压的路径 
    "train_list_path": "./train.txt",                # train_data.txt路径
    "eval_list_path": "./eval.txt",                  # eval_data.txt路径
    "label_dict":{},                                 # 标签字典
    "readme_path": "/home/aistudio/data/readme.json",# readme.json路径
    "num_epochs":10,                                 # 训练轮数
    "train_batch_size": 32,                          # 批次的大小
    "learning_strategy": {                           # 优化函数相关的配置
        "lr": 0.0005                                  # 超参数学习率
    } 
}


复制代码

def unzip_data(src_path,target_path):
    '''
    解压原始数据集，将src_path路径下的zip包解压至data/dataset目录下
    ''' 
    if(not os.path.isdir(target_path)):    
        z = zipfile.ZipFile(src_path, 'r')
        z.extractall(path=target_path)
        z.close()
    else:
        print("文件已解压")
复制代码

def get_data_list(target_path,train_list_path,eval_list_path):
    '''
    生成数据列表
    '''
    # 获取所有类别保存的文件夹名称
    data_list_path=target_path
    class_dirs = os.listdir(data_list_path) 
    if '__MACOSX' in class_dirs:
        class_dirs.remove('__MACOSX')
    # 存储要写进eval.txt和train.txt中的内容
    trainer_list=[]
    eval_list=[]
    class_label=0
    i = 0
    
    for class_dir in class_dirs:   
        path = os.path.join(data_list_path,class_dir)
        # 获取所有图片
        img_paths = os.listdir(path)
        for img_path in img_paths:                                        # 遍历文件夹下的每个图片
            i += 1
            name_path = os.path.join(path,img_path)                       # 每张图片的路径
            if i % 10 == 0:                                                
                eval_list.append(name_path + "\t%d" % class_label + "\n")
            else: 
                trainer_list.append(name_path + "\t%d" % class_label + "\n") 
        
        train_parameters['label_dict'][str(class_label)]=class_dir
        class_label += 1
            
    #乱序  
    random.shuffle(eval_list)
    with open(eval_list_path, 'a') as f:
        for eval_image in eval_list:
            f.write(eval_image) 
    #乱序        
    random.shuffle(trainer_list) 
    with open(train_list_path, 'a') as f2:
        for train_image in trainer_list:
            f2.write(train_image) 
 
    print ('生成数据列表完成！')
复制代码

# 参数初始化
src_path=train_parameters['src_path']
target_path=train_parameters['target_path']
train_list_path=train_parameters['train_list_path']
eval_list_path=train_parameters['eval_list_path']
batch_size=train_parameters['train_batch_size']

# 解压原始数据到指定路径
unzip_data(src_path,target_path)

#每次生成数据列表前，首先清空train.txt和eval.txt
with open(train_list_path, 'w') as f: 
    f.seek(0)
    f.truncate() 
with open(eval_list_path, 'w') as f: 
    f.seek(0)
    f.truncate() 
    
#生成数据列表   
get_data_list(target_path,train_list_path,eval_list_path)
复制代码

文件已解压
生成数据列表完成！
复制代码

class Reader(Dataset):
    def __init__(self, data_path, mode='train'):
        """
        数据读取器
        :param data_path: 数据集所在路径
        :param mode: train or eval
        """
        super().__init__()
        self.data_path = data_path
        self.img_paths = []
        self.labels = []

        if mode == 'train':
            with open(os.path.join(self.data_path, "train.txt"), "r", encoding="utf-8") as f:
                self.info = f.readlines()
            for img_info in self.info:
                img_path, label = img_info.strip().split('\t')
                self.img_paths.append(img_path)
                self.labels.append(int(label))

        else:
            with open(os.path.join(self.data_path, "eval.txt"), "r", encoding="utf-8") as f:
                self.info = f.readlines()
            for img_info in self.info:
                img_path, label = img_info.strip().split('\t')
                self.img_paths.append(img_path)
                self.labels.append(int(label))


    def __getitem__(self, index):
        """
        获取一组数据
        :param index: 文件索引号
        :return:
        """
        # 第一步打开图像文件并获取label值
        img_path = self.img_paths[index]
        img = Image.open(img_path)
        if img.mode != 'RGB':
            img = img.convert('RGB') 
        img = img.resize((224, 224), Image.BILINEAR)
        img = np.array(img).astype('float32')
        img = img.transpose((2, 0, 1)) / 255
        label = self.labels[index]
        label = np.array([label], dtype="int64")
        return img, label

    def print_sample(self, index: int = 0):
        print("文件名", self.img_paths[index], "\t标签值", self.labels[index])

    def __len__(self):
        return len(self.img_paths)
复制代码

#训练数据加载
train_dataset = Reader('/home/aistudio/',mode='train')
train_loader = paddle.io.DataLoader(train_dataset, batch_size=16, shuffle=True)
#测试数据加载
eval_dataset = Reader('/home/aistudio/',mode='eval')
eval_loader = paddle.io.DataLoader(eval_dataset, batch_size = 8, shuffle=False)
复制代码

train_dataset.print_sample(200)
print(train_dataset.__len__())
eval_dataset.print_sample(0)
print(eval_dataset.__len__())
print(eval_dataset.__getitem__(10)[0].shape)
print(eval_dataset.__getitem__(10)[1].shape)
复制代码

文件名 /home/aistudio/data/dataset/Fluorite/fluorite_35.jpg 	标签值 14
730
文件名 /home/aistudio/data/dataset/Danburite/danburite_15.jpg 	标签值 18
81
(3, 224, 224)
(1,)
复制代码

Batch=0
Batchs=[]
all_train_accs=[]
def draw_train_acc(Batchs, train_accs):
    title="training accs"
    plt.title(title, fontsize=24)
    plt.xlabel("batch", fontsize=14)
    plt.ylabel("acc", fontsize=14)
    plt.plot(Batchs, train_accs, color='green', label='training accs')
    plt.legend()
    plt.grid()
    plt.show()

all_train_loss=[]
def draw_train_loss(Batchs, train_loss):
    title="training loss"
    plt.title(title, fontsize=24)
    plt.xlabel("batch", fontsize=14)
    plt.ylabel("loss", fontsize=14)
    plt.plot(Batchs, train_loss, color='red', label='training loss')
    plt.legend()
    plt.grid()
    plt.show()
复制代码

2、定义模型

# 定义卷积神经网络实现宝石识别
class MyCNN(nn.Layer): 
    def __init__(self):
        super(MyCNN,self).__init__()
        self.conv0=nn.Conv2D(in_channels=3, out_channels=64, kernel_size=3, stride=1)
        self.pool0=nn.MaxPool2D(kernel_size=2, stride=2)
        self.conv1=nn.Conv2D(in_channels=64, out_channels=128, kernel_size=4, stride=1)
        self.pool1=nn.MaxPool2D(kernel_size=2, stride=2)
        self.conv2=nn.Conv2D(in_channels=128, out_channels=50, kernel_size=5)
        self.pool2=nn.MaxPool2D(kernel_size=2, stride=2)
        self.conv3=nn.Conv2D(in_channels=50, out_channels=50, kernel_size=5)
        self.pool3=nn.MaxPool2D(kernel_size=2, stride=2)
        self.conv4=nn.Conv2D(in_channels=50, out_channels=50, kernel_size=5)
        self.pool4=nn.MaxPool2D(kernel_size=2, stride=2)
        self.fc1=nn.Linear(in_features=50*3*3, out_features=25)

    def forward(self,input): 
        print("input.shape:",input.shape)
        x=self.conv0(input)
        print("x.shape:",x.shape)
        x=self.pool0(x)        
        print('x0.shape:',x.shape)
        x=self.conv1(x)
        print(x.shape)
        x=self.pool1(x)
        print('x1.shape:',x.shape)
        x=self.conv2(x)
        print(x.shape)
        x=self.pool2(x)
        print('x2.shape:',x.shape)
        x=self.conv3(x)
        print(x.shape)
        x=self.pool3(x)
        print('x3.shape:',x.shape)
        x=self.conv4(x)
        print(x.shape)
        x=self.pool4(x)
        print('x4.shape:',x.shape)
        x=paddle.reshape(x, shape=[-1, 50*3*3])
        print('x3.shape:',x.shape)
        y=self.fc1(x)
        print('y.shape:', y.shape)
        # input.shape: [16, 3, 224, 224]
        # x.shape: [16, 64, 222, 222]
        # x0.shape: [16, 64, 111, 111]
        # [16, 128, 108, 108]
        # x1.shape: [16, 128, 54, 54]
        # [16, 50, 50, 50]
        # x2.shape: [16, 50, 25, 25]
        # [16, 50, 21, 21]
        # x3.shape: [16, 50, 10, 10]
        # [16, 50, 6, 6]
        # x4.shape: [16, 50, 3, 3]
        return y       
        
复制代码

# 定义卷积神经网络实现宝石识别
class MyCNN2(nn.Layer): 
    def __init__(self):
        super(MyCNN2,self).__init__()
        self.conv0 = nn.Conv2D(in_channels= 3,out_channels=64, kernel_size=3,stride=1)
        self.pool0 = nn.MaxPool2D(kernel_size=2,stride=2)
        self.conv1 = nn.Conv2D(in_channels = 64,out_channels=128,kernel_size=4,stride = 1)
        self.pool1 = nn.MaxPool2D(kernel_size=2,stride=2)
        self.conv2 = nn.Conv2D(in_channels= 128,out_channels=50,kernel_size=5)
        self.pool2 = nn.MaxPool2D(kernel_size=2,stride=2)
        self.fc1 = nn.Linear(in_features=50*25*25,out_features=25)


    def forward(self,input): 
        x = self.conv0(input)
        x = self.pool0(x)
        print("x:", x.shape)
        x = self.conv1(x)
        x = self.pool1(x)
        print("x:", x.shape)
        x = self.conv2(x)
        x = self.pool2(x)
        print("x:", x.shape)
        x = paddle.reshape(x,shape=[-1,50*25*25])
        y = self.fc1(x)
        
        return y
复制代码

3、训练模型-CNN

model=MyCNN() # 模型实例化
model.train() # 训练模式
cross_entropy = paddle.nn.CrossEntropyLoss()
opt=paddle.optimizer.SGD(learning_rate=train_parameters['learning_strategy']['lr'],\
                                                    parameters=model.parameters())

epochs_num=train_parameters['num_epochs'] #迭代次数
for pass_num in range(train_parameters['num_epochs']):
    for batch_id,data in enumerate(train_loader()):
        image = data[0]
        label = data[1]
        predict=model(image) #数据传入model
        loss=cross_entropy(predict,label)
        acc=paddle.metric.accuracy(predict,label)#计算精度
        if batch_id!=0 and batch_id%5==0:
            Batch = Batch+5 
            Batchs.append(Batch)
            all_train_loss.append(loss.numpy()[0])
            all_train_accs.append(acc.numpy()[0]) 
            print("epoch:{},step:{},train_loss:{},train_acc:{}".format(pass_num,batch_id,loss.numpy(),acc.numpy()))        
        loss.backward()       
        opt.step()
        opt.clear_grad()   #opt.clear_grad()来重置梯度
paddle.save(model.state_dict(),'MyCNN')#保存模型
draw_train_acc(Batchs,all_train_accs)
draw_train_loss(Batchs,all_train_loss)
复制代码

y.shape: [16, 25]


---------------------------------------------------------------------------

KeyboardInterrupt                         Traceback (most recent call last)

/tmp/ipykernel_3036/2664449083.py in <module>
      7 epochs_num=train_parameters['num_epochs'] #迭代次数
      8 for pass_num in range(train_parameters['num_epochs']):
----> 9     for batch_id,data in enumerate(train_loader()):
     10         image = data[0]
     11         label = data[1]


/opt/conda/envs/python35-paddle120-env/lib/python3.7/site-packages/paddle/fluid/dataloader/dataloader_iter.py in __next__(self)
    349         try:
    350             if in_dygraph_mode():
--> 351                 return self._reader.read_next_var_list()
    352             else:
    353                 if self._return_list:


KeyboardInterrupt: 
复制代码

4、模型评估-CNN

#模型评估
para_state_dict = paddle.load("MyCNN") 
model = MyCNN()
model.set_state_dict(para_state_dict) #加载模型参数
model.eval() #验证模式

accs = []

for batch_id,data in enumerate(eval_loader()):#测试集
    image=data[0]
    label=data[1]     
    predict=model(image)       
    acc=paddle.metric.accuracy(predict,label)
    accs.append(acc.numpy()[0])
avg_acc = np.mean(accs)
print("当前模型在验证集上的准确率为:",avg_acc)
复制代码

5、模型预测-CNN

def unzip_infer_data(src_path,target_path):
    '''
    解压预测数据集
    '''
    if(not os.path.isdir(target_path)):     
        z = zipfile.ZipFile(src_path, 'r')
        z.extractall(path=target_path)
        z.close()


def load_image(img_path):
    '''
    预测图片预处理
    '''
    img = Image.open(img_path) 
    if img.mode != 'RGB': 
        img = img.convert('RGB') 
    img = img.resize((224, 224), Image.BILINEAR)
    img = np.array(img).astype('float32') 
    img = img.transpose((2, 0, 1))  # HWC to CHW 
    img = img/255                # 像素值归一化 
    return img


infer_src_path = '/home/aistudio/data/data55032/archive_test.zip'
infer_dst_path = '/home/aistudio/data/archive_test'
unzip_infer_data(infer_src_path,infer_dst_path)

para_state_dict = paddle.load("MyCNN")
model = MyCNN()
model.set_state_dict(para_state_dict) #加载模型参数
model.eval() #验证模式

#展示预测图片
infer_path='data/archive_test/alexandrite_3.jpg'
img = Image.open(infer_path)
plt.imshow(img)          #根据数组绘制图像
plt.show()               #显示图像
#对预测图片进行预处理
infer_imgs = []
infer_imgs.append(load_image(infer_path))
infer_imgs = np.array(infer_imgs)
label_dic = train_parameters['label_dict']
for i in range(len(infer_imgs)):
    data = infer_imgs[i]
    dy_x_data = np.array(data).astype('float32')
    dy_x_data=dy_x_data[np.newaxis,:, : ,:]
    img = paddle.to_tensor (dy_x_data)
    out = model(img)
    lab = np.argmax(out.numpy())  #argmax():返回最大数的索引
    print("第{}个样本,被预测为：{},真实标签为：{}".format(i+1,label_dic[str(lab)],infer_path.split('/')[-1].split("_")[0]))       
print("结束")

复制代码