# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: py35-paddle1.2.0
# ---

#解压数据集
# !unzip -oq data/data129645/RockData.zip -d data/
#运行一次就够了
#运行一次就够了
#运行一次就够了

#查看数据集目录树
# !tree data/RockData/

# 引入依赖库
import cv2
import os 
import numpy as np
import matplotlib.pyplot as plt

# +
#图片抽样

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
# %matplotlib inline

#读取数据集中一个文件夹的路径
file_dir = '/home/aistudio/data/RockData/Coal/'
#设立列表存储图片名
filesum = []
#读取图片名
for root, dirs, files in os.walk(file_dir): 
    filesum.append(files)
filesum = filesum[0]
#打印图片名看看效果符不符合预期
#print(filesum)
#定义画布大小
plt.figure(figsize=(8, 8))
#循环读取图片并显示
for i in range(1,5):
    plt.subplot(2,2,i)
    plt.title(filesum[i])
    image = file_dir+filesum[i]
    print(f">>>{image}")
    plt.imshow(cv2.imread(image,1))
plt.tight_layout()
plt.show()
# -

#使用pip安装方式安装2.1.0版本：
# !pip install paddlex==2.1.0 -i https://mirror.baidu.com/pypi/simple

#在模型进行训练时，我们需要划分训练集，验证集和测试集
#因此需要对如上数据进行划分，直接使用paddlex命令即可将数据集随机划分成70%训练集，20%验证集和10%测试集
#划分好的数据集会额外生成labels.txt, train_list.txt, val_list.txt, test_list.txt四个文件，之后可直接进行训练。
# !paddlex --split_dataset --format ImageNet --dataset_dir /home/aistudio/data/RockData --val_value 0.2 --test_value 0.1

# +
#在训练和验证过程中，数据的处理
from paddlex import transforms as T
train_transforms = T.Compose([
    T.RandomCrop(crop_size=224),
    T.Normalize()])

eval_transforms = T.Compose([
    T.ResizeByShort(short_size=256),
    T.CenterCrop(crop_size=224),
    T.Normalize()
])
# -

import paddlex as pdx
#定义数据集，pdx.datasets.ImageNet表示读取ImageNet格式的分类数据集：
train_dataset = pdx.datasets.ImageNet(
    data_dir='/home/aistudio/data/RockData',
    file_list='/home/aistudio/data/RockData/train_list.txt',
    label_list='/home/aistudio/data/RockData/labels.txt',
    transforms=train_transforms,
    shuffle=True)
eval_dataset = pdx.datasets.ImageNet(
    data_dir='/home/aistudio/data/RockData',
    file_list='/home/aistudio/data/RockData/val_list.txt',
    label_list='/home/aistudio/data/RockData/labels.txt',
    transforms=eval_transforms)

# +
#使用百度基于蒸馏方法得到的MobileNetV3预训练模型，模型结构与MobileNetV3一致，但精度更高。


num_classes = len(train_dataset.labels)
model = pdx.cls.MobileNetV3_small(num_classes=num_classes)

model.train(num_epochs=10,
            train_dataset=train_dataset,
            train_batch_size=32,
            eval_dataset=eval_dataset,
            lr_decay_epochs=[4, 6, 8],
            save_dir='output/mobilenetv3_small',
            use_vdl=True)
#use_vdl=True表示可以启动visualdl并查看可视化的指标变化情况。
# -

#模型在训练过程中，会每间隔一定轮数保存一次模型
#在验证集上评估效果最好的一轮会保存在save_dir目录下的best_model文件夹
#加载模型，进行预测：
import paddlex as pdx
model = pdx.load_model('output/mobilenetv3_small/best_model')
result = model.predict('/home/aistudio/data/RockData/Coal/Coal0.jpg')
plt.imshow(cv2.imread('/home/aistudio/data/RockData/Coal/Coal0.jpg',1))
print("Predict Result: ", result)

StarcoderdataJupyterScriptsDedupFiltered