class SH_UNet(nn.Module):
def __init__(self, path_to_shape_net_weights='', n_classes=2):
super(SH_UNet, self).__init__()
self.unet = EESPNet_Seg(classes=2, s=2)
self.shapeUNet = ShapeUNet((2, 1024, 1024))
self.softmax = nn.Softmax(dim=1)
if path_to_shape_net_weights:
self.shapeUNet.load_state_dict(
torch.load(path_to_shape_net_weights))
def forward(self, x, only_encode=False):
if only_encode:
_, encoded_mask = self.shapeUNet(x)
return encoded_mask
if self.training:
unet_prediction, unet_prediction_1 = self.unet(x)
else:
unet_prediction = self.unet(x)
softmax_unet_prediction = self.softmax(unet_prediction) #.detach()
shape_net_final_prediction, shape_net_encoded_prediction = self.shapeUNet(
softmax_unet_prediction)
if self.training:
return unet_prediction, unet_prediction_1, shape_net_encoded_prediction, shape_net_final_prediction
else:
return unet_prediction, shape_net_encoded_prediction, shape_net_final_prediction
class SH_UNet(nn.Module):
def __init__(self, path_to_shape_net_weights='', n_classes=2):
super(SH_UNet, self).__init__()
self.unet = UNet((3, 544, 544))
self.shapeUNet = ShapeUNet((4, 544, 544))
self.softmax = nn.Softmax(dim=1)
if path_to_shape_net_weights:
self.shapeUNet.load_state_dict(
torch.load(path_to_shape_net_weights))
def forward(self, x, only_encode=False):
if only_encode:
_, encoded_mask = self.shapeUNet(x)
return encoded_mask
unet_prediction = self.unet(x)
softmax_unet_prediction = self.softmax(unet_prediction) #.detach()
shape_net_final_prediction, shape_net_encoded_prediction = self.shapeUNet(
softmax_unet_prediction)
return unet_prediction, shape_net_encoded_prediction, shape_net_final_prediction
ultrasound_train = UltrasoundDataShapeNet(ROOT_DIR, partition)
train_loader = torch.utils.data.DataLoader(
ultrasound_train,
batch_size=2,
shuffle=True,
)
partition = 'val'
ultrasound_val = UltrasoundDataShapeNet(ROOT_DIR, partition)
val_loader = torch.utils.data.DataLoader(ultrasound_val,
batch_size=1,
shuffle=False)
# # Create model
model = ShapeUNet((4, 544, 544))
weights = torch.load(
'weights_shape_net/Unet_manually_corrupted_lumen0.113253.pth')
model.load_state_dict(weights)
model.to(device)
# Specify optimizer and criterion
lr = 1e-4
optimizer = Adam(model.parameters(), lr=lr)
NUM_OF_EPOCHS = 50
train(model, train_loader, val_loader, optimizer, NUM_OF_EPOCHS,
'weights_shape_net/')
# import torch
# import torch.nn as nn
# from torch.optim import Adam
# import torch.nn.functional as F
# Create Data Loaders
partition = 'train'
ultrasound_train = UltrasoundData(ROOT_DIR, partition)
train_loader = torch.utils.data.DataLoader(
ultrasound_train,
batch_size=1,
shuffle=True,
)
partition = 'val'
ultrasound_val = UltrasoundData(ROOT_DIR, partition)
val_loader = torch.utils.data.DataLoader(ultrasound_val,
batch_size=1,
shuffle=False)
# # Create models
PATH_TO_SHAPE_WEIGHT_MODEL = 'weights_shape_net/Unet_manually_corrupted_lumen0.113253.pth'
# Specify optimizer and criterion
NUM_OF_EPOCHS = 50
unet = UNet((3, 544, 544))
shapeUNet = ShapeUNet((4, 544, 544))
shapeUNet.load_state_dict(torch.load(PATH_TO_SHAPE_WEIGHT_MODEL))
unet.to(device)
shapeUNet.to(device)
lr = 1e-4
optimizer1 = Adam(shapeUNet.parameters(), lr=lr)
optimizer2 = Adam(unet.parameters(), lr=lr)
train(unet, shapeUNet, train_loader, val_loader, optimizer1, optimizer2,
NUM_OF_EPOCHS, 'weights_shape_net/')