# 环境变量配置，用于控制是否使用GPU
# 说明文档：https://paddlex.readthedocs.io/zh_CN/develop/appendix/parameters.html#gpu
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

from paddlex.det import transforms
import paddlex as pdx

# 下载和解压铝材缺陷检测数据集
aluminum_dataset = 'https://bj.bcebos.com/paddlex/examples/industrial_quality_inspection/datasets/aluminum_inspection.tar.gz'
pdx.utils.download_and_decompress(aluminum_dataset, path='./')

# 定义训练和验证时的transforms
# API说明 https://paddlex.readthedocs.io/zh_CN/develop/apis/transforms/det_transforms.html
train_transforms = transforms.Compose([
    transforms.MixupImage(mixup_epoch=250), transforms.RandomDistort(),
    transforms.RandomExpand(), transforms.RandomCrop(), transforms.Resize(
        target_size=608, interp='RANDOM'), transforms.RandomHorizontalFlip(),
    transforms.Normalize()
])

eval_transforms = transforms.Compose([
    transforms.Resize(
        target_size=608, interp='CUBIC'), transforms.Normalize()
])

# 定义训练和验证所用的数据集
# API说明：https://paddlex.readthedocs.io/zh_CN/develop/apis/datasets.html#paddlex-datasets-vocdetection
train_dataset = pdx.datasets.VOCDetection(
    data_dir='aluminum_inspection',
    file_list='aluminum_inspection/train_list.txt',
    label_list='aluminum_inspection/labels.txt',
    transforms=train_transforms,
    shuffle=True)
eval_dataset = pdx.datasets.VOCDetection(
    data_dir='aluminum_inspection',
    file_list='aluminum_inspection/val_list.txt',
    label_list='aluminum_inspection/labels.txt',
    transforms=eval_transforms)

# 初始化模型，并进行训练
# 可使用VisualDL查看训练指标，参考https://paddlex.readthedocs.io/zh_CN/develop/train/visualdl.html
num_classes = len(train_dataset.labels)

# API说明: https://paddlex.readthedocs.io/zh_CN/develop/apis/models/detection.html#paddlex-det-yolov3
model = pdx.det.YOLOv3(num_classes=num_classes, backbone='MobileNetV3_large')

# API说明: https://paddlex.readthedocs.io/zh_CN/develop/apis/models/detection.html#train
# 各参数介绍与调整说明：https://paddlex.readthedocs.io/zh_CN/develop/appendix/parameters.html
model.train(
    num_epochs=400,
    train_dataset=train_dataset,
    train_batch_size=8,
    eval_dataset=eval_dataset,
    warmup_steps=4000,
    learning_rate=0.000125,
    lr_decay_epochs=[240, 320],
    save_dir='output/yolov3_mobilenetv3',
    use_vdl=True)

[(11, 'paddlex.utils.download_and_decompress', 'pdx.utils.download_and_decompress', True, 'import paddlex as pdx\n'), (29, 'paddlex.datasets.VOCDetection', 'pdx.datasets.VOCDetection', True, 'import paddlex as pdx\n'), (35, 'paddlex.datasets.VOCDetection', 'pdx.datasets.VOCDetection', True, 'import paddlex as pdx\n'), (46, 'paddlex.det.YOLOv3', 'pdx.det.YOLOv3', True, 'import paddlex as pdx\n'), (16, 'paddlex.det.transforms.MixupImage', 'transforms.MixupImage', False, 'from paddlex.det import transforms\n'), (16, 'paddlex.det.transforms.RandomDistort', 'transforms.RandomDistort', False, 'from paddlex.det import transforms\n'), (17, 'paddlex.det.transforms.RandomExpand', 'transforms.RandomExpand', False, 'from paddlex.det import transforms\n'), (17, 'paddlex.det.transforms.RandomCrop', 'transforms.RandomCrop', False, 'from paddlex.det import transforms\n'), (17, 'paddlex.det.transforms.Resize', 'transforms.Resize', False, 'from paddlex.det import transforms\n'), (18, 'paddlex.det.transforms.RandomHorizontalFlip', 'transforms.RandomHorizontalFlip', False, 'from paddlex.det import transforms\n'), (19, 'paddlex.det.transforms.Normalize', 'transforms.Normalize', False, 'from paddlex.det import transforms\n'), (23, 'paddlex.det.transforms.Resize', 'transforms.Resize', False, 'from paddlex.det import transforms\n'), (24, 'paddlex.det.transforms.Normalize', 'transforms.Normalize', False, 'from paddlex.det import transforms\n')]