def visualize_predictions(run_dir, dataset, inference_config, pred_mask_dir, pred_info_dir, show_bbox=True, show_scores=True, show_class=True):
"""Visualizes predictions."""
# Create subdirectory for prediction visualizations
vis_dir = os.path.join(run_dir, 'vis')
utils.mkdir_if_missing(vis_dir)
# Feed images into model one by one. For each image visualize predictions
image_ids = dataset.image_ids
print('VISUALIZING PREDICTIONS')
for image_id in tqdm(image_ids):
# Load image and ground truth data and resize for net
image, _, _, _, _ = modellib.load_image_gt(dataset, inference_config, image_id, use_mini_mask=False)
if inference_config.IMAGE_CHANNEL_COUNT == 1:
image = np.repeat(image, 3, axis=2)
# load mask and info
r = np.load(os.path.join(pred_info_dir, 'image_{:06}.npy'.format(image_id))).item()
r_masks = np.load(os.path.join(pred_mask_dir, 'image_{:06}.npy'.format(image_id)))
# Must transpose from (n, h, w) to (h, w, n)
r['masks'] = np.transpose(r_masks, (1, 2, 0))
# Visualize
scores = r['scores'] if show_scores else None
fig = plt.figure(figsize=(1.7067, 1.7067), dpi=300, frameon=False)
ax = plt.Axes(fig, [0.,0.,1.,1.])
fig.add_axes(ax)
visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'], ['bg', 'obj'],
ax=ax, scores=scores, show_bbox=show_bbox, show_class=show_class)
file_name = os.path.join(vis_dir, 'vis_{:06d}'.format(image_id))
fig.savefig(file_name, transparent=True, dpi=300)
plt.close()
def visualize_gts(run_dir, dataset, inference_config, show_bbox=True, show_scores=False, show_class=True):
"""Visualizes gts."""
# Create subdirectory for gt visualizations
vis_dir = os.path.join(run_dir, 'gt_vis')
utils.mkdir_if_missing(vis_dir)
# Feed images one by one
image_ids = dataset.image_ids
print('VISUALIZING GROUND TRUTHS')
for image_id in tqdm(image_ids):
# Load image and ground truth data and resize for net
image, _, gt_class_id, gt_bbox, gt_mask = modellib.load_image_gt(dataset, inference_config, image_id,
use_mini_mask=False)
if inference_config.IMAGE_CHANNEL_COUNT == 1:
image = np.repeat(image, 3, axis=2)
# Visualize
scores = np.ones(gt_class_id.size) if show_scores else None
fig = plt.figure(figsize=(1.7067, 1.7067), dpi=300, frameon=False)
ax = plt.Axes(fig, [0.,0.,1.,1.])
fig.add_axes(ax)
visualize.display_instances(image, gt_bbox, gt_mask, gt_class_id, ['bg', 'obj'],
scores, ax=ax, show_bbox=show_bbox, show_class=show_class)
file_name = os.path.join(vis_dir, 'gt_vis_{:06d}'.format(image_id))
height, width = image.shape[:2]
fig.savefig(file_name, transparent=True, dpi=300)
plt.close()
def visualize_predictions(run_dir, dataset, inference_config, pred_mask_dir, pred_info_dir, show_bbox=True, show_scores=True, show_class=True):
"""Visualizes predictions."""
# Create subdirectory for prediction visualizations
vis_dir = os.path.join(run_dir, 'vis')
utils.mkdir_if_missing(vis_dir)
# Feed images into model one by one. For each image, predict, save, visualize?
image_ids = dataset.image_ids
print('VISUALIZING PREDICTIONS')
for image_id in tqdm(image_ids):
# Load image and ground truth data and resize for net
image, _, _, _, _ = modellib.load_image_gt(dataset, inference_config, image_id, use_mini_mask=False)
# load mask and info
r = np.load(os.path.join(pred_info_dir, 'image_{:06}.npy'.format(image_id))).item()
r_masks = np.load(os.path.join(pred_mask_dir, 'image_{:06}.npy'.format(image_id)))
# Must transpose from (n, h, w) to (h, w, n)
r['masks'] = np.transpose(r_masks, (1, 2, 0))
# Visualize
scores = r['scores'] if show_scores else None
visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'],
['bg', 'obj'], scores=scores, show_bbox=show_bbox, show_class=show_class)
file_name = os.path.join(vis_dir, 'vis_{:06d}'.format(image_id))
plt.savefig(file_name, bbox_inches='tight', pad_inches=0)
plt.close()
def __init__(self, cfgFile="cfg/maskrcnn.yaml"):
config = YamlConfig(cfgFile)
print("Benchmarking model.")
# Create new directory for outputs
output_dir = config['output_dir']
utils.mkdir_if_missing(output_dir)
# Save config in output directory
image_shape = config['model']['settings']['image_shape']
config['model']['settings']['image_min_dim'] = min(image_shape)
config['model']['settings']['image_max_dim'] = max(image_shape)
config['model']['settings']['gpu_count'] = 1
config['model']['settings']['images_per_gpu'] = 1
inference_config = MaskConfig(config['model']['settings'])
model_dir, _ = os.path.split(config['model']['path'])
self.model = modellib.MaskRCNN(mode=config['model']['mode'],
config=inference_config,
model_dir=model_dir)
print(("Loading weights from ", config['model']['path']))
self.model.load_weights(config['model']['path'], by_name=True)
self.graph = tf.get_default_graph()
print(self.model.keras_model.layers[0].dtype)
def benchmark(config):
"""Computes and stores predictions and then
evaluates them on COCO metrics and supplementary benchmarking script."""
print("Benchmarking Baseline method {}.".format(
config["detector"]["type"]))
# Create new directory for run outputs
# In what location should we put this new directory?
output_dir = config["output_dir"]
mkdir_if_missing(output_dir)
# Save config in run directory
config.save(os.path.join(output_dir, config["save_conf_name"]))
# directory of test images and segmasks
detector_type = config["detector"]["type"]
if detector_type == "euclidean" or detector_type == "region_growing":
from sd_maskrcnn.pcl.pydetect import detect
elif detector_type == "gop" or detector_type == "mcg":
from sd_maskrcnn.gop.detect import detect
else:
print("Detector type not supported")
exit()
# Create predictions and record where everything gets stored.
pred_mask_dir, pred_info_dir, gt_mask_dir = detect(
detector_type,
config["detector"][detector_type],
output_dir,
config["dataset"],
)
ap, ar = coco_benchmark(pred_mask_dir, pred_info_dir, gt_mask_dir)
if config["vis"]["predictions"]:
visualize_predictions(
output_dir,
config["dataset"],
pred_mask_dir,
pred_info_dir,
show_bbox=config["vis"]["show_bbox_pred"],
show_class=config["vis"]["show_class_pred"],
)
if config["vis"]["s_bench"]:
s_benchmark(
output_dir,
config["dataset"],
pred_mask_dir,
pred_info_dir,
gt_mask_dir,
)
print("Saved benchmarking output to {}.\n".format(output_dir))
return ap, ar
def resize_images(config):
"""Resizes all images so their maximum dimension is config['max_dim']. Saves to new directory."""
base_dir = config['dataset']['path']
# directories of images that need resizing
if config['images']['resize']:
image_dir = config['dataset']['img_dir']
image_out_dir = config['dataset']['img_out_dir']
utils.mkdir_if_missing(os.path.join(base_dir, image_out_dir))
old_im_path = os.path.join(base_dir, image_dir)
new_im_path = os.path.join(base_dir, image_out_dir)
for im_path in tqdm(os.listdir(old_im_path)):
im_old_path = os.path.join(old_im_path, im_path)
if '.npy' in im_old_path:
im = np.load(im_old_path)
else:
im = cv2.imread(im_old_path, cv2.IMREAD_UNCHANGED)
im = scale_to_square(im, dim=config['images']['max_dim'])
im_name = im_path.split('.')[0]
im_ext = config['images']['out_ext']
new_im_file = os.path.join(new_im_path, im_name+'.'+im_ext)
if im_ext == 'npy':
if config['images']['normalize']:
im = cv2.normalize(im, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F)
np.save(new_im_file, im)
elif im_ext == 'png':
if config['images']['normalize']:
im = cv2.normalize(im, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX)
cv2.imwrite(new_im_file, im, [cv2.IMWRITE_PNG_COMPRESSION, 0]) # 0 compression
if config['masks']['resize']:
mask_dir = config['dataset']['mask_dir']
mask_out_dir = config['dataset']['mask_out_dir']
utils.mkdir_if_missing(os.path.join(base_dir, mask_out_dir))
old_mask_path = os.path.join(base_dir, mask_dir)
new_mask_path = os.path.join(base_dir, mask_out_dir)
for mask_path in tqdm(os.listdir(old_mask_path)):
mask_old_path = os.path.join(old_mask_path, mask_path)
if '.npy' in mask_old_path:
mask = np.load(mask_old_path)
else:
mask = cv2.imread(mask_old_path, cv2.IMREAD_UNCHANGED)
if mask.shape[0] == 0 or mask.shape[1] == 0:
print("mask empty")
continue
mask = scale_to_square(mask, dim=config['masks']['max_dim'])
new_mask_file = os.path.join(new_mask_path, mask_path)
if '.npy' in mask_old_path:
np.save(new_mask_file, mask)
else:
cv2.imwrite(new_mask_file, mask, [cv2.IMWRITE_PNG_COMPRESSION, 0]) # 0 compression
def benchmark(config):
"""Computes and stores predictions and then
evaluates them on COCO metrics and supplementary benchmarking script."""
print("Benchmarking Baseline method {}.".format(
config['detector']['type']))
# Create new directory for run outputs
# In what location should we put this new directory?
output_dir = config['output_dir']
mkdir_if_missing(output_dir)
# Save config in run directory
config.save(os.path.join(output_dir, config['save_conf_name']))
# directory of test images and segmasks
detector_type = config['detector']['type']
if detector_type == 'euclidean' or detector_type == 'region_growing':
from sd_maskrcnn.pcl.pydetect import detect
elif detector_type == 'gop' or detector_type == 'mcg':
from sd_maskrcnn.gop.detect import detect
else:
print('Detector type not supported')
exit()
# Create predictions and record where everything gets stored.
pred_mask_dir, pred_info_dir, gt_mask_dir = \
detect(detector_type, config['detector'][detector_type], output_dir, config['dataset'])
ap, ar = coco_benchmark(pred_mask_dir, pred_info_dir, gt_mask_dir)
if config['vis']['predictions']:
visualize_predictions(output_dir,
config['dataset'],
pred_mask_dir,
pred_info_dir,
show_bbox=config['vis']['show_bbox_pred'],
show_class=config['vis']['show_class_pred'])
if config['vis']['s_bench']:
s_benchmark(output_dir, config['dataset'], pred_mask_dir,
pred_info_dir, gt_mask_dir)
print("Saved benchmarking output to {}.\n".format(output_dir))
return ap, ar
def visualize_gts(run_dir, dataset, inference_config, show_bbox=True, show_scores=False, show_class=True):
"""Visualizes predictions."""
# Create subdirectory for prediction visualizations
vis_dir = os.path.join(run_dir, 'gt_vis')
utils.mkdir_if_missing(vis_dir)
# Feed images into model one by one. For each image, predict, save, visualize?
image_ids = dataset.image_ids
print('VISUALIZING GROUND TRUTHS')
for image_id in tqdm(image_ids):
# Load image and ground truth data and resize for net
image, _, gt_class_id, gt_bbox, gt_mask = modellib.load_image_gt(dataset, inference_config, image_id,
use_mini_mask=False)
# Visualize
scores = np.ones(gt_class_id.size) if show_scores else None
visualize.display_instances(image, gt_bbox, gt_mask, gt_class_id,
['bg', 'obj'], scores, show_bbox=show_bbox, show_class=show_class)
file_name = os.path.join(vis_dir, 'gt_vis_{:06d}'.format(image_id))
plt.savefig(file_name, bbox_inches='tight', pad_inches=0)
plt.close()
def augment(config):
"""
Using provided image directory and output directory, perform data
augmentation methods on each image and save the new copy to the
output directory.
"""
img_dir = config["img_dir"]
out_dir = config["out_dir"]
mkdir_if_missing(out_dir)
print(("Augmenting data in directory {}.\n".format(img_dir)))
for img_file in tqdm(os.listdir(img_dir)):
if img_file.endswith(".png"):
# read in image
img_path = os.path.join(img_dir, img_file)
img = skimage.io.imread(img_path, as_grey=True)
# return list of augmented images and save
new_img = augment_img(img, config)
out_path = os.path.join(out_dir, img_file)
skimage.io.imsave(out_path, skimage.img_as_ubyte(new_img))
print(("Augmentation complete; files saved in {}.\n".format(out_dir)))
def train(config):
# Training dataset
dataset_train = ImageDataset(config)
dataset_train.load(config['dataset']['train_indices'], augment=True)
dataset_train.prepare()
# Validation dataset
dataset_val = ImageDataset(config)
dataset_val.load(config['dataset']['val_indices'])
dataset_val.prepare()
# Load config
image_shape = config['model']['settings']['image_shape']
config['model']['settings']['image_min_dim'] = min(image_shape)
config['model']['settings']['image_max_dim'] = max(image_shape)
train_config = MaskConfig(config['model']['settings'])
train_config.STEPS_PER_EPOCH = dataset_train.indices.size/(train_config.IMAGES_PER_GPU*train_config.GPU_COUNT)
train_config.display()
# Create directory if it doesn't currently exist
utils.mkdir_if_missing(config['model']['path'])
# Create the model.
model = modellib.MaskRCNN(mode='training', config=train_config,
model_dir=config['model']['path'])
# Select weights file to load
if config['model']['weights'].lower() == "coco":
weights_path = os.path.join(config['model']['path'], 'mask_rcnn_coco.h5')
# Download weights file
if not os.path.exists(weights_path):
utilslib.download_trained_weights(weights_path)
elif config['model']['weights'].lower() == "last":
# Find last trained weights
weights_path = model.find_last()
elif config['model']['weights'].lower() == "imagenet":
# Start from ImageNet trained weights
weights_path = model.get_imagenet_weights()
else:
weights_path = config['model']['weights']
# Load weights
exclude_layers = []
print("Loading weights ", weights_path)
if config['model']['weights'].lower() == "coco":
# Exclude the last layers because they require a matching
# number of classes
if config['model']['settings']['image_channel_count'] == 1:
exclude_layers = ['conv1']
exclude_layers += ["mrcnn_class_logits", "mrcnn_bbox_fc", "mrcnn_bbox", "mrcnn_mask"]
model.load_weights(weights_path, by_name=True, exclude=exclude_layers)
elif config['model']['weights'].lower() == "imagenet":
if config['model']['settings']['image_channel_count'] == 1:
exclude_layers = ['conv1']
model.load_weights(weights_path, by_name=True, exclude=exclude_layers)
elif config['model']['weights'].lower() != "new":
model.load_weights(weights_path, by_name=True)
# save config in run folder
config.save(os.path.join(config['model']['path'], config['save_conf_name']))
# train and save weights to model_path
model.train(dataset_train, dataset_val, learning_rate=train_config.LEARNING_RATE,
epochs=config['model']['epochs'], layers='all')
# save in the models folder
current_datetime = time.strftime("%Y%m%d-%H%M%S")
model_path = os.path.join(config['model']['path'], "mask_rcnn_{}_{}.h5".format(train_config.NAME, current_datetime))
model.keras_model.save_weights(model_path)
def detect(detector_type, config, run_dir, test_config):
"""Run RGB Object Proposal-based detection on a color-image-based dataset.
Parameters
----------
detector_type : str
type of detector (either mcg or gop)
config : dict
config for a GOP/MCG detector
run_dir : str
Directory to save outputs in. Output will be saved in pred_masks, pr$
and modal_segmasks_processed subdirectories.
test_config : dict
config containing dataset information
"""
##################################################################
# Set up output directories
##################################################################
# Create subdirectory for prediction masks
pred_dir = os.path.join(run_dir, 'pred_masks')
mkdir_if_missing(pred_dir)
# Create subdirectory for prediction scores & bboxes
pred_info_dir = os.path.join(run_dir, 'pred_info')
mkdir_if_missing(pred_info_dir)
# Create subdirectory for transformed GT segmasks
resized_segmask_dir = os.path.join(run_dir, 'modal_segmasks_processed')
mkdir_if_missing(resized_segmask_dir)
##################################################################
# Set up input directories
##################################################################
dataset_dir = test_config['path']
indices_arr = np.load(os.path.join(dataset_dir, test_config['indices']))
# Input depth image data (numpy files, not .pngs)
rgb_dir = os.path.join(dataset_dir, test_config['images'])
# Input GT binary masks dir
gt_mask_dir = os.path.join(dataset_dir, test_config['masks'])
# Input binary mask data
if 'bin_masks' in test_config.keys():
bin_mask_dir = os.path.join(dataset_dir, test_config['bin_masks'])
image_ids = np.arange(indices_arr.size)
##################################################################
# Process each image
##################################################################
for image_id in tqdm(image_ids):
base_name = 'image_{:06d}'.format(indices_arr[image_id])
output_name = 'image_{:06d}'.format(image_id)
rgb_image_fn = os.path.join(rgb_dir, base_name + '.png')
# Run GOP detector
if detector_type == 'gop':
detector = GOP()
elif detector_type == 'mcg':
mcg_dir = os.path.join(dataset_dir, 'mcg', config['mode'])
detector = MCG(mcg_dir, nms_thresh=config['nms_thresh'])
pred_mask = detector.detect(rgb_image_fn)
# Save out ground-truth mask as array of shape (n, h, w)
indiv_gt_masks = []
gt_mask = cv2.imread(os.path.join(gt_mask_dir, base_name +
'.png')).astype(np.uint8)[:, :, 0]
num_gt_masks = np.max(gt_mask)
for i in range(1, num_gt_masks + 1):
indiv_gt_masks.append(gt_mask == i)
gt_mask_output = np.stack(indiv_gt_masks)
np.save(os.path.join(resized_segmask_dir, output_name + '.npy'),
gt_mask_output)
# Set up predicted masks and metadata
indiv_pred_masks = []
r_info = {
'rois': [],
'scores': [],
'class_ids': [],
}
if bin_mask_dir:
mask_im = BinaryImage.open(
os.path.join(bin_mask_dir, base_name + '.png'), 'phoxi')
bin_mask = cv2.resize(mask_im.data,
(pred_mask.shape[1], pred_mask.shape[0]))
# Number of predictions to use (larger number means longer time)
num_pred_masks = min(pred_mask.shape[2], 100)
# num_pred_masks = pred_mask.shape[2]
for i in range(1, num_pred_masks + 1):
# Extract individual mask
indiv_pred_mask = pred_mask[:, :, i - 1]
if not np.any(indiv_pred_mask):
continue
if bin_mask_dir:
inter = np.logical_and(bin_mask, indiv_pred_mask)
frac_overlap = np.sum(inter) / np.sum(indiv_pred_mask)
if frac_overlap <= 0.5:
continue
inter = np.logical_and(indiv_pred_mask,
np.sum(indiv_pred_masks, axis=0))
frac_overlap = np.sum(inter) / np.sum(indiv_pred_mask)
if frac_overlap >= 0.5:
continue
indiv_pred_masks.append(indiv_pred_mask)
# Compute bounding box, score, class_id
nonzero_pix = np.nonzero(indiv_pred_mask)
min_x, max_x = np.min(nonzero_pix[1]), np.max(nonzero_pix[1])
min_y, max_y = np.min(nonzero_pix[0]), np.max(nonzero_pix[0])
r_info['rois'].append([min_y, min_x, max_y, max_x])
if detector.mock_score:
# Generates a meaningful mock score for MCG (first region scores
# highest, etc.)
r_info['scores'].append(-i)
else:
r_info['scores'].append(1.0)
r_info['class_ids'].append(1)
r_info['rois'] = np.array(r_info['rois'])
r_info['scores'] = np.array(r_info['scores'])
r_info['class_ids'] = np.array(r_info['class_ids'])
# Write the predicted masks and metadata
pred_mask_output = np.stack(indiv_pred_masks).astype(
np.uint8) if indiv_pred_masks else np.array([])
np.save(os.path.join(pred_dir, output_name + '.npy'), pred_mask_output)
np.save(os.path.join(pred_info_dir, output_name + '.npy'), r_info)
pred_mask_output = np.stack(indiv_pred_masks).astype(np.uint8)
print('Saved prediction masks to:\t {}'.format(pred_dir))
print('Saved prediction info (bboxes, scores, classes) to:\t {}'.format(
pred_info_dir))
print('Saved transformed GT segmasks to:\t {}'.format(resized_segmask_dir))
return pred_dir, pred_info_dir, resized_segmask_dir
def detect(detector_type, config, run_dir, test_config):
"""Run PCL-based detection on a depth-image-based dataset.
Parameters
----------
config : dict
config for a PCL detector
run_dir : str
Directory to save outputs in. Output will be saved in pred_masks, pred_info,
and modal_segmasks_processed subdirectories.
test_config : dict
config containing dataset information
"""
##################################################################
# Set up output directories
##################################################################
# Create subdirectory for prediction masks
pred_dir = os.path.join(run_dir, 'pred_masks')
mkdir_if_missing(pred_dir)
# Create subdirectory for prediction scores & bboxes
pred_info_dir = os.path.join(run_dir, 'pred_info')
mkdir_if_missing(pred_info_dir)
# Create subdirectory for transformed GT segmasks
resized_segmask_dir = os.path.join(run_dir, 'modal_segmasks_processed')
mkdir_if_missing(resized_segmask_dir)
##################################################################
# Set up input directories
##################################################################
dataset_dir = test_config['path']
indices_arr = np.load(os.path.join(dataset_dir, test_config['indices']))
# Input depth image data (numpy files, not .pngs)
depth_dir = os.path.join(dataset_dir, test_config['images'])
# Input GT binary masks dir
gt_mask_dir = os.path.join(dataset_dir, test_config['masks'])
# Input binary mask data
if 'bin_masks' in test_config.keys():
bin_mask_dir = os.path.join(dataset_dir, test_config['bin_masks'])
# Input camera intrinsics
camera_intrinsics_fn = os.path.join(dataset_dir, 'camera_intrinsics.intr')
camera_intrs = CameraIntrinsics.load(camera_intrinsics_fn)
image_ids = np.arange(indices_arr.size)
##################################################################
# Process each image
##################################################################
for image_id in tqdm(image_ids):
base_name = 'image_{:06d}'.format(indices_arr[image_id])
output_name = 'image_{:06d}'.format(image_id)
depth_image_fn = base_name + '.npy'
# Extract depth image
depth_data = np.load(os.path.join(depth_dir, depth_image_fn))
depth_im = DepthImage(depth_data, camera_intrs.frame)
depth_im = depth_im.inpaint(0.25)
# Mask out bin pixels if appropriate/necessary
if bin_mask_dir:
mask_im = BinaryImage.open(
os.path.join(bin_mask_dir, base_name + '.png'),
camera_intrs.frame)
mask_im = mask_im.resize(depth_im.shape[:2])
depth_im = depth_im.mask_binary(mask_im)
point_cloud = camera_intrs.deproject(depth_im)
point_cloud.remove_zero_points()
pcl_cloud = pcl.PointCloud(point_cloud.data.T.astype(np.float32))
tree = pcl_cloud.make_kdtree()
if detector_type == 'euclidean':
segmentor = pcl_cloud.make_EuclideanClusterExtraction()
segmentor.set_ClusterTolerance(config['tolerance'])
elif detector_type == 'region_growing':
segmentor = pcl_cloud.make_RegionGrowing(ksearch=50)
segmentor.set_NumberOfNeighbours(config['n_neighbors'])
segmentor.set_CurvatureThreshold(config['curvature'])
segmentor.set_SmoothnessThreshold(config['smoothness'])
else:
print('PCL detector type not supported')
exit()
segmentor.set_MinClusterSize(config['min_cluster_size'])
segmentor.set_MaxClusterSize(config['max_cluster_size'])
segmentor.set_SearchMethod(tree)
cluster_indices = segmentor.Extract()
# Set up predicted masks and metadata
indiv_pred_masks = []
r_info = {
'rois': [],
'scores': [],
'class_ids': [],
}
for i, cluster in enumerate(cluster_indices):
points = pcl_cloud.to_array()[cluster]
indiv_pred_mask = camera_intrs.project_to_image(
PointCloud(points.T, frame=camera_intrs.frame)).to_binary()
indiv_pred_mask.data[indiv_pred_mask.data > 0] = 1
indiv_pred_masks.append(indiv_pred_mask.data)
# Compute bounding box, score, class_id
nonzero_pix = np.nonzero(indiv_pred_mask.data)
min_x, max_x = np.min(nonzero_pix[1]), np.max(nonzero_pix[1])
min_y, max_y = np.min(nonzero_pix[0]), np.max(nonzero_pix[0])
r_info['rois'].append([min_y, min_x, max_y, max_x])
r_info['scores'].append(1.0)
r_info['class_ids'].append(1)
r_info['rois'] = np.array(r_info['rois'])
r_info['scores'] = np.array(r_info['scores'])
r_info['class_ids'] = np.array(r_info['class_ids'])
# Write the predicted masks and metadata
if indiv_pred_masks:
pred_mask_output = np.stack(indiv_pred_masks).astype(np.uint8)
else:
pred_mask_output = np.array(indiv_pred_masks).astype(np.uint8)
# Save out ground-truth mask as array of shape (n, h, w)
indiv_gt_masks = []
gt_mask = cv2.imread(os.path.join(gt_mask_dir, base_name + '.png'))
gt_mask = cv2.resize(gt_mask,
(depth_im.shape[1], depth_im.shape[0])).astype(
np.uint8)[:, :, 0]
num_gt_masks = np.max(gt_mask)
for i in range(1, num_gt_masks + 1):
indiv_gt_mask = (gt_mask == i)
if np.any(indiv_gt_mask):
indiv_gt_masks.append(gt_mask == i)
gt_mask_output = np.stack(indiv_gt_masks)
np.save(os.path.join(resized_segmask_dir, output_name + '.npy'),
gt_mask_output)
np.save(os.path.join(pred_dir, output_name + '.npy'), pred_mask_output)
np.save(os.path.join(pred_info_dir, output_name + '.npy'), r_info)
print('Saved prediction masks to:\t {}'.format(pred_dir))
print('Saved prediction info (bboxes, scores, classes) to:\t {}'.format(
pred_info_dir))
print('Saved transformed GT segmasks to:\t {}'.format(resized_segmask_dir))
return pred_dir, pred_info_dir, resized_segmask_dir
def benchmark(config):
"""Benchmarks a model, computes and stores model predictions and then
evaluates them on COCO metrics and supplementary benchmarking script."""
print("Benchmarking model.")
# Create new directory for outputs
output_dir = config['output_dir']
utils.mkdir_if_missing(output_dir)
# Save config in output directory
config.save(os.path.join(output_dir, config['save_conf_name']))
image_shape = config['model']['settings']['image_shape']
config['model']['settings']['image_min_dim'] = min(image_shape)
config['model']['settings']['image_max_dim'] = max(image_shape)
config['model']['settings']['gpu_count'] = 1
config['model']['settings']['images_per_gpu'] = 1
inference_config = MaskConfig(config['model']['settings'])
model_dir, _ = os.path.split(config['model']['path'])
model = modellib.MaskRCNN(mode=config['model']['mode'], config=inference_config,
model_dir=model_dir)
# Load trained weights
print("Loading weights from ", config['model']['path'])
model.load_weights(config['model']['path'], by_name=True)
# Create dataset
test_dataset = ImageDataset(config)
test_dataset.load(config['dataset']['indices'])
test_dataset.prepare()
vis_config = copy(config)
vis_config['dataset']['images'] = 'depth_ims'
vis_config['dataset']['masks'] = 'modal_segmasks'
vis_dataset = ImageDataset(config)
vis_dataset.load(config['dataset']['indices'])
vis_dataset.prepare()
######## BENCHMARK JUST CREATES THE RUN DIRECTORY ########
# code that actually produces outputs should be plug-and-play
# depending on what kind of benchmark function we run.
# If we want to remove bin pixels, pass in the directory with
# those masks.
if config['mask']['remove_bin_pixels']:
bin_mask_dir = os.path.join(config['dataset']['path'], config['mask']['bin_masks'])
overlap_thresh = config['mask']['overlap_thresh']
else:
bin_mask_dir = False
overlap_thresh = 0
# Create predictions and record where everything gets stored.
pred_mask_dir, pred_info_dir, gt_mask_dir = \
detect(config['output_dir'], inference_config, model, test_dataset, bin_mask_dir, overlap_thresh)
ap, ar = coco_benchmark(pred_mask_dir, pred_info_dir, gt_mask_dir)
if config['vis']['predictions']:
visualize_predictions(config['output_dir'], vis_dataset, inference_config, pred_mask_dir, pred_info_dir,
show_bbox=config['vis']['show_bbox_pred'], show_scores=config['vis']['show_scores_pred'], show_class=config['vis']['show_class_pred'])
if config['vis']['ground_truth']:
visualize_gts(config['output_dir'], vis_dataset, inference_config, show_scores=False, show_bbox=config['vis']['show_bbox_gt'], show_class=config['vis']['show_class_gt'])
if config['vis']['s_bench']:
s_benchmark(config['output_dir'], vis_dataset, inference_config, pred_mask_dir, pred_info_dir)
print("Saved benchmarking output to {}.\n".format(config['output_dir']))
return ap, ar
def detect(run_dir, inference_config, model, dataset, bin_mask_dir=False, overlap_thresh=0.5):
"""
Given a run directory, a MaskRCNN config object, a MaskRCNN model object,
and a Dataset object,
- Loads and processes ground-truth masks, saving them to a new directory
for annotation
- Makes predictions on images
- Saves prediction masks in a certain directory
- Saves other prediction info (scores, bboxes) in a separate directory
Returns paths to directories for prediction masks, prediction info, and
modified GT masks.
If bin_mask_dir is specified, then we will be checking predictions against
the "bin-vs-no bin" mask for the test case.
For each predicted instance, if less than overlap_thresh of the mask actually
consists of non-bin pixels, we will toss out the mask.
"""
# Create subdirectory for prediction masks
pred_dir = os.path.join(run_dir, 'pred_masks')
utils.mkdir_if_missing(pred_dir)
# Create subdirectory for prediction scores & bboxes
pred_info_dir = os.path.join(run_dir, 'pred_info')
utils.mkdir_if_missing(pred_info_dir)
# Create subdirectory for transformed GT segmasks
resized_segmask_dir = os.path.join(run_dir, 'modal_segmasks_processed')
utils.mkdir_if_missing(resized_segmask_dir)
# Feed images into model one by one. For each image, predict and save.
image_ids = dataset.image_ids
indices = dataset.indices
times = []
print('MAKING PREDICTIONS')
for image_id in tqdm(image_ids):
# Load image and ground truth data and resize for net
image, _, _, _, gt_mask =\
modellib.load_image_gt(dataset, inference_config, image_id,
use_mini_mask=False)
# Run object detection
results = model.detect([image], verbose=0)
r = results[0]
times.append(r['time'])
# If we choose to mask out bin pixels, load the bin masks and
# transform them properly.
# Then, delete the mask, score, class id, and bbox corresponding
# to each mask that is entirely bin pixels.
if bin_mask_dir:
name = 'image_{:06d}.png'.format(indices[image_id])
bin_mask = io.imread(os.path.join(bin_mask_dir, name))[:,:,np.newaxis]
bin_mask, _, _, _, _ = utilslib.resize_image(
bin_mask,
max_dim=inference_config.IMAGE_MAX_DIM,
min_dim=inference_config.IMAGE_MIN_DIM,
mode=inference_config.IMAGE_RESIZE_MODE
)
bin_mask = bin_mask.squeeze()
deleted_masks = [] # which segmasks are gonna be tossed?
num_detects = r['masks'].shape[2]
for k in range(num_detects):
# compute the area of the overlap.
inter = np.logical_and(bin_mask, r['masks'][:,:,k])
frac_overlap = np.sum(inter) / np.sum(r['masks'][:,:,k])
if frac_overlap <= overlap_thresh:
deleted_masks.append(k)
r['masks'] = [r['masks'][:,:,k] for k in range(num_detects) if k not in deleted_masks]
r['masks'] = np.stack(r['masks'], axis=2) if r['masks'] else np.array([])
r['rois'] = [r['rois'][k,:] for k in range(num_detects) if k not in deleted_masks]
r['rois'] = np.stack(r['rois'], axis=0) if r['rois'] else np.array([])
r['class_ids'] = np.array([r['class_ids'][k] for k in range(num_detects)
if k not in deleted_masks])
r['scores'] = np.array([r['scores'][k] for k in range(num_detects)
if k not in deleted_masks])
# Save copy of transformed GT segmasks to disk in preparation for annotations
mask_name = 'image_{:06d}'.format(image_id)
mask_path = os.path.join(resized_segmask_dir, mask_name)
# save the transpose so it's (n, h, w) instead of (h, w, n)
np.save(mask_path, gt_mask.transpose(2, 0, 1))
# Save masks
save_masks = np.stack([r['masks'][:,:,i] for i in range(r['masks'].shape[2])]) if np.any(r['masks']) else np.array([])
save_masks_path = os.path.join(pred_dir, 'image_{:06d}.npy'.format(image_id))
np.save(save_masks_path, save_masks)
# Save info
r_info = {
'rois': r['rois'],
'scores': r['scores'],
'class_ids': r['class_ids']
}
r_info_path = os.path.join(pred_info_dir, 'image_{:06d}.npy'.format(image_id))
np.save(r_info_path, r_info)
print('Took {} s'.format(sum(times)))
print('Saved prediction masks to:\t {}'.format(pred_dir))
print('Saved prediction info (bboxes, scores, classes) to:\t {}'.format(pred_info_dir))
print('Saved transformed GT segmasks to:\t {}'.format(resized_segmask_dir))
return pred_dir, pred_info_dir, resized_segmask_dir
def detect(run_dir,
inference_config,
model,
dataset,
bin_mask_dir=False,
overlap_thresh=0.5):
"""
Given a run directory, a MaskRCNN config object, a MaskRCNN model object,
and a Dataset object,
- Loads and processes ground-truth masks, saving them to a new directory
for annotation
- Makes predictions on images
- Saves prediction masks in a certain directory
- Saves other prediction info (scores, bboxes) in a separate directory
Returns paths to directories for prediction masks, prediction info, and
modified GT masks.
If bin_mask_dir is specified, then we will be checking predictions against
the "bin-vs-no bin" mask for the test case.
For each predicted instance, if less than overlap_thresh of the mask actually
consists of non-bin pixels, we will toss out the mask.
"""
# Create subdirectory for prediction masks
pred_dir = os.path.join(run_dir, 'pred_masks')
utils.mkdir_if_missing(pred_dir)
# Create subdirectory for prediction scores & bboxes
pred_info_dir = os.path.join(run_dir, 'pred_info')
utils.mkdir_if_missing(pred_info_dir)
# Create subdirectory for transformed GT segmasks
resized_segmask_dir = os.path.join(run_dir, 'gt_annot')
utils.mkdir_if_missing(resized_segmask_dir)
# Feed images into model one by one. For each image, predict and save.
image_ids = dataset.image_ids
indices = dataset.indices
times = []
print('MAKING PREDICTIONS')
for image_id in tqdm(image_ids):
# Load image and ground truth data and resize for net
image, _, _, _, gt_mask =\
modellib.load_image_gt(dataset, inference_config, image_id)
# Run object detection
results = model.detect([image], verbose=0)
r = results[0]
times.append(r['time'])
# Save copy of transformed GT segmasks to disk in preparation for annotations
mask_name = 'image_{:06d}'.format(image_id)
mask_path = os.path.join(resized_segmask_dir, mask_name)
# save the transpose so it's (n, h, w) instead of (h, w, n)
np.save(mask_path, gt_mask.transpose(2, 0, 1))
# Save masks
save_masks = np.stack([
r['masks'][:, :, i] for i in range(r['masks'].shape[2])
]) if np.any(r['masks']) else np.array([])
save_masks_path = os.path.join(pred_dir,
'image_{:06d}.npy'.format(image_id))
np.save(save_masks_path, save_masks)
# Save info
r_info = {
'rois': r['rois'],
'scores': r['scores'],
'class_ids': r['class_ids']
}
r_info_path = os.path.join(pred_info_dir,
'image_{:06d}.npy'.format(image_id))
np.save(r_info_path, r_info)
print('Took {} s'.format(sum(times)))
print('Saved prediction masks to:\t {}'.format(pred_dir))
print('Saved prediction info (bboxes, scores, classes) to:\t {}'.format(
pred_info_dir))
print('Saved transformed GT segmasks to:\t {}'.format(resized_segmask_dir))
return pred_dir, pred_info_dir, resized_segmask_dir
