【Paddle打比赛】百度网盘AI大赛-图像处理挑战赛文档检测优化赛ONNX方案

# clone PaddleSeg的项目!git clone https://gitee.com/paddlepaddle/PaddleSeg# 安装依赖!pip install  -r PaddleSeg/requirements.txt

!pip install  -r ./work/requirements.txt

#解压数据集!mkdir dataset
!unzip -oq data/data140391/train_datasets_document_detection_0411.zip -d dataset/
!unzip -oq data/data140389/testA_datasets_document_detection.zip -d dataset/# 数据处理，将255变为1（PaddleSeg默认255为忽略数据）!python work/data_process.py# 数据处理，生成数据列表，tarin.txt,val.txt,test.txt,labels.txt!python PaddleSeg/tools/split_dataset_list.py dataset/train_datasets_document_detection_0411 images label --split 0.9 0.1 0 --format jpg png

# # 复制segformer_b3配置文件%cp ./work/segformer_b3_doc.yml ./PaddleSeg/configs/

# 开始训练，如果想自己训练权重，去掉下面的注释即可# !python PaddleSeg/train.py \#        --config PaddleSeg/configs/segformer_b3_doc.yml \#        --do_eval \#        --use_vdl \#        --save_interval 500 \#        --save_dir output_segformer_b3_train

# 如果你只想看看结果，运行这段代码!python PaddleSeg/predict.py \
       --config PaddleSeg/configs/segformer_b3_doc.yml \
       --model_path work/model.pdparams \
       --image_path dataset/testA_datasets_document_detection/images# 输出效果图在output/result/added_prediction下

# 如果你训练更多的epoch，把下面代码的注释去掉，用自己训练的权重预测# !python PaddleSeg/predict.py \#        --config PaddleSeg/configs/segformer_b3_doc.yml \#        --model_path output_segformer_b3_train/best_model/model.pdparams \#        --image_path dataset/testA_datasets_document_detection/images

!pip install paddleseg
!pip install paddle2onnx
!pip install pillow==7.1.2

# 划重点，这里是转换代码，PaddleSeg转换简单，太香了！！！！！！！！！！！from paddleseg.models import SegFormerfrom paddleseg.cvlibs import managerimport paddle

model = SegFormer(num_classes=2,
                  backbone=manager.BACKBONES['MixVisionTransformer_B3'](),
                  embedding_dim=768,
                  pretrained=None)# 加载预训练模型参数model.set_dict(paddle.load('work/model.pdparams'))# 将模型设置为评估状态model.eval()# 定义输入数据input_spec = paddle.static.InputSpec(shape=[None, 3, 512, 512], dtype='float32', name='image')# ONNX模型导出paddle.onnx.export(model, './work/SegFormerB3', input_spec=[input_spec], opset_version=12)

%cp ./work/SegFormerB3.onnx ./SegFormerB3.onnx
!python ./work/predict.py dataset/testA_datasets_document_detection/images results

# 输出结果在resulrs文件夹中，可视化一下import osimport numpy as npimport matplotlib.pyplot as plt
%matplotlib inlinefrom PIL import Imageimport globimport warnings
warnings.filterwarnings("ignore")


DATADIR = 'results' lst=glob.glob(DATADIR+"/*.**g")



file0 = lst[0]
file1 = lst[1]
file2 = lst[2]# 读取图片img0 = Image.open(file0)
img0 = np.array(img0)
img1 = Image.open(file1)
img1 = np.array(img1)
img2 = Image.open(file2)
img2 = np.array(img2)# 画出读取的图片plt.figure(figsize=(16, 8))

f = plt.subplot(131)
f.set_title('0', fontsize=20)
plt.imshow(img0)

f = plt.subplot(132)
f.set_title('1', fontsize=20)
plt.imshow(img1)

f = plt.subplot(133)
f.set_title('2', fontsize=20)
plt.imshow(img2)
plt.show()

# 压缩可提交文件%cp ./work/predict.py ./predict.py
! zip submit.zip SegFormerB3.onnx predict.py# 下载上传即可