def test_visualize_object(self):
"""Given a dictionary object as follows
x['img']: numpy array of shape (height,width,colors)
x['mask']: numpy array of shape (height,width), with every element categorizing it
into one of the object ids
The method generates an image overlaying a translucent mask on the image.
"""
visualize_mask(self.test_object, self.c, self.transparency)
def segment(dataloader, segment_dir, model, core_config):
model.eval() # convert the model into evaluation mode
img_dir = os.path.join(segment_dir, 'img')
rle_dir = os.path.join(segment_dir, 'rle')
mask_dir = os.path.join(segment_dir, 'mask')
if not os.path.exists(rle_dir):
os.makedirs(rle_dir)
if not os.path.exists(img_dir):
os.makedirs(img_dir)
if not os.path.exists(mask_dir):
os.makedirs(mask_dir)
exist_ids = next(os.walk(img_dir))[2]
num_classes = core_config.num_classes
offset_list = core_config.offsets
if args.object_merge_factor is None:
args.object_merge_factor = 1.0 / len(offset_list)
segmenter_opts = SegmenterOptions(
same_different_bias=args.same_different_bias,
object_merge_factor=args.object_merge_factor,
merge_logprob_bias=args.merge_logprob_bias)
for i, (img, size, id) in enumerate(dataloader):
id = id[0] # tuple to str
if id + '.png' in exist_ids:
continue
original_height, original_width = size[0].item(), size[1].item()
print("Processing image '{}'...".format(id))
with torch.no_grad():
output = model(img)
class_pred = output[:, :num_classes, :, :]
adj_pred = output[:, num_classes:, :, :]
seg = ObjectSegmenter(class_pred[0].detach().numpy(),
adj_pred[0].detach().numpy(), num_classes,
offset_list, segmenter_opts)
seg.run_segmentation()
seg.prune(args.prune_threshold)
mask_pred, object_class = seg.output_mask()
img = np.moveaxis(img[0].detach().numpy(), 0, -1)
segmented_img = {}
segmented_img['img'] = img
segmented_img['mask'] = mask_pred
segmented_img['object_class'] = object_class
visualize_mask(segmented_img, core_config)
scipy.misc.imsave('{}/{}.png'.format(img_dir, id),
segmented_img['img_with_mask'])
filename = mask_dir + '/' + id + '.' + 'mask' + '.npy'
np.save(filename, mask_pred)
rles = list(mask_to_rles(mask_pred))
segment_rle_file = '{}/{}.rle'.format(rle_dir, id)
with open(segment_rle_file, 'w') as fh:
for obj in rles:
obj_str = ' '.join(str(n) for n in obj)
fh.write(obj_str)
fh.write('\n')
print("\n\n")
def visualize_object(x, transparency):
"""Given a dictionary object as follows
x['img']: numpy array of shape (num_class,width,height)
x['mask']: numpy array of same dimensions as image, but with every element categorizing it
into one of the object ids
The method generates an image overlaying a translucent mask on the image and displays it.
"""
c = CoreConfig()
c.num_colors = x['img'].shape[0]
visualize_mask(x,c,transparency)
return
def segment(dataloader, segment_dir, model, core_config):
model.eval() # convert the model into evaluation mode
lbl_dir = os.path.join(segment_dir, 'lbl')
img_dir = os.path.join(segment_dir, 'img')
if not os.path.exists(lbl_dir):
os.makedirs(lbl_dir)
if not os.path.exists(img_dir):
os.makedirs(img_dir)
exist_ids = next(os.walk(lbl_dir))[2]
num_classes = core_config.num_classes
offset_list = core_config.offsets
for i, (img, size, id) in enumerate(dataloader):
id = "res_" + id[0][:-4] # tuple to str
if id + '.txt' in exist_ids:
continue
original_height, original_width = size[0].item(), size[1].item()
with torch.no_grad():
output = model(img)
# class_pred = (output[:, :num_classes, :, :] + 0.001) * 0.999
# adj_pred = (output[:, num_classes:, :, :] + 0.001) * 0.999
class_pred = output[:, :num_classes, :, :]
adj_pred = output[:, num_classes:, :, :]
if args.object_merge_factor is None:
args.object_merge_factor = 1.0 / len(offset_list)
segmenter_opts = SegmenterOptions(
same_different_bias=args.same_different_bias,
object_merge_factor=args.object_merge_factor,
merge_logprob_bias=args.merge_logprob_bias)
seg = ObjectSegmenter(class_pred[0].detach().numpy(),
adj_pred[0].detach().numpy(), num_classes,
offset_list, segmenter_opts)
mask_pred, object_class = seg.run_segmentation()
mask_pred = resize(mask_pred, (original_height, original_width),
order=0,
preserve_range=True).astype(int)
image_with_mask = {}
img = np.moveaxis(img[0].detach().numpy(), 0, -1)
img = resize(img, (original_height, original_width),
preserve_range=True)
image_with_mask['img'] = img
image_with_mask['mask'] = mask_pred
image_with_mask['object_class'] = object_class
visual_mask = visualize_mask(image_with_mask,
core_config)['img_with_mask']
scipy.misc.imsave('{}/{}.png'.format(img_dir, id), visual_mask)
lbls = list(mask_to_lbls(mask_pred))
segment_lbl_file = '{}/{}.txt'.format(lbl_dir, id)
with open(segment_lbl_file, 'w') as fh:
for obj in lbls:
obj_str = ','.join(','.join(str(i) for i in point)
for point in obj)
obj_str = re.sub('[() ]', '')
fh.write(obj_str)
fh.write('\n')
def segment(dataloader, segment_dir, core_config):
rle_dir = os.path.join(segment_dir, 'rle')
img_dir = os.path.join(segment_dir, 'img')
if not os.path.exists(rle_dir):
os.makedirs(rle_dir)
if not os.path.exists(img_dir):
os.makedirs(img_dir)
num_classes = core_config.num_classes
offset_list = core_config.offsets
img, class_pred, adj_pred = next(iter(dataloader))
# By default, we use c++ version segmenter. In short, we call function
# cseg.run_segmentation(). For details, please check the "README" file
# in directory whose path is "scripts/waldo/csegmenter".
# If the c++ version segmenter is not available, we can comment out the
# python segmenter and use it.
"""
if args.object_merge_factor is None:
args.object_merge_factor = 1.0 / len(offset_list)
segmenter_opts = SegmenterOptions(same_different_bias=args.same_different_bias,
object_merge_factor=args.object_merge_factor,
merge_logprob_bias=args.merge_logprob_bias)
seg = ObjectSegmenter(class_pred[0].detach().numpy(),
adj_pred[0].detach().numpy(),
num_classes, offset_list,
segmenter_opts)
mask_pred, object_class = seg.run_segmentation()
"""
if args.object_merge_factor is None:
args.object_merge_factor = 1.0 / len(offset_list)
class_pred_in = class_pred[0].detach().numpy().astype(np.float32)
adj_pred_in = adj_pred[0].detach().numpy().astype(np.float32)
mask_pred, object_class = cseg.run_segmentation(class_pred_in, adj_pred_in,
num_classes, offset_list,
args.same_different_bias,
args.object_merge_factor,
args.merge_logprob_bias)
image_with_mask = {}
img = np.moveaxis(img[0].detach().numpy(), 0, -1)
image_with_mask['img'] = img
image_with_mask['mask'] = mask_pred
image_with_mask['object_class'] = object_class
visual_mask = visualize_mask(image_with_mask, core_config)['img_with_mask']
scipy.misc.imsave('{}/oracle.png'.format(img_dir), visual_mask)
rles = list(mask_to_rles(mask_pred))
segment_rle_file = '{}/oracle.rle'.format(rle_dir)
with open(segment_rle_file, 'w') as fh:
for obj in rles:
obj_str = ' '.join(str(n) for n in obj)
fh.write(obj_str)
fh.write('\n')
def main():
global args
args = parser.parse_args()
args.batch_size = 1 # only segment one image for experiment
core_config_path = os.path.join(args.dir, 'configs/core.config')
unet_config_path = os.path.join(args.dir, 'configs/unet.config')
core_config = CoreConfig()
core_config.read(core_config_path)
print('Using core configuration from {}'.format(core_config_path))
# loading Unet configuration
unet_config = UnetConfig()
unet_config.read(unet_config_path, args.train_image_size)
print('Using unet configuration from {}'.format(unet_config_path))
offset_list = core_config.offsets
print("offsets are: {}".format(offset_list))
# model configurations from core config
num_classes = core_config.num_classes
num_colors = core_config.num_colors
num_offsets = len(core_config.offsets)
# model configurations from unet config
start_filters = unet_config.start_filters
up_mode = unet_config.up_mode
merge_mode = unet_config.merge_mode
depth = unet_config.depth
model = UNet(num_classes,
num_offsets,
in_channels=num_colors,
depth=depth,
start_filts=start_filters,
up_mode=up_mode,
merge_mode=merge_mode)
model_path = os.path.join(args.dir, args.model)
if os.path.isfile(model_path):
print("=> loading checkpoint '{}'".format(model_path))
checkpoint = torch.load(model_path,
map_location=lambda storage, loc: storage)
model.load_state_dict(checkpoint['state_dict'])
print("loaded.")
else:
print("=> no checkpoint found at '{}'".format(model_path))
model.eval() # convert the model into evaluation mode
testset = WaldoDataset(args.test_data, core_config, args.train_image_size)
print('Total samples in the test set: {0}'.format(len(testset)))
dataloader = torch.utils.data.DataLoader(testset,
num_workers=1,
batch_size=args.batch_size)
segment_dir = '{}/segment'.format(args.dir)
if not os.path.exists(segment_dir):
os.makedirs(segment_dir)
img, class_pred, adj_pred = sample(model, dataloader, segment_dir,
core_config)
seg = ObjectSegmenter(class_pred[0].detach().numpy(),
adj_pred[0].detach().numpy(), num_classes,
offset_list)
mask_pred, object_class = seg.run_segmentation()
x = {}
# from (color, height, width) to (height, width, color)
x['img'] = np.moveaxis(img[0].numpy(), 0, -1)
x['mask'] = mask_pred.astype(int)
x['object_class'] = object_class
visualize_mask(x, core_config)
def segment(dataloader, segment_dir, model, core_config):
model.eval() # convert the model into evaluation mode
rle_dir = os.path.join(segment_dir, 'rle')
img_dir = os.path.join(segment_dir, 'img')
if not os.path.exists(rle_dir):
os.makedirs(rle_dir)
if not os.path.exists(img_dir):
os.makedirs(img_dir)
exist_ids = next(os.walk(rle_dir))[2]
num_classes = core_config.num_classes
offset_list = core_config.offsets
for i, (img, size, id) in enumerate(dataloader):
id = id[0] # tuple to str
if id + '.rle' in exist_ids:
continue
original_height, original_width = size[0].item(), size[1].item()
with torch.no_grad():
output = model(img)
# class_pred = (output[:, :num_classes, :, :] + 0.001) * 0.999
# adj_pred = (output[:, num_classes:, :, :] + 0.001) * 0.999
class_pred = output[:, :num_classes, :, :]
adj_pred = output[:, num_classes:, :, :]
# By default, we use c++ version segmenter. In short, we call function
# cseg.run_segmentation(). For details, please check the "README" file
# in directory whose path is "scripts/waldo/csegmenter".
# If the c++ version segmenter is not available, we can comment out the
# python segmenter and use it.
"""
if args.object_merge_factor is None:
args.object_merge_factor = 1.0 / len(offset_list)
segmenter_opts = SegmenterOptions(same_different_bias=args.same_different_bias,
object_merge_factor=args.object_merge_factor,
merge_logprob_bias=args.merge_logprob_bias)
seg = ObjectSegmenter(class_pred[0].detach().numpy(),
adj_pred[0].detach().numpy(),
num_classes, offset_list,
segmenter_opts)
mask_pred, object_class = seg.run_segmentation()
"""
if args.object_merge_factor is None:
args.object_merge_factor = 1.0 / len(offset_list)
mask_pred, object_class = cseg.run_segmentation(
class_pred[0].detach().numpy().astype(np.float32),
adj_pred[0].detach().numpy().astype(np.float32),
num_classes,
offset_list,
args.same_different_bias,
args.object_merge_factor,
args.merge_logprob_bias)
mask_pred = resize(mask_pred, (original_height, original_width),
order=0, preserve_range=True).astype(int)
image_with_mask = {}
img = np.moveaxis(img[0].detach().numpy(), 0, -1)
img = resize(img, (original_height, original_width),
preserve_range=True)
image_with_mask['img'] = img
image_with_mask['mask'] = mask_pred
image_with_mask['object_class'] = object_class
visual_mask = visualize_mask(image_with_mask, core_config)[
'img_with_mask']
scipy.misc.imsave('{}/{}.png'.format(img_dir, id), visual_mask)
rles = list(mask_to_rles(mask_pred))
segment_rle_file = '{}/{}.rle'.format(rle_dir, id)
with open(segment_rle_file, 'w') as fh:
for obj in rles:
obj_str = ' '.join(str(n) for n in obj)
fh.write(obj_str)
fh.write('\n')