def transform_train(self):
temp = []
temp.append(tr.Resize(self.args.input_size))
temp.append(tr.RandomHorizontalFlip())
temp.append(tr.RandomRotate(15))
temp.append(tr.RandomCrop(self.args.input_size))
temp.append(tr.ToTensor())
composed_transforms = transforms.Compose(temp)
return composed_transforms
def transform_tr(self, sample):
composed_transforms = transforms.Compose([
tr.RandomHorizontalFlip(),
tr.RandomScaleCrop(base_size=self.args.base_size, crop_size=self.args.crop_size),
tr.RandomGaussianBlur(),
tr.RandomRotate(20),
tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
tr.ToTensor()])
return composed_transforms(sample)
def transform_tr(self, sample):
"""Image transformations for training"""
targs = self.transform
method = targs["method"]
pars = targs["parameters"]
composed_transforms = transforms.Compose([
tr.FixedResize(size=pars["outSize"]),
tr.RandomRotate(degree=(90)),
tr.RandomScaleCrop(baseSize=pars["baseSize"], cropSize=pars["outSize"], fill=255),
tr.Normalize(mean=pars["mean"], std=pars["std"]),
tr.ToTensor()])
return composed_transforms(sample)
def transform_train(self, sample):
composed_transforms = transforms.Compose([
tr.RandomHorizontalFlip(),
tr.RandomVerticalFlip(),
# tr.RandomScaleCrop(base_size=self.args.base_size, crop_size=self.args.crop_size, fill=255),
# tr.FixedResize(size=self.args.crop_size),
tr.RandomRotate(),
tr.RandomGammaTransform(),
tr.RandomGaussianBlur(),
tr.RandomNoise(),
tr.Normalize(mean=(0.544650, 0.352033, 0.384602, 0.352311), std=(0.249456, 0.241652, 0.228824, 0.227583)),
tr.ToTensor()])
return composed_transforms(sample)
def transform_tr(self, sample):
temp = []
if self.args.rotate > 0:
temp.append(tr.RandomRotate(self.args.rotate))
temp.append(tr.RandomScale(rand_resize=self.args.rand_resize))
temp.append(tr.RandomCrop(self.args.input_size))
temp.append(tr.RandomHorizontalFlip())
temp.append(
tr.Normalize(mean=self.args.normal_mean, std=self.args.normal_std))
if self.args.noise_param is not None:
temp.append(
tr.GaussianNoise(mean=self.args.noise_param[0],
std=self.args.noise_param[1]))
temp.append(tr.ToTensor())
composed_transforms = transforms.Compose(temp)
return composed_transforms(sample)
def transform_tr(self, sample):
color_transforms = [
transforms.RandomApply([transforms.ColorJitter(brightness=0.1)
]), # brightness
transforms.RandomApply([transforms.ColorJitter(contrast=0.1)
]), # contrast
transforms.RandomApply([transforms.ColorJitter(saturation=0.1)
]), # saturation
transforms.RandomApply([transforms.ColorJitter(hue=0.05)])
] # hue
joint_transforms = transforms.Compose([
tr.RandomHorizontalFlip(),
tr.RandomScaleCrop(base_size=self.args.base_size,
crop_size=self.args.crop_size,
fill=255),
tr.equalize(),
tr.RandomGaussianBlur(),
tr.RandomRotate(degree=7)
])
image_transforms = transforms.Compose([
transforms.RandomOrder(color_transforms),
transforms.RandomGrayscale(p=0.3)
])
normalize_transforms = transforms.Compose([
tr.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225)),
tr.ToTensor()
])
tmp_sample = joint_transforms(sample)
tmp_sample['image'] = image_transforms(tmp_sample['image'])
tmp_sample = normalize_transforms(tmp_sample)
return tmp_sample
net.cuda()
net = torch.nn.DataParallel(net,
device_ids=range(torch.cuda.device_count()))
if resume_epoch != nEpochs:
# Logging into Tensorboard
log_dir = os.path.join(
save_dir, 'models',
datetime.now().strftime('%b%d_%H-%M-%S') + '_' + socket.gethostname())
optimizer = optim.Adam(net.parameters(), lr=p['lr'], weight_decay=p['wd'])
p['optimizer'] = str(optimizer)
composed_transforms_tr = transforms.Compose([
tr.RandomSized(512),
tr.RandomRotate(15),
tr.RandomHorizontalFlip(),
tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
tr.ToTensor()
])
composed_transforms_ts = transforms.Compose([
tr.FixedResize(size=(512, 512)),
tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
tr.ToTensor()
])
# voc_train = pascal.VOCSegmentation(split='train', transform=composed_transforms_tr)
# voc_val = pascal.VOCSegmentation(split='val', transform=composed_transforms_ts)
ROOT = 'dataset/ORIGA'
voc_train = ImageFolder(root_path=ROOT, datasets='ORIGA')
mask[mask == _validc] = self.class_map[_validc]
return mask
if __name__ == '__main__':
from dataloaders import custom_transforms as tr
from dataloaders.utils import decode_segmap
from torch.utils.data import DataLoader
from torchvision import transforms
import matplotlib.pyplot as plt
composed_transforms_tr = transforms.Compose([
tr.RandomHorizontalFlip(),
tr.RandomScale((0.5, 0.75)),
tr.RandomCrop((512, 1024)),
tr.RandomRotate(5),
tr.ToTensor()
])
cityscapes_train = CityscapesSegmentation(split='train',
transform=composed_transforms_tr)
dataloader = DataLoader(cityscapes_train,
batch_size=2,
shuffle=True,
num_workers=2)
for ii, sample in enumerate(dataloader):
for jj in range(sample["image"].size()[0]):
img = sample['image'].numpy()
gt = sample['label'].numpy()
def main():
# Add default values to all parameters
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument('-g', '--gpu', type=int, default=0, help='gpu id')
parser.add_argument('--resume', default=None, help='checkpoint path')
parser.add_argument(
'--coefficient', type=float, default=0.01, help='balance coefficient'
)
parser.add_argument(
'--boundary-exist', type=bool, default=True, help='whether or not using boundary branch'
)
parser.add_argument(
'--dataset', type=str, default='refuge', help='folder id contain images ROIs to train or validation'
)
parser.add_argument(
'--batch-size', type=int, default=12, help='batch size for training the model'
)
# parser.add_argument(
# '--group-num', type=int, default=1, help='group number for group normalization'
# )
parser.add_argument(
'--max-epoch', type=int, default=300, help='max epoch'
)
parser.add_argument(
'--stop-epoch', type=int, default=300, help='stop epoch'
)
parser.add_argument(
'--warmup-epoch', type=int, default=-1, help='warmup epoch begin train GAN'
)
parser.add_argument(
'--interval-validate', type=int, default=1, help='interval epoch number to valide the model'
)
parser.add_argument(
'--lr-gen', type=float, default=1e-3, help='learning rate',
)
parser.add_argument(
'--lr-dis', type=float, default=2.5e-5, help='learning rate',
)
parser.add_argument(
'--lr-decrease-rate', type=float, default=0.2, help='ratio multiplied to initial lr',
)
parser.add_argument(
'--weight-decay', type=float, default=0.0005, help='weight decay',
)
parser.add_argument(
'--momentum', type=float, default=0.9, help='momentum',
)
parser.add_argument(
'--data-dir',
default='./fundus/',
help='data root path'
)
parser.add_argument(
'--out-stride',
type=int,
default=16,
help='out-stride of deeplabv3+',
)
parser.add_argument(
'--sync-bn',
type=bool,
default=False,
help='sync-bn in deeplabv3+',
)
parser.add_argument(
'--freeze-bn',
type=bool,
default=False,
help='freeze batch normalization of deeplabv3+',
)
args = parser.parse_args()
args.model = 'MobileNetV2'
now = datetime.now()
args.out = osp.join(here, 'logs', args.dataset, now.strftime('%Y%m%d_%H%M%S.%f'))
os.makedirs(args.out)
# save training hyperparameters or/and settings
with open(osp.join(args.out, 'config.yaml'), 'w') as f:
yaml.safe_dump(args.__dict__, f, default_flow_style=False)
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
cuda = torch.cuda.is_available()
torch.manual_seed(2020)
if cuda:
torch.cuda.manual_seed(2020)
import random
import numpy as np
random.seed(2020)
np.random.seed(2020)
# 1. loading data
composed_transforms_train = transforms.Compose([
tr.RandomScaleCrop(512),
tr.RandomRotate(),
tr.RandomFlip(),
tr.elastic_transform(),
tr.add_salt_pepper_noise(),
tr.adjust_light(),
tr.eraser(),
tr.Normalize_tf(),
tr.ToTensor()
])
composed_transforms_val = transforms.Compose([
tr.RandomCrop(512),
tr.Normalize_tf(),
tr.ToTensor()
])
data_train = DL.FundusSegmentation(base_dir=args.data_dir, dataset=args.dataset, split='train',
transform=composed_transforms_train)
dataloader_train = DataLoader(data_train, batch_size=args.batch_size, shuffle=True, num_workers=4,
pin_memory=True)
data_val = DL.FundusSegmentation(base_dir=args.data_dir, dataset=args.dataset, split='testval',
transform=composed_transforms_val)
dataloader_val = DataLoader(data_val, batch_size=args.batch_size, shuffle=False, num_workers=4, pin_memory=True)
# domain_val = DL.FundusSegmentation(base_dir=args.data_dir, dataset=args.datasetT, split='train',
# transform=composed_transforms_ts)
# domain_loader_val = DataLoader(domain_val, batch_size=args.batch_size, shuffle=False, num_workers=2,
# pin_memory=True)
# 2. model
model_gen = DeepLab(num_classes=2, backbone='mobilenet', output_stride=args.out_stride,
sync_bn=args.sync_bn, freeze_bn=args.freeze_bn).cuda()
model_bd = BoundaryDiscriminator().cuda()
model_mask = MaskDiscriminator().cuda()
start_epoch = 0
start_iteration = 0
# 3. optimizer
optim_gen = torch.optim.Adam(
model_gen.parameters(),
lr=args.lr_gen,
betas=(0.9, 0.99)
)
optim_bd = torch.optim.SGD(
model_bd.parameters(),
lr=args.lr_dis,
momentum=args.momentum,
weight_decay=args.weight_decay
)
optim_mask = torch.optim.SGD(
model_mask.parameters(),
lr=args.lr_dis,
momentum=args.momentum,
weight_decay=args.weight_decay
)
# breakpoint recovery
if args.resume:
checkpoint = torch.load(args.resume)
pretrained_dict = checkpoint['model_state_dict']
model_dict = model_gen.state_dict()
# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict)
# 3. load the new state dict
model_gen.load_state_dict(model_dict)
pretrained_dict = checkpoint['model_bd_state_dict']
model_dict = model_bd.state_dict()
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
model_dict.update(pretrained_dict)
model_bd.load_state_dict(model_dict)
pretrained_dict = checkpoint['model_mask_state_dict']
model_dict = model_mask.state_dict()
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
model_dict.update(pretrained_dict)
model_mask.load_state_dict(model_dict)
start_epoch = checkpoint['epoch'] + 1
start_iteration = checkpoint['iteration'] + 1
optim_gen.load_state_dict(checkpoint['optim_state_dict'])
optim_bd.load_state_dict(checkpoint['optim_bd_state_dict'])
optim_mask.load_state_dict(checkpoint['optim_mask_state_dict'])
trainer = Trainer.Trainer(
cuda=cuda,
model_gen=model_gen,
model_bd=model_bd,
model_mask=model_mask,
optimizer_gen=optim_gen,
optim_bd=optim_bd,
optim_mask=optim_mask,
lr_gen=args.lr_gen,
lr_dis=args.lr_dis,
lr_decrease_rate=args.lr_decrease_rate,
train_loader=dataloader_train,
validation_loader=dataloader_val,
out=args.out,
max_epoch=args.max_epoch,
stop_epoch=args.stop_epoch,
interval_validate=args.interval_validate,
batch_size=args.batch_size,
warmup_epoch=args.warmup_epoch,
coefficient=args.coefficient,
boundary_exist=args.boundary_exist
)
trainer.epoch = start_epoch
trainer.iteration = start_iteration
trainer.train()
def main():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter, )
parser.add_argument('-g', '--gpu', type=int, default=0, help='gpu id')
parser.add_argument('--resume', default=None, help='checkpoint path')
# configurations (same configuration as original work)
# https://github.com/shelhamer/fcn.berkeleyvision.org
parser.add_argument('--datasetS',
type=str,
default='refuge',
help='test folder id contain images ROIs to test')
parser.add_argument('--datasetT',
type=str,
default='Drishti-GS',
help='refuge / Drishti-GS/ RIM-ONE_r3')
parser.add_argument('--batch-size',
type=int,
default=8,
help='batch size for training the model')
parser.add_argument('--group-num',
type=int,
default=1,
help='group number for group normalization')
parser.add_argument('--max-epoch', type=int, default=200, help='max epoch')
parser.add_argument('--stop-epoch',
type=int,
default=200,
help='stop epoch')
parser.add_argument('--warmup-epoch',
type=int,
default=-1,
help='warmup epoch begin train GAN')
parser.add_argument('--interval-validate',
type=int,
default=10,
help='interval epoch number to valide the model')
parser.add_argument(
'--lr-gen',
type=float,
default=1e-3,
help='learning rate',
)
parser.add_argument(
'--lr-dis',
type=float,
default=2.5e-5,
help='learning rate',
)
parser.add_argument(
'--lr-decrease-rate',
type=float,
default=0.1,
help='ratio multiplied to initial lr',
)
parser.add_argument(
'--weight-decay',
type=float,
default=0.0005,
help='weight decay',
)
parser.add_argument(
'--momentum',
type=float,
default=0.99,
help='momentum',
)
parser.add_argument('--data-dir',
default='/home/sjwang/ssd1T/fundus/domain_adaptation/',
help='data root path')
parser.add_argument(
'--pretrained-model',
default='../../../models/pytorch/fcn16s_from_caffe.pth',
help='pretrained model of FCN16s',
)
parser.add_argument(
'--out-stride',
type=int,
default=16,
help='out-stride of deeplabv3+',
)
parser.add_argument(
'--sync-bn',
type=bool,
default=True,
help='sync-bn in deeplabv3+',
)
parser.add_argument(
'--freeze-bn',
type=bool,
default=False,
help='freeze batch normalization of deeplabv3+',
)
args = parser.parse_args()
args.model = 'FCN8s'
now = datetime.now()
args.out = osp.join(here, 'logs', args.datasetT,
now.strftime('%Y%m%d_%H%M%S.%f'))
os.makedirs(args.out)
with open(osp.join(args.out, 'config.yaml'), 'w') as f:
yaml.safe_dump(args.__dict__, f, default_flow_style=False)
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
cuda = torch.cuda.is_available()
torch.manual_seed(1337)
if cuda:
torch.cuda.manual_seed(1337)
# 1. dataset
composed_transforms_tr = transforms.Compose([
tr.RandomScaleCrop(512),
tr.RandomRotate(),
tr.RandomFlip(),
tr.elastic_transform(),
tr.add_salt_pepper_noise(),
tr.adjust_light(),
tr.eraser(),
tr.Normalize_tf(),
tr.ToTensor()
])
composed_transforms_ts = transforms.Compose(
[tr.RandomCrop(512),
tr.Normalize_tf(),
tr.ToTensor()])
domain = DL.FundusSegmentation(base_dir=args.data_dir,
dataset=args.datasetS,
split='train',
transform=composed_transforms_tr)
domain_loaderS = DataLoader(domain,
batch_size=args.batch_size,
shuffle=True,
num_workers=2,
pin_memory=True)
domain_T = DL.FundusSegmentation(base_dir=args.data_dir,
dataset=args.datasetT,
split='train',
transform=composed_transforms_tr)
domain_loaderT = DataLoader(domain_T,
batch_size=args.batch_size,
shuffle=False,
num_workers=2,
pin_memory=True)
domain_val = DL.FundusSegmentation(base_dir=args.data_dir,
dataset=args.datasetT,
split='train',
transform=composed_transforms_ts)
domain_loader_val = DataLoader(domain_val,
batch_size=args.batch_size,
shuffle=False,
num_workers=2,
pin_memory=True)
# 2. model
model_gen = DeepLab(num_classes=2,
backbone='mobilenet',
output_stride=args.out_stride,
sync_bn=args.sync_bn,
freeze_bn=args.freeze_bn).cuda()
model_dis = BoundaryDiscriminator().cuda()
model_dis2 = UncertaintyDiscriminator().cuda()
start_epoch = 0
start_iteration = 0
# 3. optimizer
optim_gen = torch.optim.Adam(model_gen.parameters(),
lr=args.lr_gen,
betas=(0.9, 0.99))
optim_dis = torch.optim.SGD(model_dis.parameters(),
lr=args.lr_dis,
momentum=args.momentum,
weight_decay=args.weight_decay)
optim_dis2 = torch.optim.SGD(model_dis2.parameters(),
lr=args.lr_dis,
momentum=args.momentum,
weight_decay=args.weight_decay)
if args.resume:
checkpoint = torch.load(args.resume)
pretrained_dict = checkpoint['model_state_dict']
model_dict = model_gen.state_dict()
# 1. filter out unnecessary keys
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict)
# 3. load the new state dict
model_gen.load_state_dict(model_dict)
pretrained_dict = checkpoint['model_dis_state_dict']
model_dict = model_dis.state_dict()
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
model_dis.load_state_dict(model_dict)
pretrained_dict = checkpoint['model_dis2_state_dict']
model_dict = model_dis2.state_dict()
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
model_dis2.load_state_dict(model_dict)
start_epoch = checkpoint['epoch'] + 1
start_iteration = checkpoint['iteration'] + 1
optim_gen.load_state_dict(checkpoint['optim_state_dict'])
optim_dis.load_state_dict(checkpoint['optim_dis_state_dict'])
optim_dis2.load_state_dict(checkpoint['optim_dis2_state_dict'])
optim_adv.load_state_dict(checkpoint['optim_adv_state_dict'])
trainer = Trainer.Trainer(
cuda=cuda,
model_gen=model_gen,
model_dis=model_dis,
model_uncertainty_dis=model_dis2,
optimizer_gen=optim_gen,
optimizer_dis=optim_dis,
optimizer_uncertainty_dis=optim_dis2,
lr_gen=args.lr_gen,
lr_dis=args.lr_dis,
lr_decrease_rate=args.lr_decrease_rate,
val_loader=domain_loader_val,
domain_loaderS=domain_loaderS,
domain_loaderT=domain_loaderT,
out=args.out,
max_epoch=args.max_epoch,
stop_epoch=args.stop_epoch,
interval_validate=args.interval_validate,
batch_size=args.batch_size,
warmup_epoch=args.warmup_epoch,
)
trainer.epoch = start_epoch
trainer.iteration = start_iteration
trainer.train()
sample = {'image': _img, 'label': _target}
if self.transform is not None:
sample = self.transform(sample)
return sample
def _make_img_gt_point_pair(self, index):
#Read image and target
_img = Image.open(self.images[index].convert('RGB'))
_target = Image.open(self.categories[index])
return _img, _target
def __str__(self):
return 'VOC2012(split=' + str(self.split) + ')'
if __name__ == '__main__':
composed_transforms_tr = transforms.Compose([tr.RandomHorizontalFlip(), tr.RandomSized(512), tr.RandomRotate(15), tr.ToTensor()])
voc_train = PascalVOC(split='train', transform=composed_transforms_tr)
dataloader = DataLoader(voc_train, batch_size=5, shuffle=True, num_workers=2)
for ii, sample in tqdm(enumberate(dataloader)):
for jj in tqdm(range(sample["image"].size()[0])):
img = sample['image'].numpy()
gt = sample['label'].numpy()
tmp = np.array(get[jj]).astype(np.uint8)
tmp = np.squeeze(tmp, axis=0)
segmap = decode_segmap(tmp, dataset = 'pascal')
img_tmp = np.transpose(img[jj], axes=[1,2,0]).astype(np.uint8)
plt.figure()
plt.title('display')
plt.subplot(211)
