飞桨常规赛：PALM眼底彩照视盘探测与分割 - 5月第三名方案

# ! git clone https://gitee.com/paddlepaddle/PaddleSeg.git! pip -q install pattaimport sys

sys.path.append('PaddleSeg')

! unzip -oq /home/aistudio/data/data85136/常规赛：PALM眼底彩照视盘探测与分割.zip! rm -rf __MACOSX
! mv 常规赛：PALM眼底彩照视盘探测与分割 dataset

import osimport randomdef create_list():
    file_path = 'dataset/Train/fundus_image'
    imgs_name = os.listdir(file_path)
    random.shuffle(imgs_name)    with open('dataset/Train/train.txt', 'w') as tf:        with open('dataset/Train/val.txt', 'w') as ef:            for idx, img_name in enumerate(imgs_name):
                img_path = os.path.join('fundus_image', img_name)
                lab_path = img_path.replace('fundus_image', 'Disc_Masks').replace('jpg', 'bmp')                if (idx % 10 + 1) < 9:
                    tf.write(img_path + ' ' + lab_path + '\n')                else:
                    ef.write(img_path + ' ' + lab_path + '\n')

create_list()

from paddleseg.datasets import Datasetimport paddleseg.transforms as T# 构建训练集train_transforms = [
    T.Resize(target_size=(1120, 1120)),  # 修改大小
    T.Normalize()  # 归一化]
train_dataset = Dataset(
    transforms=train_transforms,
    dataset_root='dataset/Train',
    num_classes=2,
    mode='train',
    train_path='dataset/Train/train.txt',
    separator=' ',
)# 构建验证集val_transforms = [
    T.Resize(target_size=(1120, 1120)),
    T.Normalize()
]
val_dataset = Dataset(
    transforms=val_transforms,
    dataset_root='dataset/Train',
    num_classes=2,
    mode='val',
    val_path='dataset/Train/val.txt',
    separator=' ',
)

for img, lab in val_dataset:    print(img.shape, lab)    break

import paddlefrom paddleseg.models import AttentionUNetfrom paddleseg.models.losses import BCELoss, DiceLoss, MixedLoss# 模型url = Nonemodel = AttentionUNet(num_classes=2, pretrained=url)# 训练参数epochs = 20batch_size = 2iters = epochs * len(train_dataset) // batch_size# 损失函数mix_losses = [BCELoss(), DiceLoss()]
mix_coef = [1, 1]
mixloss = MixedLoss(mix_losses, mix_coef)
losses = {}
losses['types'] = [mixloss]
losses['coef'] = [1]# 学习率及优化器base_lr = 3e-4lr = paddle.optimizer.lr.CosineAnnealingDecay(base_lr, T_max=(iters // 5))
optimizer = paddle.optimizer.Adam(lr, parameters=model.parameters(), weight_decay=paddle.regularizer.L2Decay(1e-7))

from paddleseg.core import train

train(
    model=model,
    train_dataset=train_dataset,
    val_dataset=val_dataset,
    optimizer=optimizer,
    save_dir='save_output',
    iters=iters,
    batch_size=batch_size,
    save_interval=int(iters/10),
    log_iters=10,
    num_workers=0,
    losses=losses,
    use_vdl=True)

import paddlefrom paddleseg.models import AttentionUNetimport paddleseg.transforms as Tfrom paddleseg.core import inferimport osimport cv2from tqdm import tqdmfrom PIL import Imageimport numpy as npimport patta as ttadef nn_infer(model, imgs_path, is_tta=True):
    if not os.path.exists('Disc_Segmentation'):
        os.mkdir('Disc_Segmentation')    # 预测结果
    transforms = T.Compose([
        T.Resize(target_size=(1120, 1120)),
        T.Normalize()
    ])    # 循环预测和保存
    for img_path in tqdm(imgs_path):
        H, W, _ = np.asarray(Image.open(img_path)).shape
        img, _ = transforms(img_path)  # 进行数据预处理
        img = paddle.to_tensor(img[np.newaxis, :])  # C,H,W -> 1,C,H,W
        # TTA
        if is_tta == True:
            tta_pres = paddle.zeros([1, 4, 1120, 1120])            for tta_transform in tta.aliases.hflip_transform():
                tta_img = tta_transform.augment_image(img)  # TTA_transforms
                tta_pre = infer.inference(model, tta_img)  # 预测
                deaug_pre = tta_transform.deaugment_mask(tta_pre)
                tta_pres += deaug_pre
            pre = tta_pres / 2.
        else:
            pre = infer.inference(model, img)  # 预测
        pred = paddle.argmax(pre, axis=1).numpy().reshape((1120, 1120)).astype('uint8') * 255
        pred = cv2.resize(pred, (W, H), interpolation=cv2.INTER_NEAREST)
        pil_img = Image.fromarray(pred)
        pil_img.save(os.path.join('Disc_Segmentation', img_path.split('/')[-1].replace('jpg', 'png')), 'png')# 网络准备model_path='save_output/last_model/model.pdparams'model = AttentionUNet(num_classes=2, pretrained=None)
params = paddle.load(model_path)
model.set_dict(params)
model.eval()# 预测文件set_path = 'dataset/PALM-Testing400-Images'names = os.listdir(set_path)
imgs_path = []for name in names:
    imgs_path.append(os.path.join(set_path, name))# 预测nn_infer(model, imgs_path, is_tta=False)

import numpy as npfrom PIL import Imageimport matplotlib.pyplot as plt

%matplotlib inline

img_path = 'dataset/PALM-Testing400-Images/T0255.jpg'inf_path = 'Disc_Segmentation/T0255.png'img = Image.open(img_path)
inf = Image.open(inf_path)
plt.figure(figsize=(10, 5))
plt.subplot(121);plt.imshow(img);plt.title('img')
plt.subplot(122);plt.imshow(inf);plt.title('inf')
plt.show()