import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np
from sklearn.metrics import roc_curve, auc
class ImageDataset(Dataset):
def __init__(self, folder_path):
self.folder_path = folder_path
self.file_list = os.listdir(folder_path)
def __len__(self):
return len(self.file_list)
def __getitem__(self, index):
file_name = self.file_list[index]
file_path = os.path.join(self.folder_path, file_name)
image = np.load(file_path)
image = torch.from_numpy(image).float()
label = self.folder_path.split('/')[-1]
if label == 'no':
label = 0
elif label == 'sphere':
label = 1
elif label == 'vort':
label = 2
return image, label
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 32, 3)
self.conv2 = nn.Conv2d(32, 64, 3)
self.fc1 = nn.Linear(64 * 34 * 34, 128)
self.fc2 = nn.Linear(128, 3)
def forward(self, x):
x = self.conv1(x)
x = nn.functional.relu(x)
x = nn.functional.max_pool2d(x, 2)
x = self.conv2(x)
x = nn.functional.relu(x)
x = nn.functional.max_pool2d(x, 2)
x = x.view(-1, 64 * 34 * 34)
x = self.fc1(x)
x = nn.functional.relu(x)
x = self.fc2(x)
return x
train_dataset = ImageDataset('path/to/training/folder')
val_dataset = ImageDataset('path/to/validation/folder')
batch_size = 32
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
net = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)
num_epochs = 10
for epoch in range(num_epochs):
running_loss = 0.0
running_corrects = 0
net.train()
for inputs, labels in train_loader:
optimizer.zero_grad()
outputs = net(inputs.unsqueeze(1))
_, preds = torch.max(outputs, 1)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
train_loss = running_loss / len(train_dataset)
train_acc = running_corrects.double() / len(train_dataset)
running_loss = 0.0
running_corrects = 0
net.eval()
for inputs, labels in val_loader:
outputs = net(inputs.unsqueeze(1))
_, preds = torch.max(outputs, 1)
loss = criterion(outputs, labels)
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
val_loss = running_loss / len(val_dataset)
val_acc = running_corrects.double() / len(val_dataset)
print('Epoch [{}/{}], Train Loss: {:.4f}, Train Acc: {:.4f}, Val Loss: {:.4f},
Val Acc:
from sklearn.metrics import roc_curve, auc
net.eval()
y_true = []
y_scores = []
with torch.no_grad():
for inputs, labels in val_loader:
outputs = net(inputs.unsqueeze(1))
probs = torch.softmax(outputs, dim=1)
y_scores.extend(probs.cpu().numpy())
y_true.extend(labels.cpu().numpy())
y_scores = np.array(y_scores)
y_true = np.array(y_true)
n_classes = 3
fpr = dict()
tpr = dict()
roc_auc = dict()
for i in range(n_classes):
fpr[i], tpr[i], _ = roc_curve(y_true[:, i], y_scores[:, i])
roc_auc[i] = auc(fpr[i], tpr[i])
plt.figure()
colors = ['red', 'green', 'blue']
for i, color in zip(range(n_classes), colors):
plt.plot(fpr[i], tpr[i], color=color, lw=2, label='ROC curve of class {0} (AUC
= {1:.2f})'.format(i, roc_auc[i]))
plt.plot([0, 1], [0, 1], 'k--', lw=2)
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title('Receiver Operating Characteristic Curve')
plt.legend(loc="lower right")
plt.show()
