Spaces:
Build error
Build error
File size: 2,121 Bytes
8120b87 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from PIL import Image
# Define the ConvAutoencoder model structure
class ConvAutoencoder(nn.Module):
def __init__(self):
super(ConvAutoencoder, self).__init__()
# Encoder
self.encoder = nn.Sequential(
nn.Conv2d(1, 32, 3, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(32, 64, 3, stride=2, padding=1),
nn.ReLU()
)
# Decoder
self.decoder = nn.Sequential(
nn.ConvTranspose2d(64, 32, 2, stride=2, padding=1),
nn.ReLU(),
nn.ConvTranspose2d(32, 1, 2, stride=2, padding=1),
nn.Sigmoid()
)
def forward(self, x):
x = self.encoder(x)
x = self.decoder(x)
return x
# Load and resize the training slice
processed_data_path = 'BraTS20_Training_001_best_slice.npy' # Replace with your actual path
best_slice = np.load(processed_data_path)
# Resize the best slice to (256, 256)
transform = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((256, 256)), # Ensures fixed input size
transforms.ToTensor()
])
# Apply transformation to ensure correct size
best_slice = transform(best_slice).unsqueeze(0) # Shape: [1, 1, 256, 256]
# Initialize model, loss function, and optimizer
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = ConvAutoencoder().to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# Training loop
num_epochs = 100
best_slice = best_slice.to(device)
for epoch in range(num_epochs):
model.train()
optimizer.zero_grad()
output = model(best_slice)
loss = criterion(output, best_slice)
loss.backward()
optimizer.step()
if (epoch + 1) % 10 == 0:
print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}")
# Save the trained model weights
torch.save(model.state_dict(), "conv_autoencoder_model.pth")
print("Model training complete and saved as conv_autoencoder_model.pth")
|