【深度學習-圖像識別】腦PET圖像分析和疾病預測Baseline

這裏提供一個baseline；

比賽簡介：

比賽地址：http://challenge.xfyun.cn/topic/info?type=PET

數據長這個樣子：

Baseline：

1. 讀取數據

base = './train'
count = 0
with open(os.path.join('train_list.txt'), 'w') as f:
    for i, cls_fold in enumerate(['AD', 'CN']):
        try:
            cls_base = os.path.join(base, cls_fold)
            files = os.listdir(cls_base)
            print('{} train num:'.format(cls_fold), len(files))
            for pt in files:
                img = os.path.join(cls_fold, pt)
                info = img + ' ' + str(i) + '\n'
                f.write(info)
                count += 1
        except Exception as e:
            continue
print(count)
with open(os.path.join('labels.txt'), 'w') as f:
    for i, cls_fold in enumerate(['AD', 'CN']):
        f.write(cls_fold+'\n')

print('done.')

2. 定義數據加載器：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.utils.data import DataLoader
import torchvision
from torchvision.transforms import transforms
from torchvision import models
from torchvision.models import ResNet
import numpy as np
import matplotlib.pyplot as plt
import os


data_dir = './'

train_dataset = torchvision.datasets.ImageFolder(root=os.path.join(data_dir, 'train'),
                          transform=transforms.Compose(
                          [
                           transforms.Resize(299),
                           transforms.ToTensor(),
                           transforms.Normalize(
                           mean=(0.485, 0.456, 0.406),
                           std=(0.229, 0.224, 0.225))
                          ]))

val_dataset = torchvision.datasets.ImageFolder(root=os.path.join(data_dir, 'train'),
                          transform=transforms.Compose(
                          [
                           transforms.Resize(299),
                           transforms.ToTensor(),
                           transforms.Normalize(
                           mean=(0.485, 0.456, 0.406),
                           std=(0.229, 0.224, 0.225))
                          ]))

train_dataloader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
val_dataloader = DataLoader(dataset=val_dataset, batch_size=1, shuffle=True)

# 類別名稱
class_names = train_dataset.classes
print('class_names:{}'.format(class_names))

# 訓練設備  CPU/GPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('trian_device:{}'.format(device.type))

3. 構建模型：

這裏使用vgg16：

# -------------------------模型選擇，優化方法， 學習率策略----------------------
model_name = 'vgg16'
model = models.vgg16(pretrained=True)

num_classes = 2

# 改變全連接層，2分類問題，out_features = 2
# model.aux_logits = False
# model.AuxLogits.fc = nn.Linear(768, num_classes)
# model.fc = nn.Linear(2048, num_classes)
model.classifier[6] = nn.Linear(4096,num_classes)

# 模型遷移到CPU/GPU
model = model.to(device)

# 定義損失函數
loss_fc = nn.CrossEntropyLoss()

# 選擇優化方法
optimizer = optim.Adam(model.parameters(), lr=2e-4)

# 學習率調整策略
exp_lr_scheduler = lr_scheduler.StepLR(optimizer=optimizer, step_size=30, gamma=0.5)  # step_size

print('done.')

4. 訓練模型：

# ----------------訓練過程-----------------
num_epochs = 101

for epoch in range(num_epochs):

    running_loss = 0.0
    exp_lr_scheduler.step()

    for i, sample_batch in enumerate(train_dataloader):
        inputs = sample_batch[0]
        labels = sample_batch[1]

        model.train()

        # GPU/CPU
        inputs = inputs.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()

        # foward
        outputs = model(inputs)

        # loss
        loss = loss_fc(outputs, labels)

        # loss求導，反向
        loss.backward()

        # 優化
        optimizer.step()

        #
        running_loss += loss.item()

        # 測試
        if epoch % 5 == 0 and i == 1:
            correct = 0
            total = 0
            model.eval()
            for images_test, labels_test in val_dataloader:
                images_test = images_test.to(device)
                labels_test = labels_test.to(device)

                outputs_test = model(images_test)
                _, prediction = torch.max(outputs_test, 1)
                correct += (torch.sum((prediction == labels_test))).item()
               # print(prediction, labels_test, correct)
                total += labels_test.size(0)
            print('[{}, {}] running_loss = {:.5f} accurcay = {:.5f}'.format(epoch + 1, i + 1, running_loss / 20,
                                                                        correct / total))
            running_loss = 0.0

print('training finish !')
torch.save(model.state_dict(), './model.pth')

5. 生成預測結果：

test_base = './test/AD&CN'
import pandas as pd
import cv2
from PIL import Image
from torchvision import transforms
from torch.autograd import Variable as V
import torch as t
import pickle as pkl

# model_name = 'vgg16_bn'

trans=transforms.Compose(
               [
               transforms.Resize(299),
               transforms.ToTensor(),
               transforms.Normalize(
               mean=(0.485, 0.456, 0.406),
               std=(0.229, 0.224, 0.225))
               ]
            )

model.eval()

submission = {'uuid':[], 'label':[]}

for i in tqdm(range(1, 1000+1)):
    im_pt = os.path.join(test_base, '{}.png'.format(i))
    img = cv2.imread(im_pt)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = Image.fromarray(img)
    img = trans(img)
    img = img.unsqueeze(0)
    img = V(img.cuda())
    predict = model(img)
    predict = F.softmax(predict)[0]
    temp = {'AD':float(predict[0]), 'CN':float(predict[1]), 'id':i}
    _, pred = torch.max(predict, 0)
    pred = class_names[int(pred)]
    submission['uuid'].append(i)
    submission['label'].append(pred)

data = pd.DataFrame(submission)
data.to_csv('./submission_{}.csv'.format(model_name), index=False)

print('done.')

最終提交結果：

大概 0.71 左右；