def validate(args):
# Setup Dataloader
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
loader = data_loader(data_path, split=args.split, is_transform=True, img_size=(args.img_rows, args.img_cols))
n_classes = loader.n_classes
valloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=4)
running_metrics = runningScore(n_classes)
# Setup Model
model = get_model(args.model_path[:args.model_path.find('_')], n_classes)
state = convert_state_dict(torch.load(args.model_path)['model_state'])
model.load_state_dict(state)
model.eval()
for i, (images, labels) in tqdm(enumerate(valloader)):
model.cuda()
images = Variable(images.cuda(), volatile=True)
labels = Variable(labels.cuda(), volatile=True)
outputs = model(images)
pred = outputs.data.max(1)[1].cpu().numpy()
gt = labels.data.cpu().numpy()
running_metrics.update(gt, pred)
score, class_iou = running_metrics.get_scores()
for k, v in score.items():
print(k, v)
for i in range(n_classes):
print(i, class_iou[i])
def train(args):
# Setup Dataloader
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset, config_file=args.config_file)
loader = data_loader(data_path,
is_transform=True,
img_size=(args.img_rows, args.img_cols))
n_classes = loader.n_classes
# must use 1 worker for AWS sagemaker without ipc="host" or larger shared memory size
trainloader = data.DataLoader(loader,
batch_size=args.batch_size,
num_workers=1,
shuffle=True)
# Setup Model
model = get_model(args.arch, n_classes)
# Setup log dir / logging
if not os.path.exists(args.log_dir):
os.makedirs(args.log_dir)
configure(args.log_dir)
model = torch.nn.DataParallel(model,
device_ids=range(torch.cuda.device_count()))
model.cuda()
optimizer = torch.optim.SGD(model.parameters(),
lr=args.l_rate,
momentum=0.9,
weight_decay=5e-4)
step = 0
for epoch in range(args.n_epoch):
start_time = time.time()
for i, (images, labels) in enumerate(trainloader):
images = Variable(images.cuda())
labels = Variable(labels.cuda())
optimizer.zero_grad()
outputs = model(images)
loss = cross_entropy2d(outputs, labels)
loss.backward()
optimizer.step()
log_value('Loss', loss.data[0], step)
step += 1
if (i + 1) % 20 == 0:
print("Epoch [%d/%d] Loss: %.4f" %
(epoch + 1, args.n_epoch, loss.data[0]),
flush=True)
end_time = time.time()
print('Epoch run time: %s' % (end_time - start_time))
torch.save(
model, args.log_dir + "{}_{}_{}_{}.pt".format(
args.arch, args.dataset, args.feature_scale, epoch))
def get_dataset_loader():
args=edict()
args.dataset_name = 'cityscapes'
args.config_path = os.path.join('/home/yzbx/git/gnu/pytorch-semseg', 'config.json')
args.img_rows=224
args.img_cols=224
args.img_norm=True
args.batch_size=2
data_loader = get_loader(args.dataset_name)
data_path = get_data_path(args.dataset_name, args.config_path)
t_loader = data_loader(data_path, is_transform=True, split='train', img_size=(
args.img_rows, args.img_cols), augmentations=None, img_norm=args.img_norm)
v_loader = data_loader(data_path, is_transform=True, split='val', img_size=(
args.img_rows, args.img_cols), img_norm=args.img_norm)
n_classes = t_loader.n_classes
print('class number is',n_classes)
trainloader = data.DataLoader(
t_loader, batch_size=args.batch_size, num_workers=8, shuffle=True)
valloader = data.DataLoader(
v_loader, batch_size=args.batch_size, num_workers=8)
return trainloader,valloader
def test(args, cfg):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_file_name = os.path.split(args.model_path)[1]
model_name = model_file_name[:model_file_name.find("_")]
IMG_Path = Path(args.img_path)
IMG_File = natsort.natsorted(list(IMG_Path.glob("*.png")),
alg=natsort.PATH)
IMG_Str = []
for i in IMG_File:
IMG_Str.append(str(i))
# Setup image
print("Read Input Image from : {}".format(args.img_path))
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset, config_file=cfg)
loader = data_loader(data_path, is_transform=True, img_norm=args.img_norm)
n_classes = loader.n_classes
# Setup Model
model = get_model(cfg['model'], n_classes)
state = convert_state_dict(torch.load(args.model_path)["model_state"])
# state=torch.load(args.model_path)["model_state"]
model.load_state_dict(state)
model.eval()
model.to(device)
for j in tqdm(range(len(IMG_Str))):
img_path = IMG_Str[j]
img = misc.imread(img_path)
# img = img[:, :, ::-1]
img = img.astype(np.float64)
# img -= loader.mean
if args.img_norm:
img = img.astype(float) / 255.0
# NHWC -> NCHW
img = img.transpose(2, 0, 1)
img = np.expand_dims(img, 0)
img = torch.from_numpy(img).float()
images = img.to(device)
outputs = model(images)
outputs_probability = F.softmax(outputs)
data = outputs_probability.data
data_max = data.max(1)
prob = data_max[0]
prob_img_format = np.squeeze(prob.cpu().numpy(), axis=0)
avg_prob = np.mean(prob_img_format)
print("Confidence Score for %s: \n%f" % (img_path, avg_prob))
pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=0)
decoded = loader.decode_segmap(pred)
out_path = "test_out/test_confidence/out/" + Path(img_path).name
decoded_bgr = cv.cvtColor(decoded, cv.COLOR_RGB2BGR)
# misc.imsave(out_path, decoded)
cv.imwrite(out_path, decoded_bgr)
def test(args, cfg):
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_file_name = os.path.split(args.model_path)[1]
model_name = model_file_name[:model_file_name.find("_")]
IMG_Path = Path(args.img_path)
IMG_File = natsort.natsorted(list(IMG_Path.glob("*.tif")),
alg=natsort.PATH)
IMG_Str = []
for i in IMG_File:
IMG_Str.append(str(i))
# Setup image
print("Read Input Image from : {}".format(args.img_path))
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset, config_file=cfg)
loader = data_loader(data_path, is_transform=True, img_norm=args.img_norm)
n_classes = loader.n_classes
v_loader = data_loader(
data_path,
is_transform=True,
split=cfg['data']['val_split'],
img_size=(cfg['data']['img_rows'], cfg['data']['img_cols']),
)
valloader = data.DataLoader(v_loader,
batch_size=cfg['training']['batch_size'],
num_workers=cfg['training']['n_workers'])
# Setup Model
model = get_model(cfg['model'], n_classes)
state = convert_state_dict(torch.load(args.model_path)["model_state"])
# state=torch.load(args.model_path)["model_state"]
model.load_state_dict(state)
model.eval()
model.to(device)
with torch.no_grad():
for i_val, (img_path, images_val,
labels_val) in tqdm(enumerate(valloader)):
img_name = img_path[0]
images_val = images_val.to(device)
outputs = model(images_val)
pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=0)
decoded = loader.decode_segmap(pred)
out_path = "test_out/CAN_res50_4band_data07/" + Path(
img_name).stem + ".png"
decoded_bgr = cv.cvtColor(decoded, cv.COLOR_RGB2BGR)
# misc.imsave(out_path, decoded)
cv.imwrite(out_path, decoded_bgr)
def process_args_from_loaded_cfg(cfg, args):
"""
process arguments from loaded CFG.
:param cfg:
:param args:
:return:
"""
# If calling this, it will use CFG for sure.
# if 'runet' in cfg['model']['arch']:
# args.hidden_size = cfg['model']['hidden_size']
#
r = re.compile('-h\d+-')
r_d = re.compile('\d+')
_h = r.findall(cfg['logdir'])
if len(_h) > 0:
res = int(r_d.findall(_h[0])[0])
if res > 1: # avoid the represent h=init
args.hidden_size = res
else:
res = args.hidden_size
if cfg['model'].get('hidden_size'):
assert cfg['model']['hidden_size'] == res
else:
cfg['model']['hidden_size'] = res
args.gate = cfg['model'].get('gate') or args.gate
if args.is_recurrent is not None:
args.is_recurrent = cfg['training']['loss']['name'] \
in ['multi_step_cross_entropy']
out_path = args.out_path or os.path.join(
'results', cfg['data']['dataset'],
os.path.basename(os.path.dirname(cfg['logdir'])))
cfg['eval_out_path'] = out_path
# Process for unet_level
if cfg['model']['arch'] == "runet":
if cfg['model'].get('unet_level'):
args.unet_level = cfg['model']['unet_level']
else:
unet_level = args.hidden_size // 32
args.unet_level = unet_level
cfg['model']['unet_level'] = args.unet_level
cfg['model']['recurrent_level'] = args.recurrent_level
if not os.path.exists(cfg['data']['path']):
cfg['data']['path'] = get_data_path(
cfg['data']['path'],
config_file=f"configs/dataset/{cfg['data']['dataset'].replace('_', '')}.yml")
IPython.embed()
return cfg, args
def setup(self, pre_encode=False):
sbd_path = get_data_path('sbd')
voc_path = get_data_path('pascal')
target_path = self.root + '/SegmentationClass/pre_encoded/'
if not os.path.exists(target_path):
os.makedirs(target_path)
# Load SBD train set
sbd_train_list = tuple(open(sbd_path + 'dataset/train.txt', 'r'))
sbd_train_list = [id_.rstrip() for id_ in sbd_train_list]
# Load SBD val set
sbd_val_list = tuple(open(sbd_path + 'dataset/val.txt', 'r'))
sbd_val_list = [id_.rstrip() for id_ in sbd_val_list]
# Join everything
self.files[
'train_aug'] = self.files['train'] + sbd_train_list + sbd_val_list
# Remove duplicates and intersection with Pascal VOC validation set
self.files['train_aug'] = list(
set(self.files['train_aug']) - set(self.files['val']))
self.files['train_aug'].sort()
if pre_encode:
print("Pre-encoding segmentation masks...")
lbl_dir = os.path.join(sbd_path, 'dataset', 'cls')
lbl_list = [f for f in os.listdir(lbl_dir) if f.endswith('.mat')]
for i in tqdm(lbl_list):
lbl_path = os.path.join(lbl_dir, i)
lbl = io.loadmat(
lbl_path)['GTcls'][0]['Segmentation'][0].astype(np.int32)
lbl = m.toimage(lbl, high=lbl.max(), low=lbl.min())
m.imsave(target_path + os.path.splitext(i)[0] + '.png', lbl)
for i in tqdm(self.files['trainval']):
lbl_path = self.root + '/SegmentationClass/' + i + '.png'
lbl = self.encode_segmap(m.imread(lbl_path))
lbl = m.toimage(lbl, high=lbl.max(), low=lbl.min())
m.imsave(target_path + i + '.png', lbl)
def validate(args):
# Setup Dataloader
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
loader = data_loader(data_path,
split=args.split,
is_transform=True,
img_size=(args.img_rows, args.img_cols))
n_classes = loader.n_classes
valloader = data.DataLoader(loader,
batch_size=args.batch_size,
num_workers=4)
# Setup Model
model = torch.load(args.model_path)
model.eval()
if torch.cuda.is_available():
model.cuda(args.gpu)
gts, preds = [], []
for i, (images, labels) in tqdm(enumerate(valloader)):
if torch.cuda.is_available():
images = Variable(images.cuda(args.gpu))
labels = Variable(labels.cuda(args.gpu))
else:
images = Variable(images)
labels = Variable(labels)
outputs = model(images)
pred = outputs.data.max(1)[1].cpu().numpy()
pred = np.squeeze(pred)
pred = cv2.resize(pred,
labels.size()[1:][::-1],
interpolation=cv2.INTER_NEAREST)
pred = np.expand_dims(pred, axis=0)
gt = labels.data.cpu().numpy()
for gt_, pred_ in zip(gt, pred):
gts.append(gt_)
preds.append(pred_)
score, class_iou = scores(gts, preds, n_class=n_classes)
for k, v in score.items():
print k, v
for i in range(n_classes):
print i, class_iou[i]
def test(args):
model_file_name = os.path.split(args.model_path)[1]
model_name = model_file_name[:model_file_name.find('_')]
print("Building " + model_name)
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
loader = data_loader(data_path, phase='test')
im_paths = loader.im_paths()
n_classes = loader.n_classes
testloader = data.DataLoader(loader,
batch_size=1,
num_workers=1,
shuffle=False)
# Setup Model
model = get_model(model_name, n_classes)
state = torch.load(args.model_path)['model_state']
model.load_state_dict(state)
model.eval()
model.cuda()
# Run test for KITTI Road dataset
for i, (image, tr_image, lidar, tr_lidar) in enumerate(testloader):
im_name_splits = im_paths[i].split('/')[-1].split('.')[0].split('_')
task = im_name_splits[0]
print('processing %d-th image' % i)
t0 = time.time()
orig_h, orig_w = image.shape[1:3]
with torch.no_grad():
tr_image = Variable(tr_image.cuda())
tr_lidar = Variable(tr_lidar.cuda())
outputs = model([tr_image, tr_lidar])
outputs = outputs.cpu().numpy().transpose((2, 3, 1, 0)).squeeze()
outputs = cv2.resize(outputs, (orig_w, orig_h))
outputs = outputs[:, :, 1]
print('Time({:d}'.format(i) + ') {0:.3f}'.format(time.time() - t0))
output_fg = outputs * 255.
output_fg[output_fg > 255] = 255
output_fg = output_fg.astype(np.uint8)
cv2.imwrite(
'./outputs/results/' + im_name_splits[0] + '_road_' +
im_name_splits[1] + '.png', output_fg)
print('write to ./outputs/results/' + im_name_splits[0] + '_road_' +
im_name_splits[1] + '.png')
def validate():
# Setup Dataloader
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
loader = data_loader(data_path, img_size=args.img_size)
n_classes = loader.n_classes
n_channels = loader.n_channels
valloader = data.DataLoader(loader,
batch_size=args.batch_size,
num_workers=4,
shuffle=True)
# Setup Model
model = torch.load(args.model_path)
model.eval()
if torch.cuda.is_available():
model.cuda(0)
gts, preds = [], []
for i, (images, labels) in tqdm(enumerate(valloader)):
if i >= args.max_samples:
break
if torch.cuda.is_available():
images = Variable(images.cuda(0))
labels = Variable(labels.cuda(0))
else:
images = Variable(images)
labels = Variable(labels)
outputs = model(images)
pred = np.squeeze((torch.max(outputs.data, 1,
keepdim=True))[1].cpu().numpy())
gt = np.squeeze(labels.data.cpu().numpy())
for gt_, pred_ in zip(gt, pred):
gts.append(gt_)
preds.append(pred_)
score, class_iou = scores(gts, preds, n_class=n_classes)
for k, v in score.items():
print(k, v)
for i in range(n_classes):
print(i, class_iou[i])
def test(args):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_file_name = os.path.split(args.model_path)[1]
model_name = model_file_name[:model_file_name.find("_")]
# Setup image
print("Read Input Image from : {}".format(args.img_path))
img = misc.imread(args.img_path)
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
loader = data_loader(data_path, is_transform=True, img_norm=args.img_norm)
n_classes = loader.n_classes
img = img[:, :, ::-1]
img = img.astype(np.float64)
img -= loader.mean
if args.img_norm:
img = img.astype(float) / 255.0
# NHWC -> NCHW
img = img.transpose(2, 0, 1)
img = np.expand_dims(img, 0)
img = torch.from_numpy(img).float()
# Setup Model
model_dict = {"arch": model_name}
model = get_model(model_dict, n_classes, version=args.dataset)
state = convert_state_dict(torch.load(args.model_path)["model_state"])
model.load_state_dict(state)
model.eval()
model.to(device)
images = img.to(device)
outputs = model(images)
if args.mask_path:
print("Read Image Mask from : {}".format(args.mask_path))
mask = torch.load(args.mask_path)
mask = mask.to(device)
outputs = to_super_to_pixels(outputs, mask)
pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=0)
decoded = loader.decode_segmap(pred)
print("Classes found: ", np.unique(pred))
misc.imsave(args.out_path, decoded)
print("Segmentation Mask Saved at: {}".format(args.out_path))
def test(args):
# Setup image
print("Read Input Image from : ", args.img_path)
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
loader = data_loader(data_path, is_transform=True)
n_classes = loader.n_classes
img = img[:, :, ::-1]
img = img.astype(np.float64)
img -= loader.mean
# print(loader.img_size[0], loader.img_size[1])
# img = misc.imresize(img, (loader.img_size[0], loader.img_size[1]))
img = img.astype(float) / 255.0
# NHWC -> NCWH
img = img.transpose(2, 0, 1)
img = np.expand_dims(img, 0)
img = torch.from_numpy(img).float()
# Setup Model
model = torch.load(args.model_path)
model.eval()
if torch.cuda.is_available():
model.cuda(0)
images = Variable(img.cuda(0))
else:
images = Variable(img)
outputs = model(images)
log('the size of outputs is '.format(outputs.size()))
log('the size of outputs is ', outputs.data.max(1)[1].size())
pred = np.squeeze(outputs.data.cpu().numpy(), axis=0)
pred = pred[0, :, :] > pred[1, :, :]
pred = pred * 255
pred = outputs.data.cpu().numpy()
log('The prediction shape is {}'.format(pred.shape))
# print('the size of pred is ', pred.shape())
# fig = plt.figure()
# plt.imshow(pred[1,:,:])
# plt.title('testing result')
# plt.show()
pred = pred * 255
misc.imsave(args.out_path, pred[0, :, :])
print("Segmentation Mask Saved at: ", args.out_path)
def test(args):
# Setup image
print("Read Input Image from : {}".format(args.img_path))
orig_img = misc.imread(args.img_path)
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
loader = data_loader(data_path, is_transform=True)
n_classes = loader.n_classes
img = orig_img[:, :, ::-1]
img = img.astype(np.float64)
img -= loader.mean
img = misc.imresize(img, (loader.img_size[0], loader.img_size[1]))
img = img.astype(float) / 255.0
# NHWC -> NCWH
img = img.transpose(2, 0, 1)
img = np.expand_dims(img, 0)
img = torch.from_numpy(img).float()
# Setup Model
model = get_model(args.arch, n_classes)
model.load_state_dict(torch.load(args.model_path)['state_dict'])
model = torch.nn.DataParallel(model,
device_ids=range(
torch.cuda.device_count())).cuda()
model.eval()
if torch.cuda.is_available():
model.cuda(0)
images = Variable(img.cuda(0))
else:
images = Variable(img)
outputs = model(images)
pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=0)
decoded = loader.decode_segmap(pred)
if args.alpha_blend:
orig_img = misc.imresize(orig_img,
(loader.img_size[0], loader.img_size[1]))
out_img = ALPHA * orig_img + (1 - ALPHA) * decoded
else:
out_img = decoded
print(np.unique(pred))
misc.imsave(args.out_path, out_img)
print("Segmentation Mask Saved at: {}".format(args.out_path))
def validate(args):
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
loader = data_loader(data_path,
split=args.split,
is_transform=True,
img_size=None,
img_norm=args.img_norm)
n_classes = loader.n_classes
valloader = data.DataLoader(loader,
batch_size=args.batch_size,
num_workers=4)
running_metrics = runningScore(n_classes)
# Setup Model
from pytorchgo.model.deeplab_resnet import Res_Deeplab
model = Res_Deeplab(NoLabels=n_classes, pretrained=False)
state = torch.load("/home/hutao/MS_DeepLab_resnet_trained_VOC.pth")
model.load_state_dict(state)
model.eval()
model.cuda()
print "pra"
for i, (images, labels) in tqdm(enumerate(valloader), desc="validation"):
start_time = timeit.default_timer()
img_large = torch.Tensor(np.zeros((1, 3, 513, 513)))
img_large[:, :, :images.shape[2], :images.shape[3]] = images
output = model(Variable(img_large, volatile=True).cuda())
output = output.data.max(1)[1].cpu().numpy()
pred = output[:, :images.shape[2], :images.shape[3]]
gt = labels.numpy()
if args.measure_time:
elapsed_time = timeit.default_timer() - start_time
print('Inference time (iter {0:5d}): {1:3.5f} fps'.format(
i + 1, pred.shape[0] / elapsed_time))
running_metrics.update(gt, pred)
score, class_iou = running_metrics.get_scores()
for k, v in score.items():
print(k, v)
for i in range(n_classes):
print(i, class_iou[i])
def validate(args):
# Setup Dataloader
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
loader = data_loader(data_path,
split=args.split,
is_transform=True,
img_size=(args.img_rows, args.img_cols))
n_classes = loader.n_classes
valloader = data.DataLoader(loader,
batch_size=args.batch_size,
num_workers=4)
# Setup Model
model = torch.load(args.model_path)
model.eval()
gts, preds = [], []
for i, (images, labels) in tqdm(enumerate(valloader)):
if torch.cuda.is_available():
model = torch.nn.DataParallel(model,
device_ids=range(
torch.cuda.device_count()))
images = Variable(images.cuda(0))
labels = Variable(labels.cuda(0))
else:
images = Variable(images)
labels = Variable(labels)
outputs = model(images)
pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=1)
gt = labels.data.cpu().numpy()
for gt_, pred_ in zip(gt, pred):
gts.append(gt_)
preds.append(pred_)
score, class_iou = scores(gts, preds, n_class=n_classes)
for k, v in score.items():
print(k, v)
for i in range(n_classes):
print(i, class_iou[i])
def test(args):
args.img_path = '/home/shehabk/dataSets/CKPLUS/unet_256/A/test/S077_007_00000029.png'
args.dataset = 'ckplus'
args.out_path = '/home/shehabk/Desktop/' + os.path.basename(args.img_path)
# args.arch = 'segnet'
args.model_path = 'unet2_ckplus_1_99.pkl'
# Setup image
print("Read Input Image from : {}".format(args.img_path))
img = misc.imread(args.img_path,mode = 'RGB')
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
loader = data_loader(data_path, is_transform=True)
n_classes = loader.n_classes
img = img[:, :, ::-1]
img = img.astype(np.float64)
img -= loader.mean
img = misc.imresize(img, (loader.img_size[0], loader.img_size[1]))
img = img.astype(float) / 255.0
# NHWC -> NCWH
img = img.transpose(2, 0, 1)
img = np.expand_dims(img, 0)
img = torch.from_numpy(img).float()
# Setup Model
model = torch.load(args.model_path)
model.eval()
if torch.cuda.is_available():
model.cuda(0)
images = Variable(img.cuda(0))
else:
images = Variable(img)
outputs = model(images)
pred = np.squeeze(outputs.data.max(1 , keepdim=True )[1].cpu().numpy(), axis=1)
decoded = loader.decode_segmap(pred[0])
print(np.unique(pred))
misc.imsave(args.out_path, decoded)
print("Segmentation Mask Saved at: {}".format(args.out_path))
def test(args, cfg):
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
data_loader = get_loader(cfg['data']['dataset'])
data_path = get_data_path(cfg['data']['dataset'], config_file=cfg)
loader = data_loader(data_path, is_transform=True, img_norm=args.img_norm)
n_classes = loader.n_classes
t_loader = data_loader(
data_path,
is_transform=True,
split='test',
img_size=(cfg['data']['img_rows'], cfg['data']['img_cols']),
)
testloader = data.DataLoader(t_loader,
batch_size=1,
num_workers=cfg['training']['n_workers'])
# Setup Model
model = get_model(cfg['model'], n_classes)
state = convert_state_dict(torch.load(args.model_path)["model_state"])
# state=torch.load(args.model_path)["model_state"]
model.load_state_dict(state)
model.eval()
model.to(device)
with torch.no_grad():
for i_val, (img_path, image_src,
image_dst) in tqdm(enumerate(testloader)):
img_name = img_path[0]
image_src = image_src.to(device)
image_dst = image_dst.to(device)
outputs = model(image_src, image_dst)
pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=0)
decoded = loader.decode_segmap(pred)
out_path = "test_out/changenet_change_det/" + Path(
img_name).stem + ".png"
decoded_bgr = cv.cvtColor(decoded, cv.COLOR_RGB2BGR)
# misc.imsave(out_path, decoded)
cv.imwrite(out_path, decoded_bgr)
def validate(args):
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus
# Setup Dataloader
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
loader = data_loader(data_path,
split=args.split,
is_transform=True,
img_size=(args.img_rows, args.img_cols))
n_classes = loader.n_classes
valloader = data.DataLoader(loader,
batch_size=args.batch_size,
num_workers=4)
running_metrics = runningScore(n_classes)
# Setup Model
model = get_model(args.arch, n_classes)
checkpoint = torch.load(args.model_path)
state = convert_state_dict(checkpoint['model_state'])
model.load_state_dict(state)
print("Loaded checkpoint '{}' (epoch {})".format(args.model_path,
checkpoint['epoch']))
model.eval()
for i, (images, labels) in tqdm(enumerate(valloader)):
model.cuda()
images = Variable(images.cuda(), volatile=True)
labels = Variable(labels.cuda(), volatile=True)
outputs = model(images)
pred = outputs.data.max(1)[1].cpu().numpy()
gt = labels.data.cpu().numpy()
running_metrics.update(gt, pred)
score, class_iou = running_metrics.get_scores()
for k, v in score.items():
print(k, v)
for i in range(n_classes):
print(i, classes[i], class_iou[i])
print('\t'.join([str(class_iou[i]) for i in range(n_classes)]))
def test(args):
# Setup image
print("Read Input Image from : {}".format(args.img_path))
img = misc.imread(args.img_path)
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset, config_file=args.config_file)
loader = data_loader(data_path, is_transform=True)
n_classes = loader.n_classes
img = img[:, :, ::-1]
img = img.astype(np.float64)
img -= loader.mean
img = misc.imresize(img, (loader.img_size[0], loader.img_size[1]))
img = img.astype(float) / 255.0
# NHWC -> NCWH
img = img.transpose(2, 0, 1)
img = np.expand_dims(img, 0)
img = torch.from_numpy(img).float()
# Setup Model
model = torch.load(args.model_path)
model.eval()
#if torch.cuda.is_available():
# model.cuda(0)
# images = Variable(img.cuda(0))
#else:
images = Variable(img)
import time
start_time = time.clock()
outputs = model(images)
print('Time: {time}'.format(time=(time.clock() - start_time)))
pred = outputs[0].cpu().data.numpy()
if args.label:
pred = pred[args.label]
else:
pred = pred.argmax(0)
misc.imsave(args.out_path, pred)
print("Segmentation Mask Saved at: {}".format(args.out_path))
def test():
# Setup Dataloader
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
loader = data_loader(data_path, img_size=args.img_size)
n_classes = loader.n_classes
n_channels = loader.n_channels
valloader = data.DataLoader(loader,
batch_size=args.batch_size,
num_workers=4,
shuffle=True)
# Setup Model
model = torch.load(args.model_path)
model.eval()
if torch.cuda.is_available():
model.cuda(0)
for i, (images, labels) in enumerate(tqdm(valloader)):
if torch.cuda.is_available():
images = Variable(images.cuda(0))
labels = Variable(labels.cuda(0))
else:
images = Variable(images)
labels = Variable(labels)
outputs = model(images)
pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=1)
gt = labels.data.cpu().numpy()
for gt_, pred_ in zip(gt, pred):
gt_path = args.out_dir + "gt{}.png".format(i)
pred_path = args.out_dir + "pred{}.png".format(i)
decoded_gt = loader.decode_segmap(gt_)
decoded_pred = loader.decode_segmap(pred_)
misc.imsave(gt_path, decoded_gt)
misc.imsave(pred_path, decoded_pred)
def test(args):
# Setup image
print("Read Input Image from : {}".format(args.img_path))
img = misc.imread(args.img_path)
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
loader = data_loader(data_path, is_transform=True)
n_classes = loader.n_classes
img = img[:, :, ::-1]
img = img.astype(np.float64)
img -= loader.mean
img = misc.imresize(img, (loader.img_size[0], loader.img_size[1]))
img = img.astype(float) / 255.0
# NHWC -> NCWH
img = img.transpose(2, 0, 1)
img = np.expand_dims(img, 0)
img = torch.from_numpy(img).float()
# Setup Model
model = torch.load(args.model_path)
model.eval()
if torch.cuda.is_available():
model.cuda(0)
images = Variable(img.cuda(0))
else:
images = Variable(img)
outputs = model(images)
pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=1)
decoded = loader.decode_segmap(pred[0])
print(np.unique(pred))
misc.imsave(args.out_path, decoded)
print("Segmentation Mask Saved at: {}".format(args.out_path))
def train(args):
# Setup Augmentations
data_aug= Compose([RandomRotate(10),
RandomHorizontallyFlip()])
# Setup Dataloader
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
t_loader = data_loader(data_path, is_transform=True, img_size=(args.img_rows, args.img_cols), augmentations=data_aug)
v_loader = data_loader(data_path, is_transform=True, split='val', img_size=(args.img_rows, args.img_cols))
n_classes = t_loader.n_classes
trainloader = data.DataLoader(t_loader, batch_size=args.batch_size, num_workers=8, shuffle=True)
valloader = data.DataLoader(v_loader, batch_size=args.batch_size, num_workers=8)
# Setup Metrics
running_metrics = runningScore(n_classes)
# Setup visdom for visualization
if args.visdom:
vis = visdom.Visdom()
loss_window = vis.line(X=torch.zeros((1,)).cpu(),
Y=torch.zeros((1)).cpu(),
opts=dict(xlabel='minibatches',
ylabel='Loss',
title='Training Loss',
legend=['Loss']))
# Setup Model
model = get_model(args.arch, n_classes)
model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))
model.cuda()
# Check if model has custom optimizer / loss
if hasattr(model.module, 'optimizer'):
optimizer = model.module.optimizer
else:
optimizer = torch.optim.SGD(model.parameters(), lr=args.l_rate, momentum=0.99, weight_decay=5e-4)
if hasattr(model.module, 'loss'):
print('Using custom loss')
loss_fn = model.module.loss
else:
loss_fn = cross_entropy2d
if args.resume is not None:
if os.path.isfile(args.resume):
print("Loading model and optimizer from checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['model_state'])
optimizer.load_state_dict(checkpoint['optimizer_state'])
print("Loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
else:
print("No checkpoint found at '{}'".format(args.resume))
best_iou = -100.0
for epoch in range(args.n_epoch):
model.train()
for i, (images, labels) in enumerate(trainloader):
images = Variable(images.cuda())
labels = Variable(labels.cuda())
optimizer.zero_grad()
outputs = model(images)
loss = loss_fn(input=outputs, target=labels)
loss.backward()
optimizer.step()
if args.visdom:
vis.line(
X=torch.ones((1, 1)).cpu() * i,
Y=torch.Tensor([loss.data[0]]).unsqueeze(0).cpu(),
win=loss_window,
update='append')
if (i+1) % 20 == 0:
print("Epoch [%d/%d] Loss: %.4f" % (epoch+1, args.n_epoch, loss.data[0]))
model.eval()
for i_val, (images_val, labels_val) in tqdm(enumerate(valloader)):
images_val = Variable(images_val.cuda(), volatile=True)
labels_val = Variable(labels_val.cuda(), volatile=True)
outputs = model(images_val)
pred = outputs.data.max(1)[1].cpu().numpy()
gt = labels_val.data.cpu().numpy()
running_metrics.update(gt, pred)
score, class_iou = running_metrics.get_scores()
for k, v in score.items():
print(k, v)
running_metrics.reset()
if score['Mean IoU : \t'] >= best_iou:
best_iou = score['Mean IoU : \t']
state = {'epoch': epoch+1,
'model_state': model.state_dict(),
'optimizer_state' : optimizer.state_dict(),}
torch.save(state, "{}_{}_best_model.pkl".format(args.arch, args.dataset))
def test(args, cfg):
# os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# device=torch.device("cuda:0")
# device_1=torch.device("cpu")
model_file_name = os.path.split(args.model_path)[1]
model_name = model_file_name[:model_file_name.find("_")]
IMG_Path = Path(args.img_path)
IMG_File = natsort.natsorted(list(IMG_Path.glob("*.png")),
alg=natsort.PATH)
IMG_Str = []
for i in IMG_File:
IMG_Str.append(str(i))
# Setup image
print("Read Input Image from : {}".format(args.img_path))
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset, config_file=cfg)
loader = data_loader(data_path, is_transform=True, img_norm=args.img_norm)
n_classes = loader.n_classes
# Setup Model
model = get_model(cfg['model'], n_classes)
state = convert_state_dict(torch.load(args.model_path)["model_state"])
# state=torch.load(args.model_path)["model_state"]
model.load_state_dict(state)
model.eval()
model.to(device)
for j in tqdm(range(len(IMG_Str))):
img_path = IMG_Str[j]
img_input = misc.imread(img_path)
sp = list(img_input.shape)
#shape height*width*channel
sp = sp[0:2]
ori_size = tuple(sp)
# img = img[:, :, ::-1]
# multiscale
# img_125=cv.resize(img,dsize=(0,0),fx=1.25,fy=1.25,interpolation=cv.INTER_LINEAR)
# img_075=cv.resize(img,dsize=(0,0),fx=0.75,fy=0.75,interpolation=cv.INTER_LINEAR)
# scale_list=[2.0,1.75,1.5,1.25,1,0.75,0.5]
scale_list = [1.5, 1.25, 0.75, 0.5]
# scale_list=[1.4,1.2,0.8,0.6]
# scale_list=[2.0]
multi_avg = torch.zeros((1, 6, 512, 512),
dtype=torch.float32).to(device)
# torch.zeros(batch-size,num-classes,height,width)
for scale in scale_list:
if scale != 1:
img = cv.resize(img_input,
dsize=(0, 0),
fx=scale,
fy=scale,
interpolation=cv.INTER_LINEAR)
else:
img = img_input
img = img.astype(np.float64)
# img -= loader.mean
if args.img_norm:
img = img.astype(float) / 255.0
# NHWC -> NCHW
img = img.transpose(2, 0, 1)
img = np.expand_dims(img, 0)
img = torch.from_numpy(img).float()
images = img.to(device)
outputs = model(images)
# del images
# bilinear is ok for both upsample and downsample
if scale != 1:
outputs = F.upsample(outputs,
ori_size,
mode='bilinear',
align_corners=False)
# outputs=outputs.to(device)
multi_avg = multi_avg + outputs
# del outputs
# outputs=multi_avg/len(scale_list)
outputs = multi_avg
out_path = "test_out/mv3_1_true_2_res50_data10_MS/mv3_1_true_2_res50_data10_MS_7/" + Path(
img_path).stem + "_S4_not_1.pt"
torch.save(outputs, out_path)
def train(args):
# Setup Augmentations
data_aug = Compose([RandomRotate(10), RandomHorizontallyFlip()])
# Setup Dataloader
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
t_loader = data_loader(data_path,
is_transform=True,
img_size=(args.img_rows, args.img_cols),
augmentations=data_aug,
img_norm=args.img_norm)
v_loader = data_loader(data_path,
is_transform=True,
split='val',
img_size=(args.img_rows, args.img_cols),
img_norm=args.img_norm)
n_classes = t_loader.n_classes
trainloader = data.DataLoader(t_loader,
batch_size=args.batch_size,
num_workers=8,
shuffle=True)
valloader = data.DataLoader(v_loader,
batch_size=args.batch_size,
num_workers=8)
# Setup Metrics
running_metrics = runningScore(n_classes)
# Setup visdom for visualization
if args.visdom:
vis = visdom.Visdom()
loss_window = vis.line(X=torch.zeros((1, )).cpu(),
Y=torch.zeros((1)).cpu(),
opts=dict(xlabel='minibatches',
ylabel='Loss',
title='Training Loss',
legend=['Loss']))
# Setup Model
model = get_model(args.arch, n_classes)
model = torch.nn.DataParallel(model,
device_ids=range(torch.cuda.device_count()))
model.cuda()
# Check if model has custom optimizer / loss
if hasattr(model.module, 'optimizer'):
optimizer = model.module.optimizer
else:
optimizer = torch.optim.SGD(model.parameters(),
lr=args.l_rate,
momentum=0.99,
weight_decay=5e-4)
if hasattr(model.module, 'loss'):
print('Using custom loss')
loss_fn = model.module.loss
else:
loss_fn = cross_entropy2d
if args.resume is not None:
if os.path.isfile(args.resume):
print("Loading model and optimizer from checkpoint '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['model_state'])
optimizer.load_state_dict(checkpoint['optimizer_state'])
print("Loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("No checkpoint found at '{}'".format(args.resume))
best_iou = -100.0
for epoch in range(args.n_epoch):
model.train()
for i, (images, labels) in enumerate(trainloader):
images = images.cuda()
labels = labels.cuda()
optimizer.zero_grad()
outputs = model(images)
loss = loss_fn(input=outputs, target=labels)
loss.backward()
optimizer.step()
if args.visdom:
vis.line(X=torch.ones((1, )).cpu() * i,
Y=torch.Tensor([loss.item()]).cpu(),
win=loss_window,
update='append')
if (i + 1) % 20 == 0:
print("Epoch [%d/%d] Loss: %.4f" %
(epoch + 1, args.n_epoch, loss.item()))
model.eval()
with torch.no_grad():
for i_val, (images_val, labels_val) in tqdm(enumerate(valloader)):
images_val = images_val.cuda()
labels_val = labels_val.cuda()
outputs = model(images_val)
pred = outputs.detach().max(1)[1].cpu().numpy()
gt = labels_val.detach().cpu().numpy()
running_metrics.update(gt, pred)
score, class_iou = running_metrics.get_scores()
for k, v in score.items():
print(k, v)
running_metrics.reset()
if score['Mean IoU : \t'] >= best_iou:
best_iou = score['Mean IoU : \t']
state = {
'epoch': epoch + 1,
'model_state': model.state_dict(),
'optimizer_state': optimizer.state_dict(),
}
torch.save(state,
"{}_{}_best_model.pkl".format(args.arch, args.dataset))
def validate(cfg, args):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_file_name = os.path.split(args.model_path)[1]
model_name = model_file_name[:model_file_name.find("_")]
# Setup Dataloader
data_loader = get_loader(cfg['data']['dataset'])
data_path = get_data_path(cfg['data']['dataset'])
loader = data_loader(
data_path,
split=cfg['data']['val_split'],
is_transform=True,
img_size=(cfg['data']['img_rows'], cfg['data']['img_rows']),
)
n_classes = loader.n_classes
valloader = data.DataLoader(loader,
batch_size=cfg['training']['batch_size'],
num_workers=8)
running_metrics = runningScore(n_classes)
# Setup Model
model = get_model(model_name, n_classes, version=cfg['data']['dataset'])
state = convert_state_dict(torch.load(args.model_path)["model_state"])
model.load_state_dict(state)
model.eval()
model.to(device)
for i, (images, labels) in enumerate(valloader):
start_time = timeit.default_timer()
images = images.to(device)
if args.eval_flip:
outputs = model(images)
# Flip images in numpy (not support in tensor)
outputs = outputs.data.cpu().numpy()
flipped_images = np.copy(images.data.cpu().numpy()[:, :, :, ::-1])
flipped_images = torch.from_numpy(flipped_images).float().to(
device)
outputs_flipped = model(flipped_images)
outputs_flipped = outputs_flipped.data.cpu().numpy()
outputs = (outputs + outputs_flipped[:, :, :, ::-1]) / 2.0
pred = np.argmax(outputs, axis=1)
else:
outputs = model(images)
pred = outputs.data.max(1)[1].cpu().numpy()
gt = labels.numpy()
if args.measure_time:
elapsed_time = timeit.default_timer() - start_time
print("Inference time \
(iter {0:5d}): {1:3.5f} fps".format(
i + 1, pred.shape[0] / elapsed_time))
running_metrics.update(gt, pred)
score, class_iou = running_metrics.get_scores()
for k, v in score.items():
print(k, v)
for i in range(n_classes):
print(i, class_iou[i])
def test(args):
model_file_name = os.path.split(args.model_path)[1]
model_name = model_file_name[:model_file_name.find('_')]
# Setup image
print("Read Input Image from : {}".format(args.img_path))
img = misc.imread(args.img_path)
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
loader = data_loader(data_path, is_transform=True, img_norm=args.img_norm)
n_classes = loader.n_classes
resized_img = misc.imresize(img, (loader.img_size[0], loader.img_size[1]), interp='bicubic')
orig_size = img.shape[:-1]
if model_name in ['pspnet', 'icnet', 'icnetBN']:
img = misc.imresize(img, (orig_size[0]//2*2+1, orig_size[1]//2*2+1)) # uint8 with RGB mode, resize width and height which are odd numbers
else:
img = misc.imresize(img, (loader.img_size[0], loader.img_size[1]))
img = img[:, :, ::-1]
img = img.astype(np.float64)
img -= loader.mean
if args.img_norm:
img = img.astype(float) / 255.0
# NHWC -> NCHW
img = img.transpose(2, 0, 1)
img = np.expand_dims(img, 0)
img = torch.from_numpy(img).float()
# Setup Model
model = get_model(model_name, n_classes, version=args.dataset)
state = convert_state_dict(torch.load(args.model_path)['model_state'])
model.load_state_dict(state)
model.eval()
if torch.cuda.is_available():
model.cuda(0)
images = Variable(img.cuda(0), volatile=True)
else:
images = Variable(img, volatile=True)
outputs = model(images)
#outputs = F.softmax(outputs, dim=1)
if args.dcrf:
unary = outputs.data.cpu().numpy()
unary = np.squeeze(unary, 0)
unary = -np.log(unary)
unary = unary.transpose(2, 1, 0)
w, h, c = unary.shape
unary = unary.transpose(2, 0, 1).reshape(loader.n_classes, -1)
unary = np.ascontiguousarray(unary)
resized_img = np.ascontiguousarray(resized_img)
d = dcrf.DenseCRF2D(w, h, loader.n_classes)
d.setUnaryEnergy(unary)
d.addPairwiseBilateral(sxy=5, srgb=3, rgbim=resized_img, compat=1)
q = d.inference(50)
mask = np.argmax(q, axis=0).reshape(w, h).transpose(1, 0)
decoded_crf = loader.decode_segmap(np.array(mask, dtype=np.uint8))
dcrf_path = args.out_path[:-4] + '_drf.png'
misc.imsave(dcrf_path, decoded_crf)
print("Dense CRF Processed Mask Saved at: {}".format(dcrf_path))
pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=0)
if model_name in ['pspnet', 'icnet', 'icnetBN']:
pred = pred.astype(np.float32)
pred = misc.imresize(pred, orig_size, 'nearest', mode='F') # float32 with F mode, resize back to orig_size
decoded = loader.decode_segmap(pred)
print('Classes found: ', np.unique(pred))
misc.imsave(args.out_path, decoded)
print("Segmentation Mask Saved at: {}".format(args.out_path))
def encode_segmap(self, mask):
#Put all void classes to zero
for _voidc in self.void_classes:
mask[mask==_voidc] = 0
for _validc in self.valid_classes:
mask[mask==_validc] = self.class_map[_validc]
return mask
# Run it as a standalone file from project root directory
if __name__ == '__main__':
import torchvision
import matplotlib.pyplot as plt
from ptsemseg.loader import get_loader, get_data_path
dst = cityscapesLoader(get_data_path('cityscapes'), is_transform=True)
trainloader = data.DataLoader(dst, batch_size=4, num_workers=0)
for i, data in enumerate(trainloader):
imgs, labels = data
img = imgs.numpy()[0, ::-1, :, :]
img = np.transpose(img, [1,2,0])
f, axarr = plt.subplots(2,1)
axarr[0].imshow(img)
axarr[1].imshow(dst.decode_segmap(labels.numpy()[0]))
plt.show()
a = input()
if a == 'ex':
break
else:
plt.close()
def test(args):
# Setup image
print("Read Input Image from : {}".format(args.img_path))
img = misc.imread(args.img_path)
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
loader = data_loader(data_path, is_transform=True)
n_classes = loader.n_classes
resized_img = misc.imresize(img, (loader.img_size[0], loader.img_size[1]), interp='bicubic')
img = img[:, :, ::-1]
img = img.astype(np.float64)
img -= loader.mean
img = misc.imresize(img, (loader.img_size[0], loader.img_size[1]))
img = img.astype(float) / 255.0
# NHWC -> NCWH
img = img.transpose(2, 0, 1)
img = np.expand_dims(img, 0)
img = torch.from_numpy(img).float()
# Setup Model
model = get_model(args.model_path[:args.model_path.find('_')], n_classes)
state = convert_state_dict(torch.load(args.model_path)['model_state'])
model.load_state_dict(state)
model.eval()
model.cuda(0)
images = Variable(img.cuda(0), volatile=True)
outputs = F.softmax(model(images), dim=1)
if args.dcrf == "True":
unary = outputs.data.cpu().numpy()
unary = np.squeeze(unary, 0)
unary = -np.log(unary)
unary = unary.transpose(2, 1, 0)
w, h, c = unary.shape
unary = unary.transpose(2, 0, 1).reshape(loader.n_classes, -1)
unary = np.ascontiguousarray(unary)
resized_img = np.ascontiguousarray(resized_img)
d = dcrf.DenseCRF2D(w, h, loader.n_classes)
d.setUnaryEnergy(unary)
d.addPairwiseBilateral(sxy=5, srgb=3, rgbim=resized_img, compat=1)
q = d.inference(50)
mask = np.argmax(q, axis=0).reshape(w, h).transpose(1, 0)
decoded_crf = loader.decode_segmap(np.array(mask, dtype=np.uint8))
dcrf_path = args.out_path[:-4] + '_drf.png'
misc.imsave(dcrf_path, decoded_crf)
print("Dense CRF Processed Mask Saved at: {}".format(dcrf_path))
if torch.cuda.is_available():
model.cuda(0)
images = Variable(img.cuda(0), volatile=True)
else:
images = Variable(img, volatile=True)
pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=0)
decoded = loader.decode_segmap(pred)
print('Classes found: ', np.unique(pred))
misc.imsave(args.out_path, decoded)
print("Segmentation Mask Saved at: {}".format(args.out_path))
def validate_bayesian(args, model, split, labeled_index=None, verbose=False):
# Setup Data
data_loader = get_loader('camvid')
data_path = get_data_path('camvid')
dataset = data_loader(data_path,
split,
is_transform=True,
labeled_index=labeled_index)
valloader = data.DataLoader(dataset,
batch_size=1,
num_workers=0,
shuffle=False)
# Setup Model
model.eval()
model.apply(set_dropout)
# Uncertainty Hyperparameter
T = args.sample_num
inference_time = 0
gts, preds, uncts = [], [], []
uncts_r, uncts_e, uncts_v, uncts_b = [], [], [], []
image_names_all = []
for i, (images, labels, image_names) in enumerate(valloader):
torch.cuda.synchronize()
t1 = time.time()
if torch.cuda.is_available():
model.cuda()
images = Variable(images.cuda(), volatile=True)
labels_var = Variable(labels.cuda(async=True), volatile=True)
else:
images = Variable(images, volatile=True)
labels_var = Variable(labels, volatile=True)
output_list = []
# MC dropout
for t in range(T):
output = F.softmax(model(images))
if t == 0:
output_mean = output * 0
output_square = output * 0
entropy_mean = output.mean(1) * 0
output_mean += output
output_square += output.pow(2)
entropy_mean += acquisition_func('e', output)
output_mean = output_mean / T
output_square = output_square / T
entropy_mean = entropy_mean / T
# Uncertainty estimation
if args.acqu_func != 'all':
unc_map = acquisition_func(args.acqu_func, output_mean,\
square_mean=output_square, entropy_mean=entropy_mean)
else:
unc_map_r = acquisition_func('r', output_mean,\
square_mean=output_square, entropy_mean=entropy_mean)
unc_map_e = acquisition_func('e', output_mean,\
square_mean=output_square, entropy_mean=entropy_mean)
unc_map_b = acquisition_func('b', output_mean,\
square_mean=output_square, entropy_mean=entropy_mean)
unc_map_v = acquisition_func('v', output_mean,\
square_mean=output_square, entropy_mean=entropy_mean)
pred = torch.max(output_mean, 1)[1]
torch.cuda.synchronize()
t2 = time.time()
gts += list(labels.numpy())
preds += list(pred.data.cpu().numpy())
if args.acqu_func != 'all':
uncts += list(unc_map.data.cpu().numpy())
else:
uncts_r += list(unc_map_r.data.cpu().numpy())
uncts_e += list(unc_map_e.data.cpu().numpy())
uncts_b += list(unc_map_b.data.cpu().numpy())
uncts_v += list(unc_map_v.data.cpu().numpy())
image_names_all += list(image_names)
inference_time += t2 - t1
if verbose:
print '[Info] evaluate ', image_names
print '[Time] Average Inference Time = ', inference_time / (i + 1)
# Save unct_map and pred_map
if args.save_output:
for index in range(len(preds)):
out_name = os.path.basename(image_names_all[index]).replace(
'.png', '')
np.save(os.path.join(args.out_pred_dir, out_name), preds[index])
if args.acqu_func != 'all':
np.save(os.path.join(args.out_unct_dir, out_name),
uncts[index])
else:
np.save(os.path.join(args.out_unct_dir_r, out_name),
uncts_r[index])
np.save(os.path.join(args.out_unct_dir_e, out_name),
uncts_e[index])
np.save(os.path.join(args.out_unct_dir_b, out_name),
uncts_b[index])
np.save(os.path.join(args.out_unct_dir_v, out_name),
uncts_v[index])
return gts, preds, uncts
def merge(args):
result_root_path = make_result_dir(args.dataset, args.split)
# Setup Dataloader
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
loader = data_loader(data_path,
split=args.split,
is_transform=True,
img_size=(args.img_rows, args.img_cols),
img_norm=args.img_norm,
no_gt=args.no_gt)
n_classes = loader.n_classes
testloader = data.DataLoader(loader,
batch_size=args.batch_size,
num_workers=4,
pin_memory=True)
running_metrics = runningScore(n_classes)
rm = 0
pred_dict = {}
map = AverageMeter()
with open('list_test_18000') as f:
id_list = f.read().splitlines()
# Average all the probability maps from all folds
all_prob = np.zeros((len(testloader), args.img_rows, args.img_cols),
dtype=np.float32)
for i in range(1, args.max_k_split + 1):
prob = np.load('prob-{}_{}_{}.npy'.format(args.split, i,
args.max_k_split))
all_prob = all_prob + prob
all_prob = all_prob / args.max_k_split
np.save('prob-{}_avg_{}.npy'.format(args.split, args.max_k_split),
all_prob)
for i, id in tqdm(enumerate(id_list)):
lbl = id + '.png'
pred = all_prob[i, :, :]
if not args.no_gt:
gt_path = os.path.join(data_path, args.split, 'masks', lbl)
gt = loader.encode_segmap(cv2.imread(gt_path,
cv2.IMREAD_GRAYSCALE))
pred = np.where(pred < args.pred_thr, 0, 1)
if pred.sum(
) <= loader.lbl_thr: # Remove salt masks for sum of salts <= a threshold lbl_thr
pred = np.zeros((args.img_rows, args.img_cols), dtype=np.uint8)
rm = rm + 1
decoded = loader.decode_segmap(pred)
rle_mask = loader.RLenc(decoded)
pred_dict[id] = rle_mask
save_result_path = os.path.join(result_root_path, id + '.png')
cv2.imwrite(save_result_path, decoded)
if not args.no_gt:
map_val = running_metrics.comput_map(gt, pred)
map.update(map_val.mean(), n=map_val.size)
if not args.no_gt:
print('Mean Average Precision: {:.5f}'.format(map.avg))
# Create final submission
sub = pd.DataFrame.from_dict(pred_dict, orient='index')
sub.index.names = ['id']
sub.columns = ['rle_mask']
sub.to_csv(args.split + '.csv')
print('To black: ', rm)
def train(args):
logger.auto_set_dir()
os.environ['CUDA_VISIBLE_DEVICES'] = '3'
# Setup Augmentations
data_aug = Compose([RandomRotate(10), RandomHorizontallyFlip()])
# Setup Dataloader
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
t_loader = data_loader(data_path,
is_transform=True,
img_size=(args.img_rows, args.img_cols),
epoch_scale=4,
augmentations=data_aug,
img_norm=args.img_norm)
v_loader = data_loader(data_path,
is_transform=True,
split='val',
img_size=(args.img_rows, args.img_cols),
img_norm=args.img_norm)
n_classes = t_loader.n_classes
trainloader = data.DataLoader(t_loader,
batch_size=args.batch_size,
num_workers=8,
shuffle=True)
valloader = data.DataLoader(v_loader,
batch_size=args.batch_size,
num_workers=8)
# Setup Metrics
running_metrics = runningScore(n_classes)
# Setup Model
from model_zoo.deeplabv1 import VGG16_LargeFoV
model = VGG16_LargeFoV(class_num=n_classes,
image_size=[args.img_cols, args.img_rows],
pretrained=True)
#model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))
model.cuda()
# Check if model has custom optimizer / loss
if hasattr(model, 'optimizer'):
logger.warn("don't have customzed optimizer, use default setting!")
optimizer = model.module.optimizer
else:
optimizer = torch.optim.SGD(model.parameters(),
lr=args.l_rate,
momentum=0.99,
weight_decay=5e-4)
optimizer_summary(optimizer)
if args.resume is not None:
if os.path.isfile(args.resume):
logger.info(
"Loading model and optimizer from checkpoint '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['model_state'])
optimizer.load_state_dict(checkpoint['optimizer_state'])
logger.info("Loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
logger.info("No checkpoint found at '{}'".format(args.resume))
best_iou = -100.0
for epoch in tqdm(range(args.n_epoch), total=args.n_epoch):
model.train()
for i, (images, labels) in tqdm(enumerate(trainloader),
total=len(trainloader),
desc="training epoch {}/{}".format(
epoch, args.n_epoch)):
cur_iter = i + epoch * len(trainloader)
cur_lr = adjust_learning_rate(optimizer,
args.l_rate,
cur_iter,
args.n_epoch * len(trainloader),
power=0.9)
#if i > 10:break
images = Variable(images.cuda())
labels = Variable(labels.cuda())
optimizer.zero_grad()
outputs = model(images)
#print(np.unique(outputs.data[0].cpu().numpy()))
loss = CrossEntropyLoss2d_Seg(input=outputs,
target=labels,
class_num=n_classes)
loss.backward()
optimizer.step()
if (i + 1) % 100 == 0:
logger.info("Epoch [%d/%d] Loss: %.4f, lr: %.7f" %
(epoch + 1, args.n_epoch, loss.data[0], cur_lr))
model.eval()
for i_val, (images_val, labels_val) in tqdm(enumerate(valloader),
total=len(valloader),
desc="validation"):
images_val = Variable(images_val.cuda(), volatile=True)
labels_val = Variable(labels_val.cuda(), volatile=True)
outputs = model(images_val)
pred = outputs.data.max(1)[1].cpu().numpy()
gt = labels_val.data.cpu().numpy()
running_metrics.update(gt, pred)
score, class_iou = running_metrics.get_scores()
for k, v in score.items():
logger.info("{}: {}".format(k, v))
running_metrics.reset()
if score['Mean IoU : \t'] >= best_iou:
best_iou = score['Mean IoU : \t']
state = {
'epoch': epoch + 1,
'mIoU': best_iou,
'model_state': model.state_dict(),
'optimizer_state': optimizer.state_dict(),
}
torch.save(state,
os.path.join(logger.get_logger_dir(), "best_model.pkl"))
args.deep_features_size = 256
elif args.backend == 'resnet50' or args.backend == 'resnet101' or args.backend == 'resnet152':
args.psp_size = 2048
args.deep_features_size = 1024
# Setup Dataloader
val_transforms = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(256),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
val_dataset = data_loader(data_path, split='val', transform=val_transforms)
args.n_classes = val_dataset.n_classes
valloader = data.DataLoader(val_dataset,
batch_size=args.batch_size,
num_workers=args.num_workers,
shuffle=False,
pin_memory=True)
# Setup Model
model = get_model(args)
model.load_state_dict(torch.load(args.model_path)['state_dict'])
model = torch.nn.DataParallel(model,
device_ids=range(
torch.cuda.device_count())).cuda()
