def main():
args = parse_args()
random.seed(0)
torch.manual_seed(0)
if not args.nogpu:
torch.cuda.manual_seed_all(0)
if args.no_norm:
imgtr = [ToTensor()]
else:
imgtr = [ToTensor(),NormalizeOwn()]
# softmax
labtr = [IgnoreLabelClass(),ToTensorLabel()]
# labtr = [IgnoreLabelClass(),ToTensorLabel(tensor_type=torch.FloatTensor)]
# cotr = [RandomSizedCrop((320,320))] # (321,321)
cotr = [RandomSizedCrop3((320,320))]
print("dataset_dir: ", args.dataset_dir)
trainset_l = Corrosion(home_dir,args.dataset_dir,img_transform=Compose(imgtr),
label_transform=Compose(labtr),co_transform=Compose(cotr),
split=args.split,labeled=True)
trainloader_l = DataLoader(trainset_l,batch_size=args.batch_size,shuffle=True,
num_workers=2,drop_last=True)
if args.mode == 'semi':
trainset_u = Corrosion(home_dir,args.dataset_dir,img_transform=Compose(imgtr),
label_transform=Compose(labtr),co_transform=Compose(cotr),
split=args.split,labeled=False)
trainloader_u = DataLoader(trainset_u,batch_size=args.batch_size,shuffle=True,
num_workers=2,drop_last=True)
#########################
# Validation Dataloader #
########################
if args.val_orig:
if args.no_norm:
imgtr = [ZeroPadding(),ToTensor()]
else:
imgtr = [ZeroPadding(),ToTensor(),NormalizeOwn()]
labtr = [IgnoreLabelClass(),ToTensorLabel()]
# labtr = [IgnoreLabelClass(),ToTensorLabel(tensor_type=torch.FloatTensor)]
cotr = []
else:
if args.no_norm:
imgtr = [ToTensor()]
else:
imgtr = [ToTensor(),NormalizeOwn()]
labtr = [IgnoreLabelClass(),ToTensorLabel()]
# labtr = [IgnoreLabelClass(),ToTensorLabel(tensor_type=torch.FloatTensor)]
# cotr = [RandomSizedCrop3((320,320))] # (321,321)
cotr = [RandomSizedCrop3((320,320))]
valset = Corrosion(home_dir,args.dataset_dir,img_transform=Compose(imgtr), \
label_transform = Compose(labtr),co_transform=Compose(cotr),train_phase=False)
valoader = DataLoader(valset,batch_size=1)
#############
# GENERATOR #
#############
# generator = deeplabv2.ResDeeplab()
# softmax generator: in_chs=3, out_chs=2
generator = unet.AttU_Net()
# model_summary = generator.cuda()
init_weights(generator,args.init_net)
if args.init_net != 'unet':
optimG = optim.SGD(filter(lambda p: p.requires_grad, \
generator.parameters()),lr=args.g_lr,momentum=0.9,\
weight_decay=0.0001,nesterov=True)
else:
optimG = optim.Adam(filter(lambda p: p.requires_grad, \
generator.parameters()),args.g_lr, [0.5, 0.999])
if not args.nogpu:
generator = nn.DataParallel(generator).cuda()
#################
# DISCRIMINATOR #
################
if args.mode != "base":
# softmax generator
discriminator = Dis(in_channels=2)
# model_summary = discriminator.cuda()
# summary(model_summary, (2, 320, 320))
if args.d_optim == 'adam':
optimD = optim.Adam(filter(lambda p: p.requires_grad, \
discriminator.parameters()),lr = args.d_lr,weight_decay=0.0001)
else:
optimD = optim.SGD(filter(lambda p: p.requires_grad, \
discriminator.parameters()),lr=args.d_lr,weight_decay=0.0001,momentum=0.5,nesterov=True)
if not args.nogpu:
discriminator = nn.DataParallel(discriminator).cuda()
if args.mode == 'base':
train_base(generator,optimG,trainloader_l,valoader,args)
elif args.mode == 'adv':
train_adv(generator,discriminator,optimG,optimD,trainloader_l,valoader,args)
elif args.mode == 'semi':
train_semi(generator,discriminator,optimG,optimD,trainloader_l,trainloader_u,valoader,args)
else:
# train_semir(generator,discriminator,optimG,optimD,trainloader_l,valoader,args)
print("training mode incorrect")
def main():
args = parse_args()
CUR_DIR = os.getcwd()
with open(osp.join(CUR_DIR, "utils/config_crf.yaml")) as f:
CRF_CONFIG = Dict(yaml.safe_load(f))
random.seed(0)
torch.manual_seed(0)
if not args.nogpu:
torch.cuda.manual_seed_all(0)
if args.no_norm:
imgtr = [ToTensor()]
else:
imgtr = [ToTensor(),NormalizeOwn()]
# softmax
labtr = [IgnoreLabelClass(),ToTensorLabel()]
# labtr = [IgnoreLabelClass(),ToTensorLabel(tensor_type=torch.FloatTensor)]
# cotr = [RandomSizedCrop((320,320))] # (321,321)
cotr = [RandomSizedCrop3((320,320))]
print("dataset_dir: ", args.dataset_dir)
if args.mode == 'semi':
split_ratio = 0.8
else:
split_ratio = 1.0
trainset_l = Corrosion(home_dir,args.dataset_dir,img_transform=Compose(imgtr),
label_transform=Compose(labtr),co_transform=Compose(cotr),
split=split_ratio,labeled=True)
trainloader_l = DataLoader(trainset_l,batch_size=args.batch_size,shuffle=True,
num_workers=2,drop_last=True)
if args.mode == 'semi':
trainset_u = Corrosion(home_dir,args.dataset_dir,img_transform=Compose(imgtr),
label_transform=Compose(labtr),co_transform=Compose(cotr),
split=split_ratio,labeled=False)
trainloader_u = DataLoader(trainset_u,batch_size=args.batch_size,shuffle=True,
num_workers=2,drop_last=True)
postprocessor = DenseCRF(
iter_max=CRF_CONFIG.CRF.ITER_MAX,
pos_xy_std=CRF_CONFIG.CRF.POS_XY_STD,
pos_w=CRF_CONFIG.CRF.POS_W,
bi_xy_std=CRF_CONFIG.CRF.BI_XY_STD,
bi_rgb_std=CRF_CONFIG.CRF.BI_RGB_STD,
bi_w=CRF_CONFIG.CRF.BI_W,
)
#########################
# Validation Dataloader #
########################
if args.val_orig:
if args.no_norm:
imgtr = [ZeroPadding(),ToTensor()]
else:
imgtr = [ZeroPadding(),ToTensor(),NormalizeOwn()]
labtr = [IgnoreLabelClass(),ToTensorLabel()]
# labtr = [IgnoreLabelClass(),ToTensorLabel(tensor_type=torch.FloatTensor)]
cotr = []
else:
if args.no_norm:
imgtr = [ToTensor()]
else:
imgtr = [ToTensor(),NormalizeOwn()]
labtr = [IgnoreLabelClass(),ToTensorLabel()]
# labtr = [IgnoreLabelClass(),ToTensorLabel(tensor_type=torch.FloatTensor)]
# cotr = [RandomSizedCrop3((320,320))] # (321,321)
cotr = [RandomSizedCrop3((320,320))]
valset = Corrosion(home_dir,args.dataset_dir,img_transform=Compose(imgtr), \
label_transform = Compose(labtr),co_transform=Compose(cotr),train_phase=False)
valoader = DataLoader(valset,batch_size=1)
#############
# GENERATOR #
#############
# generator = deeplabv2.ResDeeplab()
# softmax generator: in_chs=3, out_chs=2
generator = unet.AttU_Net()
# model_summary = generator.cuda()
init_weights(generator,args.init_net)
if args.init_net != 'unet':
optimG = optim.SGD(filter(lambda p: p.requires_grad, \
generator.parameters()),lr=args.g_lr,momentum=0.9,\
weight_decay=0.0001,nesterov=True)
else:
optimG = optim.Adam(filter(lambda p: p.requires_grad, \
generator.parameters()),args.g_lr, [0.9, 0.999])
if not args.nogpu:
generator = nn.DataParallel(generator).cuda()
#################
# DISCRIMINATOR #
################
if args.mode != "base":
# softmax generator
discriminator = DisSigmoid(in_channels=2)
init_weights(generator,args.init_net)
# model_summary = discriminator.cuda()
# summary(model_summary, (2, 320, 320))
if args.d_optim == 'adam':
optimD = optim.Adam(filter(lambda p: p.requires_grad, \
discriminator.parameters()),args.d_lr,[0.9,0.999])
# discriminator.parameters()),[0.9,0.999],lr = args.d_lr,weight_decay=0.0001)
else:
optimD = optim.SGD(filter(lambda p: p.requires_grad, \
discriminator.parameters()),lr=args.d_lr,weight_decay=0.0001,momentum=0.9,nesterov=True)
if not args.nogpu:
discriminator = nn.DataParallel(discriminator).cuda()
if args.mode == 'base':
train_base(generator,optimG,trainloader_l,valoader,args)
elif args.mode == 'adv':
train_adv(generator,discriminator,optimG,optimD,trainloader_l,valoader,postprocessor,args)
elif args.mode == 'semi':
train_semi(generator,discriminator,optimG,optimD,trainloader_l,trainloader_u,valoader,args)
else:
# train_semir(generator,discriminator,optimG,optimD,trainloader_l,valoader,args)
print("training mode incorrect")
def evaluate_discriminator():
home_dir = os.path.dirname(os.path.realpath(__file__))
parser = argparse.ArgumentParser()
parser.add_argument("dataset_dir",
help="A directory containing img (Images) \
and cls (GT Segmentation) folder")
parser.add_argument("snapshot_g",
help="Snapshot with the saved generator model")
parser.add_argument("snapshot_d",
help="Snapshot with the saved discriminator model")
parser.add_argument("--val_orig",
help="Do Inference on original size image.\
Otherwise, crop to 320x320 like in training ",
action='store_true')
parser.add_argument("--norm",help="Normalize the test images",\
action='store_true')
args = parser.parse_args()
# print(args.val_orig, args.norm)
if args.val_orig:
img_transform = transforms.Compose([ToTensor()])
if args.norm:
img_transform = transforms.Compose(
[ToTensor(), NormalizeOwn(dataset='corrosion')])
label_transform = transforms.Compose(
[IgnoreLabelClass(), ToTensorLabel()])
# co_transform = transforms.Compose([RandomSizedCrop((320,320))])
co_transform = transforms.Compose([ResizedImage3((320, 320))])
testset = Corrosion(home_dir, args.dataset_dir,img_transform=img_transform, \
label_transform = label_transform,co_transform=co_transform,train_phase=False)
testloader = DataLoader(testset, batch_size=1)
else:
img_transform = transforms.Compose([ZeroPadding(), ToTensor()])
if args.norm:
img_transform = img_transform = transforms.Compose(
[ZeroPadding(),
ToTensor(),
NormalizeOwn(dataset='corrosion')])
label_transform = transforms.Compose(
[IgnoreLabelClass(), ToTensorLabel()])
testset = Corrosion(home_dir,args.dataset_dir,img_transform=img_transform, \
label_transform=label_transform,train_phase=False)
testloader = DataLoader(testset, batch_size=1)
# generator = deeplabv2.ResDeeplab()
# generatro = fcn.FCN8s_soft()
generator = unet.AttU_Net()
print(args.snapshot_g)
assert (os.path.isfile(args.snapshot_g))
snapshot_g = torch.load(args.snapshot_g)
discriminator = Dis(in_channels=2)
print(args.snapshot_d)
assert (os.path.isfile(args.snapshot_d))
snapshot_d = torch.load(args.snapshot_d)
saved_net = {
k.partition('module.')[2]: v
for i, (k, v) in enumerate(snapshot_g['state_dict'].items())
}
print('Generator Snapshot Loaded')
generator.load_state_dict(saved_net)
generator.eval()
generator = nn.DataParallel(generator).cuda()
print('Generator Loaded')
saved_net_d = {
k.partition('module.')[2]: v
for i, (k, v) in enumerate(snapshot_d['state_dict'].items())
}
print('Discriminator Snapshot Loaded')
discriminator.load_state_dict(saved_net_d)
discriminator.eval()
discriminator = nn.DataParallel(discriminator).cuda()
print('discriminator Loaded')
n_classes = 2
gts, preds = [], []
print('Prediction Goint to Start')
colorize = VOCColorize()
palette = make_palette(2)
# print(palette)
IMG_DIR = osp.join(args.dataset_dir, 'corrosion/JPEGImages')
# TODO: Crop out the padding before prediction
for img_id, (img, gt_mask, _, gte_mask, name) in enumerate(testloader):
print("Generating Predictions for Image {}".format(img_id))
gt_mask = gt_mask.numpy()[0]
img = Variable(img.cuda())
# img.cpu().numpy()[0]
img_path = osp.join(IMG_DIR, name[0] + '.jpg')
print(img_path)
img_array = cv2.imread(img_path)
img_array = cv2.resize(img_array, (320, 320),
interpolation=cv2.INTER_AREA)
img_array = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
out_pred_map = generator(img)
# print(out_pred_map.size())
# Get hard prediction
soft_pred = out_pred_map.data.cpu().numpy()[0]
# print("gen: ", soft_pred.shape)
# print(soft_pred.shape)
soft_pred = soft_pred[:, :gt_mask.shape[0], :gt_mask.shape[1]]
# print("gen: ", soft_pred.shape)
# print(soft_pred.shape)
hard_pred = np.argmax(soft_pred, axis=0).astype(np.uint8)
# print("gen: ", hard_pred.shape)
# Get discriminator prediction
dis_conf = discriminator(out_pred_map)
dis_confsmax = nn.Softmax2d()(dis_conf)
# print(dis_conf.size())
dis_soft_pred = dis_confsmax.data.cpu().numpy()[0]
# dis_soft_pred[dis_soft_pred<=0.2] = 0
# dis_soft_pred[dis_soft_pred>0.2] = 1
# print("dis: ", dis_soft_pred.shape)
dis_hard_pred = np.argmax(dis_soft_pred, axis=0).astype(np.uint8)
# print("dis: ", dis_hard_pred.shape)
# dis_pred = dis_pred[:,:gt_mask.shape[0],:gt_mask.shape[1]]
# print(soft_pred.shape)
# dis_hard_pred = np.argmax(dis_pred,axis=0).astype(np.uint8)
# print(hard_pred.shape, name)
output = np.asarray(hard_pred, dtype=np.int)
# print("gen: ", output.shape)
filename = os.path.join('results', '{}.png'.format(name[0]))
color_file = Image.fromarray(
colorize(output).transpose(1, 2, 0), 'RGB')
color_file.save(filename)
masked_im = Image.fromarray(vis_seg(img_array, output, palette))
masked_im.save(filename[0:-4] + '_vis.png')
# discriminator output
dis_output = np.asarray(dis_hard_pred, dtype=np.int)
# print("dis: ", dis_output.shape)
dis_filename = os.path.join('results',
'{}_dis.png'.format(name[0][0:-4]))
dis_color_file = Image.fromarray(
colorize(dis_output).transpose(1, 2, 0), 'RGB')
dis_color_file.save(dis_filename)
for gt_, pred_ in zip(gt_mask, hard_pred):
gts.append(gt_)
preds.append(pred_)
# input('s')
score, class_iou = scores(gts, preds, n_class=n_classes)
print("Mean IoU: {}".format(score))
def main():
args = parse_args()
random.seed(0)
torch.manual_seed(0)
if not args.nogpu:
torch.cuda.manual_seed_all(0)
if args.no_norm:
imgtr = [ToTensor()]
else:
imgtr = [ToTensor(), NormalizeOwn()]
if len(args.lr_step) != 0:
steps = list(map(lambda x: int(x), args.lr_step.split(',')))
# softmax
labtr = [IgnoreLabelClass(), ToTensorLabel()]
cotr = [RandomSizedCrop4((512, 512))]
print("dataset_dir: ", args.dataset_dir)
trainset_l = BoxSet(home_dir,
args.dataset_dir,
img_transform=Compose(imgtr),
label_transform=Compose(labtr),
co_transform=Compose(cotr),
split=args.split,
labeled=True,
label_correction=True)
trainloader_l = DataLoader(trainset_l,
batch_size=args.batch_size,
shuffle=True,
num_workers=2,
drop_last=True)
if args.split != 1:
trainset_u = BoxSet(home_dir,
args.dataset_dir,
img_transform=Compose(imgtr),
label_transform=Compose(labtr),
co_transform=Compose(cotr),
split=args.split,
labeled=False,
label_correction=True)
trainloader_u = DataLoader(trainset_l,
batch_size=args.batch_size,
shuffle=True,
num_workers=2,
drop_last=True)
#########################
# Validation Dataloader #
########################
if args.val_orig:
if args.no_norm:
imgtr = [ZeroPadding(), ToTensor()]
else:
imgtr = [ZeroPadding(), ToTensor(), NormalizeOwn()]
# softmax
labtr = [IgnoreLabelClass(), ToTensorLabel()]
cotr = []
else:
if args.no_norm:
imgtr = [ToTensor()]
else:
imgtr = [ToTensor(), NormalizeOwn()]
# softmax
labtr = [IgnoreLabelClass(), ToTensorLabel()]
cotr = [ResizedImage4((512, 512))]
valset = BoxSet(home_dir,args.dataset_dir,img_transform=Compose(imgtr), \
label_transform = Compose(labtr),co_transform=Compose(cotr),train_phase=False)
valoader = DataLoader(valset, batch_size=1)
#############
# GENERATOR #
#############
generator = unet.AttU_Net(output_ch=7, Centroids=False)
if osp.isfile(args.snapshot):
print("load checkpoint => ", args.snapshot)
checkpoint = torch.load(args.snapshot)
generator_dict = generator.state_dict()
saved_net = {
k.partition('module.')[2]: v
for i, (k, v) in enumerate(checkpoint['state_dict'].items())
if k.partition('module.')[2] in generator_dict
}
generator_dict.update(saved_net)
generator.load_state_dict(saved_net)
else:
init_weights(generator, args.init_net)
if args.init_net != 'unet':
optimG = optim.Adam(filter(lambda p: p.requires_grad, \
generator.parameters()),args.g_lr, [0.5, 0.999])
else:
optimG = optim.Adam(filter(lambda p: p.requires_grad, \
generator.parameters()),args.g_lr, [0.5, 0.999])
"""
optimG = optim.SGD(filter(lambda p: p.requires_grad, \
generator.parameters()),lr=args.g_lr,momentum=0.9,\
weight_decay=0.0001,nesterov=True)
"""
if not args.nogpu:
generator = nn.DataParallel(generator).cuda()
if args.mode == 'base':
train_base(generator, optimG, trainloader_l, valoader, args)
elif args.mode == 'label_correction':
train_box_cluster(generator, steps, optimG, trainloader_l, valoader,
args)
else:
print("training mode incorrect")