def build_model(self):
"""Build generator and discriminator."""
if self.model_type == 'UNet':
self.unet = UNet(n_channels=1, n_classes=1)
elif self.model_type == 'R2U_Net':
self.unet = R2U_Net(img_ch=3, output_ch=1, t=self.t) # TODO: changed for green image channel
elif self.model_type == 'AttU_Net':
self.unet = AttU_Net(img_ch=1, output_ch=1)
elif self.model_type == 'R2AttU_Net':
self.unet = R2AttU_Net(img_ch=3, output_ch=1, t=self.t)
elif self.model_type == 'Iternet':
self.unet = Iternet(n_channels=1, n_classes=1)
elif self.model_type == 'AttUIternet':
self.unet = AttUIternet(n_channels=1, n_classes=1)
elif self.model_type == 'R2UIternet':
self.unet = R2UIternet(n_channels=3, n_classes=1)
elif self.model_type == 'NestedUNet':
self.unet = NestedUNet(in_ch=1, out_ch=1)
elif self.model_type == "AG_Net":
self.unet = AG_Net(n_classes=1, bn=True, BatchNorm=False)
self.optimizer = optim.Adam(list(self.unet.parameters()),
self.lr,
betas=tuple(self.beta_list))
self.unet.to(self.device)
def build_model(self):
"""Build generator and discriminator."""
if self.model_type == 'U_Net':
self.unet = U_Net(img_ch=3, output_ch=1)
elif self.model_type == 'R2U_Net':
self.unet = R2U_Net(img_ch=3, output_ch=1, t=self.t)
elif self.model_type == 'AttU_Net':
self.unet = AttU_Net(img_ch=3, output_ch=1)
elif self.model_type == 'R2AttU_Net':
self.unet = R2AttU_Net(img_ch=3, output_ch=1, t=self.t)
elif self.model_type == 'ASM_U_Net':
self.unet = Structure_U_Net(img_ch=3, output_ch=1)
self.analyzer = Analyzer_U_Net(img_ch=1, output_ch=1)
self.optimizer_unet = optim.Adam(params=list(self.unet.parameters()),
lr=self.lr,
weight_decay=1e-5)
self.optimizer_analyzer = optim.Adam(params=list(
self.analyzer.parameters()),
lr=self.lr * 0.2,
weight_decay=1e-5)
self.unet.to(self.device)
self.unet = torch.nn.DataParallel(self.unet)
self.analyzer.to(self.device)
self.analyzer = torch.nn.DataParallel(self.analyzer)
def build_model(self):
"""Build our deep learning model."""
if self.model_type == 'U_Net':
self.unet = U_Net(img_ch=1,
output_ch=1,
first_layer_numKernel=self.first_layer_numKernel)
elif self.model_type == 'R2U_Net':
self.unet = R2U_Net(
img_ch=1,
output_ch=1,
t=self.t,
first_layer_numKernel=self.first_layer_numKernel)
elif self.model_type == 'AttU_Net':
self.unet = AttU_Net(
img_ch=1,
output_ch=1,
first_layer_numKernel=self.first_layer_numKernel)
elif self.model_type == 'R2AttU_Net':
self.unet = R2AttU_Net(
img_ch=1,
output_ch=1,
t=self.t,
first_layer_numKernel=self.first_layer_numKernel)
elif self.model_type == 'ResAttU_Net':
self.unet = ResAttU_Net(
UnetLayer=self.UnetLayer,
img_ch=1,
output_ch=1,
first_layer_numKernel=self.first_layer_numKernel)
if self.initialization != 'NA':
init_weights(self.unet, init_type=self.initialization)
self.unet.to(self.device)
def build_model(self):
"""Build generator and discriminator."""
if self.model_type == 'U_Net':
self.unet = U_Net(img_ch=self.img_ch, output_ch=1)
elif self.model_type == 'R2U_Net':
self.unet = R2U_Net(img_ch=self.img_ch, output_ch=1, t=self.t)
elif self.model_type == 'AttU_Net':
self.unet = AttU_Net(img_ch=self.img_ch, output_ch=1)
elif self.model_type == 'R2AttU_Net':
self.unet = R2AttU_Net(img_ch=self.img_ch, output_ch=1, t=self.t)
elif self.model_type == 'MixU_Net':
self.unet = MixU_Net(img_ch=self.img_ch, output_ch=1)
elif self.model_type == 'MixAttU_Net':
self.unet = MixAttU_Net(img_ch=self.img_ch, output_ch=1)
elif self.model_type == 'MixR2U_Net':
self.unet = MixR2U_Net(img_ch=self.img_ch, output_ch=1)
elif self.model_type == 'MixR2AttU_Net':
self.unet = MixR2AttU_Net(img_ch=self.img_ch, output_ch=1)
elif self.model_type == 'GhostU_Net':
self.unet = GhostU_Net(img_ch=self.img_ch, output_ch=1)
elif self.model_type == 'GhostU_Net1':
self.unet = GhostU_Net1(img_ch=self.img_ch, output_ch=1)
elif self.model_type == 'GhostU_Net2':
self.unet = GhostU_Net2(img_ch=self.img_ch, output_ch=1)
#pytorch_total_params = sum(p.numel() for p in self.unet.parameters() if p.requires_grad)
#print (pytorch_total_params)
#raise
self.optimizer = optim.Adam(list(self.unet.parameters()), self.lr,
[self.beta1, self.beta2])
self.unet.to(self.device)
def build_model(self):
"""Build generator and discriminator."""
if self.model_type == 'U_Net':
self.unet = U_Net(img_ch=3, output_ch=2)
elif self.model_type == 'R2U_Net':
self.unet = R2U_Net(img_ch=3, output_ch=1, t=self.t)
elif self.model_type == 'AttU_Net':
self.unet = AttU_Net(img_ch=3, output_ch=1)
elif self.model_type == 'R2AttU_Net':
self.unet = R2AttU_Net(img_ch=3, output_ch=1, t=self.t)
self.optimizer = optim.Adam(list(self.unet.parameters()), self.lr,
[self.beta1, self.beta2])
self.unet.to(self.device)
def build_model(self):
# Load required model
if self.model_type == 'U_Net':
self.unet = U_Net(img_ch=3, output_ch=1)
elif self.model_type == 'R2U_Net':
self.unet = R2U_Net(img_ch=3, output_ch=1, t=self.t)
elif self.model_type == 'AttU_Net':
self.unet = AttU_Net(img_ch=3, output_ch=1)
elif self.model_type == 'R2AttU_Net':
self.unet = R2AttU_Net(img_ch=3, output_ch=1, t=self.t)
# Load optimizer
self.optimizer = optim.Adam(list(self.unet.parameters()), self.lr,
[self.beta1, self.beta2])
# Move model to device
self.unet.to(self.device)
def build_model(self, config):
"""Build generator and discriminator."""
if self.model_type == 'U_Net':
self.unet = U_Net(img_ch=3, output_ch=1)
elif self.model_type == 'R2U_Net':
self.unet = R2U_Net(img_ch=3, output_ch=1, t=self.t)
elif self.model_type == 'AttU_Net':
self.unet = AttU_Net(img_ch=3, output_ch=1)
elif self.model_type == 'R2AttU_Net':
self.unet = R2AttU_Net(img_ch=3, output_ch=1, t=self.t)
# # Load the pretrained Encoder
# if os.path.isfile(config.pretrained):
# self.unet.load_state_dict(torch.load(config.pretrained))
# print('%s is Successfully Loaded from %s'%(self.model_type,config.pretrained))
self.optimizer = optim.Adam(list(self.unet.parameters()), self.lr,
[self.beta1, self.beta2])
self.unet.to(self.device)
def build_model(self):
"""Build generator and discriminator."""
if self.model_type == 'U_Net':
self.unet = U_Net(UnetLayer=self.UnetLayer,
img_ch=self.img_ch,
output_ch=self.output_ch,
first_layer_numKernel=self.first_layer_numKernel)
elif self.model_type == 'R2U_Net':
self.unet = R2U_Net(
img_ch=self.img_ch,
output_ch=self.output_ch,
t=self.t,
first_layer_numKernel=self.first_layer_numKernel)
elif self.model_type == 'AttU_Net':
self.unet = AttU_Net(
img_ch=self.img_ch,
output_ch=self.output_ch,
first_layer_numKernel=self.first_layer_numKernel)
elif self.model_type == 'R2AttU_Net':
self.unet = R2AttU_Net(
img_ch=self.img_ch,
output_ch=self.output_ch,
t=self.t,
first_layer_numKernel=self.first_layer_numKernel)
elif self.model_type == 'ResAttU_Net':
self.unet = ResAttU_Net(
UnetLayer=self.UnetLayer,
img_ch=self.img_ch,
output_ch=self.output_ch,
first_layer_numKernel=self.first_layer_numKernel)
if self.optimizer_choice == 'Adam':
self.optimizer = optim.Adam(list(self.unet.parameters()),
self.initial_lr,
[self.beta1, self.beta2])
elif self.optimizer_choice == 'SGD':
self.optimizer = optim.SGD(list(self.unet.parameters()),
self.initial_lr, self.momentum)
else:
pass
self.unet.to(self.device)
def build_model(self):
"""Build generator and discriminator."""
if self.model_type == 'U_Net':
self.unet = U_Net(img_ch=3, output_ch=3)
elif self.model_type == 'R2U_Net':
print("------> using R2U <--------")
self.unet = R2U_Net(img_ch=3, output_ch=3, t=self.t)
elif self.model_type == 'AttU_Net':
print("------> using AttU <--------")
self.unet = AttU_Net(img_ch=3, output_ch=3)
elif self.model_type == 'R2AttU_Net':
print("------> using R2-AttU <--------")
self.unet = R2AttU_Net(img_ch=3, output_ch=3, t=self.t)
elif self.model_type == 'ABU_Net':
print("------> using ABU_Net <--------")
self.unet = U_Net_AB(img_ch=3, output_ch=1)
elif self.model_type == 'Multi_Task':
print("------> using Multi_Task Learning <--------")
model = torch.hub.load('pytorch/vision',
'mobilenet_v2',
pretrained=True)
model_infeatures_final_layer = model.classifier[1].in_features
model.classifier = torch.nn.Sequential(
*list(model.classifier.children())[:-1])
for param in model.parameters():
param.requires_grad = True
for param in model.features[18].parameters():
param.requires_grad = True
for param in model.classifier.parameters():
param.requires_grad = True
model_trained_mobilenet = model
print("All trainable parameters of model are")
for name, param in model_trained_mobilenet.named_parameters():
if param.requires_grad:
print(name, param.shape)
self.unet = multi_task_model_classification(
model_trained_mobilenet)
self.optimizer = optim.AdamW(list(self.unet.parameters()), self.lr,
[self.beta1, self.beta2])
self.unet.to(self.device)
def build_model(self):
"""Build generator and discriminator."""
if self.model_type == 'U_Net':
self.unet = U_Net(img_ch=1, output_ch=1)
elif self.model_type == 'R2U_Net':
self.unet = R2U_Net(img_ch=1, output_ch=1, t=self.t)
elif self.model_type == 'AttU_Net':
self.unet = AttU_Net(img_ch=1, output_ch=1)
elif self.model_type == 'R2AttU_Net':
self.unet = R2AttU_Net(img_ch=1, output_ch=1, t=self.t)
if self.optimizer_choice == 'Adam':
self.optimizer = optim.Adam(list(self.unet.parameters()),
self.initial_lr,
[self.beta1, self.beta2])
elif self.optimizer_choice == 'SGD':
self.optimizer = optim.SGD(list(self.unet.parameters()),
self.initial_lr, self.momentum)
else:
pass
self.unet.to(self.device)
out_files = args.output
return out_files
def mask_to_image(mask):
return Image.fromarray((mask * 255).astype(np.uint8))
if __name__ == "__main__":
args = get_args()
in_files = args.input
out_files = get_output_filenames(args)
#net = UNet(n_channels=3, n_classes=1)
net = R2AttU_Net()
logging.info("Loading model {}".format(args.model))
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
net.to(device=device)
net.load_state_dict(torch.load(args.model, map_location=device))
logging.info("Model loaded !")
for i, fn in enumerate(in_files):
logging.info("\nPredicting image {} ...".format(fn))
img = Image.open(fn)
mask = predict_img(net=net,