def crf_inference(opts, dataset, metrics):
metrics.reset()
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
postprocessor = DenseCRF(
iter_max=10,
pos_xy_std=1,
pos_w=3,
bi_xy_std=67,
bi_rgb_std=3,
bi_w=4,
)
def process(i):
image, gt_label, fname = dataset.__getitem__(i)
filename = os.path.join(opts.logit_dir, fname + ".npy")
logit = np.load(filename)
_, H, W = image.shape
logit = torch.FloatTensor(logit)[None, ...]
logit = F.interpolate(logit,
size=(H, W),
mode="bilinear",
align_corners=False)
prob = F.softmax(logit, dim=1)[0].numpy()
gt_label = gt_label.cpu().numpy()
image = image.permute(1, 2, 0).cpu().numpy()
image *= std
image += mean
image *= 255
image = image.astype(np.uint8)
prob = postprocessor(image, prob)
pred_label = np.argmax(prob, axis=0)
return pred_label, gt_label
# CRF in multi-process
results = joblib.Parallel(n_jobs=multiprocessing.cpu_count(),
verbose=10,
pre_dispatch="all")([
joblib.delayed(process)(i)
for i in range(len(dataset))
])
preds, gts = zip(*results)
for pred, gt in zip(preds, gts):
metrics.update(gt, pred)
score = metrics.get_results()
return score
def img_crf(img, softmax_pred):
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,
)
img = img.astype(np.uint8)
crf_output = postprocessor(img, softmax_pred)
return crf_output
def get_roi_crf(img, softmax_pred, hard_pred):
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,
)
contours, _ = cv2.findContours(hard_pred.copy(), cv2.RETR_CCOMP,
cv2.CHAIN_APPROX_NONE)
rois = []
rois_crf = []
if len(contours) > 1:
for c in contours:
x, y, w, h = cv2.boundingRect(c)
if w * h <= 40 * 40:
continue
else:
rect = get_rect_extend([x, y, x + w, y + h], img, 50)
xmin, ymin, xmax, ymax = rect[0], rect[1], rect[2], rect[3]
prob_map = softmax_pred[:, ymin:ymax, xmin:xmax]
roi = img[ymin:ymax, xmin:xmax, :]
roi = roi.astype(np.uint8)
roi_crf_output = postprocessor(roi, prob_map)
rois.append([xmin, ymin, xmax, ymax])
rois_crf.append(roi_crf_output)
elif len(contours) == 0:
rois.append([0, 0, img.shape[0], img.shape[1]])
rois_crf.append(softmax_pred)
else:
rois = []
x, y, w, h = cv2.boundingRect(contours[0])
rect = get_rect_extend([x, y, x + w, y + h], img, 80)
# rois.append([0, 0, img.shape[0], img.shape[1]])
xmin, ymin, xmax, ymax = rect[0], rect[1], rect[2], rect[3]
prob_map = softmax_pred[:, ymin:ymax, xmin:xmax]
roi = img[ymin:ymax, xmin:xmax, :]
roi = roi.astype(np.uint8)
roi_crf_output = postprocessor(roi, prob_map)
rois_crf.append(roi_crf_output)
rois.append(rect)
# rois_crf = img_crf(img, softmax_pred)
return rois_crf, rois
def roi_crf(img, softmax_pred, rois):
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,
)
rois_infer = []
for roi in rois:
xmin, ymin, xmax, ymax = roi[0], roi[1], roi[2], roi[3]
prob_map = softmax_pred[:, ymin:ymax, xmin:xmax]
roi = img[ymin:ymax, xmin:xmax, :]
roi = roi.astype(np.uint8)
roi_crf_output = postprocessor(roi, prob_map)
rois_infer.append(roi_crf_output)
return rois_infer # np.array(rois_infer)
def crf(args):
print('CRF post-processing...')
torch.set_grad_enabled(False)
if args.model_type == 'vgg16':
with open('./configs/vgg16.yaml', 'r') as f:
CONFIG = Dict(yaml.load(f, Loader=yaml.FullLoader))
elif args.model_type == 'resnet101':
with open('./configs/resnet101.yaml', 'r') as f:
CONFIG = Dict(yaml.load(f, Loader=yaml.FullLoader))
else:
print('Unsupported model structure!')
return
crop_size = CONFIG.IMAGE.SIZE.TEST
# CRF post-processor
postprocessor = DenseCRF(
iter_max=CONFIG.CRF.ITER_MAX,
pos_xy_std=CONFIG.CRF.POS_XY_STD,
pos_w=CONFIG.CRF.POS_W,
bi_xy_std=CONFIG.CRF.BI_XY_STD,
bi_rgb_std=CONFIG.CRF.BI_RGB_STD,
bi_w=CONFIG.CRF.BI_W,
)
# VOC color
palette = []
for i in range(256):
palette.extend((i, i, i))
palette[:3 * 21] = np.array(
[[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0], [0, 0, 128],
[128, 0, 128], [0, 128, 128], [128, 128, 128],
[64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0], [64, 0, 128],
[192, 0, 128], [64, 128, 128], [192, 128, 128], [0, 64, 0],
[128, 64, 0], [0, 192, 0], [128, 192, 0], [0, 64, 128]],
dtype='uint8').flatten()
logitPath = os.path.join(CONFIG.MODEL.NAME, 'logitPath')
predictPath = os.path.join(CONFIG.MODEL.NAME, 'prediction')
if not os.path.exists(predictPath):
os.makedirs(predictPath)
id_path = os.path.join('./datalists', 'val.txt')
image_ids = [i.strip() for i in open(id_path) if not i.strip() == ' ']
# Process per sample
def process(i):
image_id = image_ids[i]
gtPath = os.path.join(args.root_path, 'SegmentationClass',
image_id + '.png')
gt_label = Image.open(gtPath)
imgPath = os.path.join(args.root_path, 'JPEGImages', image_id + '.jpg')
raw_image = cv2.imread(imgPath, cv2.IMREAD_COLOR) # shape = [H, W, 3]
H, W, _ = raw_image.shape
if args.model_type == 'vgg16':
## padding raw_image to 513X513
pad_h = max(crop_size - H, 0)
pad_w = max(crop_size - W, 0)
pad_kwargs = {
"top": 0,
"bottom": pad_h,
"left": 0,
"right": pad_w,
"borderType": cv2.BORDER_CONSTANT,
}
raw_image = cv2.copyMakeBorder(raw_image,
value=[0, 0, 0],
**pad_kwargs)
raw_image = raw_image.astype(np.uint8)
H, W, _ = raw_image.shape
filename = os.path.join(logitPath, image_id + ".npy")
logit = np.load(filename)
logit = torch.FloatTensor(logit)[None, ...]
logit = F.interpolate(logit,
size=(H, W),
mode="bilinear",
align_corners=True)
logit = F.softmax(logit, dim=1)[0].numpy()
prob = postprocessor(raw_image, logit)
label = np.argmax(prob, axis=0)
# save predictions to path
w, h = gt_label.size[0], gt_label.size[1]
img_label = Image.fromarray(label.astype(np.uint8))
img_label = img_label.crop((0, 0, w, h))
label = np.asarray(img_label)
gt_label = np.asarray(gt_label)
img_label.putpalette(palette)
img_label.save(os.path.join(predictPath, image_id + '.png'))
return label, gt_label
# CRF in multi-process
results = joblib.Parallel(n_jobs=args.n_jobs,
verbose=10,
pre_dispatch="all")([
joblib.delayed(process)(i)
for i in range(len(image_ids))
])
preds, gts = zip(*results)
# Pixel Accuracy, Mean Accuracy, Class IoU, Mean IoU, Freq Weighted IoU
score = scores(gts, preds, n_class=CONFIG.DATASET.N_CLASSES)
print(score)
savePath = os.path.join(CONFIG.MODEL.NAME, 'results')
if not os.path.exists(savePath):
os.makedirs(savePath)
with open(savePath + '/results.json', "w") as f:
json.dump(score, f, indent=4, sort_keys=True)
def test(self):
epoch = 1
self.model.eval()
self.evaluator.reset()
tbar = tqdm(self.val_loader, desc='\r')
test_loss = 0.0
overallScore = 0.0
overallScore_1 = 0.0
# w1 = POS_W = [3:6]
# w2 = BI_W = 3
# Sa = BI_XY_STD = [30:100]
# Sb = BI_RGB_STD = [3:6]
# Sg = POS_XY_STD = 3
w1 = 3
w2 = 3
Sa = 30
Sb = 3
Sg = 3
postprocess = DenseCRF(iter_max=10,
pos_w=w1,
bi_w=w2,
bi_xy_std=Sa,
bi_rgb_std=Sb,
pos_xy_std=Sg)
start = time.time()
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
_, _, H, W = image.cpu().numpy().shape
if self.args.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
for j, (img, logit,
gt_label) in enumerate(zip(image, output, target)):
filename = os.path.join("output/", str(j) + ".npy")
# Pixel Labeling
_, H, W = logit.shape
original = logit.cpu().numpy()
img = img.cpu().numpy()
gt_label = gt_label.cpu().numpy()
logit = torch.FloatTensor(logit.cpu().numpy())[None, ...]
logit = F.interpolate(logit,
size=(H, W),
mode="bilinear",
align_corners=False)
prob = F.softmax(logit, dim=1)[0].cpu().numpy()
img = img.astype(np.uint8).transpose(1, 2, 0)
prob = postprocess(img, prob)
label = np.argmax(prob, axis=0)
self.evaluatorCRF.add_batch(gt_label, label)
score = scores(gt_label, label, n_class=self.nclass)
overallScore += score['Mean IoU']
mIoU_CRF = self.evaluatorCRF.Mean_Intersection_over_Union()
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
self.evaluator.add_batch(target, pred)
mIoU = self.evaluator.Mean_Intersection_over_Union()
mIoU = self.evaluator.Mean_Intersection_over_Union()
print("w1 = ", w1)
print("w2 = ", w2)
print("Sa = ", Sa)
print("Sb = ", Sb)
print("Sg = ", Sg)
print("Final-PRE -CRF =" + str(mIoU))
print("Final-POST-CRF =" + str(mIoU_CRF))
end = time.time()
def test_save(self):
epoch = 1
self.model.eval()
self.evaluator.reset()
tbar = tqdm(self.test_loader, desc='\r')
w1 = 3
w2 = 3
Sa = 30
Sb = 3
Sg = 3
postprocess = DenseCRF(iter_max=10,
pos_w=w1,
bi_w=w2,
bi_xy_std=Sa,
bi_rgb_std=Sb,
pos_xy_std=Sg)
for i, sample in enumerate(tbar):
image, image_path = sample['image'], sample['path']
_, _, H, W = image.cpu().numpy().shape
if self.args.cuda:
image = image.cuda()
with torch.no_grad():
output = self.model(image)
for j, (img, logit,
imPath) in enumerate(zip(image, output, image_path)):
filename = os.path.join("output/", str(j) + ".npy")
_, H, W = logit.shape
original = logit.cpu().numpy()
img = img.cpu().numpy()
logit = torch.FloatTensor(logit.cpu().numpy())[None, ...]
logit = F.interpolate(logit,
size=(H, W),
mode="bilinear",
align_corners=False)
prob = F.softmax(logit, dim=1)[0].cpu().numpy()
img = img.astype(np.uint8).transpose(1, 2, 0)
prob = postprocess(img, prob)
label = np.argmax(prob, axis=0)
label[label == 21] = 255
out_image = Image.fromarray(label.squeeze().astype('uint8'))
out_image.putpalette(self.colorMap)
if self.args.dataset == 'pascal':
out_image.save(
"output/test/results/VOC2012/Segmentation/comp5_test_cls/"
+ imPath[36:-4] + ".png")
elif self.args.dataset == 'cityscapes':
out_image.save("output/Cityscapes/test/" + imPath[34:-4] +
".png")
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")