def aggregate(self,
scale_cls_dets,
scale_cls_masks,
vis=False,
cache_name='cache',
vis_path=None,
vis_name=None,
pre_nms_db_divide=10,
vis_ext='.png'):
# TODO: finish the multi process version, fix bug
n_scales = len(scale_cls_dets)
assert n_scales == len(
self.cfg.TEST.VALID_RANGES
), 'A valid range should be specified for each test scale'
all_boxes = [[[] for _ in range(self.num_images)]
for _ in range(self.num_classes)]
all_masks = [[[] for _ in range(self.num_images)]
for _ in range(self.num_classes)]
nms_pool = Pool(32)
if len(scale_cls_dets) > 1:
self.show_info(
'Aggregating detections from multiple scales and applying NMS...'
)
else:
self.show_info('Performing NMS on detections...')
# Apply ranges and store detections per category
parallel_nms_args = [[] for _ in range(pre_nms_db_divide)]
parallel_nms_mask_args = [[] for _ in range(pre_nms_db_divide)]
n_roi_per_pool = math.ceil(self.num_images / float(pre_nms_db_divide))
for i in range(self.num_images):
for j in range(1, self.num_classes):
agg_dets = np.empty((0, 5), dtype=np.float32)
agg_masks = np.empty((0, 28, 28), dtype=np.float32)
for all_cls_dets, all_cls_masks, valid_range in zip(
scale_cls_dets, scale_cls_masks,
self.cfg.TEST.VALID_RANGES):
cls_dets = all_cls_dets[j][i]
cls_masks = all_cls_masks[j][i]
heights = cls_dets[:, 2] - cls_dets[:, 0]
widths = cls_dets[:, 3] - cls_dets[:, 1]
areas = widths * heights
lvalid_ids = np.where(areas > valid_range[0] * valid_range[0])[0] if valid_range[0] > 0 else \
np.arange(len(areas))
uvalid_ids = np.where(areas <= valid_range[1] * valid_range[1])[0] if valid_range[1] > 0 else \
np.arange(len(areas))
valid_ids = np.intersect1d(lvalid_ids, uvalid_ids)
cls_dets = cls_dets[
valid_ids, :] if len(valid_ids) > 0 else cls_dets
cls_masks = cls_masks[
valid_ids, :, :] if len(valid_ids) > 0 else cls_masks
agg_dets = np.vstack((agg_dets, cls_dets))
agg_masks = np.concatenate((agg_masks, cls_masks), axis=0)
parallel_nms_args[int(i / n_roi_per_pool)].append(agg_dets)
parallel_nms_mask_args[int(i /
n_roi_per_pool)].append(agg_masks)
# Divide roidb and perform NMS in parallel to reduce the memory usage
im_offset = 0
for part in tqdm(range(pre_nms_db_divide)):
final_dets = nms_pool.map(self.nms_worker.worker,
parallel_nms_args[part])
n_part_im = int(len(final_dets) / (self.num_classes - 1))
for i in range(n_part_im):
for j in range(1, self.num_classes):
if self.cfg.TEST.NMS < 0:
all_boxes[j][im_offset +
i] = final_dets[i *
(self.num_classes - 1) +
(j - 1)]
# TODO: finish the code for masks
else:
# pdb.set_trace()
# TODO: finish the multi process code
all_boxes[j][im_offset + i] = parallel_nms_args[part][
i * (self.num_classes - 1) +
(j - 1)][final_dets[i * (self.num_classes - 1) +
(j - 1)], :]
all_masks[j][im_offset + i] = parallel_nms_mask_args[
part][i * (self.num_classes - 1) +
(j - 1)][final_dets[i *
(self.num_classes - 1) +
(j - 1)], :]
im_offset += n_part_im
nms_pool.close()
# Limit number of detections to MAX_PER_IMAGE if requested and visualize if vis is True
for i in range(self.num_images):
if self.cfg.TEST.MAX_PER_IMAGE > 0:
image_scores = np.hstack([
all_boxes[j][i][:, -1] for j in range(1, self.num_classes)
])
if len(image_scores) > self.cfg.TEST.MAX_PER_IMAGE:
image_thresh = np.sort(
image_scores)[-self.cfg.TEST.MAX_PER_IMAGE]
for j in range(1, self.num_classes):
keep = np.where(
all_boxes[j][i][:, -1] >= image_thresh)[0]
all_boxes[j][i] = all_boxes[j][i][keep, :]
all_masks[j][i] = all_masks[j][i][keep, :]
if vis:
visualization_path = vis_path if vis_path else os.path.join(
self.cfg.TEST.VISUALIZATION_PATH, cache_name)
if not os.path.isdir(visualization_path):
os.makedirs(visualization_path)
import cv2
im = cv2.cvtColor(cv2.imread(self.roidb[i]['image']),
cv2.COLOR_BGR2RGB)
visualize_masks(
im, [[]] +
[all_boxes[j][i]
for j in range(1, self.num_classes)], [[]] +
[all_masks[j][i] for j in range(1, self.num_classes)],
1.0,
self.cfg.network.PIXEL_MEANS,
self.class_names,
threshold=0.5,
save_path=os.path.join(
visualization_path,
'{}{}'.format(vis_name if vis_name else i, vis_ext)),
transform=False)
if cache_name:
cache_path = os.path.join(self.result_path, cache_name)
if not os.path.isdir(cache_path):
os.makedirs(cache_path)
cache_path = os.path.join(cache_path, 'detections.pkl')
self.show_info(
'Done! Saving detections into: {}'.format(cache_path))
with open(cache_path, 'wb') as detfile:
cPickle.dump(all_boxes, detfile)
return all_boxes
def get_detections(self,
cls_thresh=1e-3,
cache_name='cache',
evaluate=False,
vis=False,
vis_path=None,
vis_ext='.png'):
all_boxes = [
[[] for _ in range(self.num_images)]
for _ in range(self.num_classes)
] # (19,num_images,0) # Probably replace num_images with batch_size?
all_masks = [[[] for _ in range(self.num_images)]
for _ in range(self.num_classes)] # (19,num_images,0)
data_counter = 0
detect_time, post_time = 0, 0
if vis:
visualization_path = vis_path if vis_path else os.path.join(
self.cfg.TEST.VISUALIZATION_PATH, cache_name)
if vis and not os.path.isdir(self.cfg.TEST.VISUALIZATION_PATH):
os.makedirs(self.cfg.TEST.VISUALIZATION_PATH)
for batch in self.test_iter:
# pdb.set_trace()
im_info = batch.data[1].asnumpy()
scales = im_info[:, 2].reshape(-1, self.batch_size)
# Run detection on the batch
stime = time.time()
scores, boxes, data, im_ids, masks = self.detect(batch, scales)
detect_time += time.time() - stime
stime = time.time()
for i, (cscores, cboxes, im_id,
cmasks) in enumerate(zip(scores, boxes, im_ids, masks)):
parallel_nms_args = []
masks_list = []
for j in range(1, self.num_classes):
# Apply the score threshold
inds = np.where(cscores[:, j] > cls_thresh)[0]
rem_scores = cscores[inds, j, np.newaxis]
rem_boxes = cboxes[inds, 0:4]
# pdb.set_trace()
cls_masks = cmasks[inds, 1, :, :]
# rem_masks = cmasks[inds, :]
cls_dets = np.hstack((rem_boxes, rem_scores))
if evaluate or vis:
parallel_nms_args.append(cls_dets)
masks_list.append(cls_masks)
else:
all_boxes[j][im_id] = cls_dets
all_masks[j][im_id] = cls_masks
# Apply nms
if evaluate or vis:
if not self.thread_pool:
self.thread_pool = ThreadPool(8)
if self.cfg.TEST.NMS < 0:
# TODO: finish vis code for mask
final_dets = self.thread_pool.map(
self.nms_worker.worker, parallel_nms_args)
pdb.set_trace()
for j in range(1, self.num_classes):
all_boxes[j][im_id] = final_dets[j - 1]
else:
keeps = self.thread_pool.map(self.nms_worker.worker,
parallel_nms_args)
for j in range(1, self.num_classes):
all_boxes[j][im_id] = all_boxes[j][im_id][
keeps[j - 1], :]
all_masks[j][im_id] = all_masks[j][im_id][
keeps[j - 1], :]
# Filter boxes based on max_per_image if needed
if evaluate and self.cfg.TEST.MAX_PER_IMAGE:
print(im_id)
image_scores = np.hstack([
all_boxes[j][im_id][:, -1]
for j in range(1, self.num_classes)
])
if len(image_scores) > self.cfg.TEST.MAX_PER_IMAGE:
image_thresh = np.sort(
image_scores)[-self.cfg.TEST.MAX_PER_IMAGE]
for j in range(1, self.num_classes):
keep = np.where(
all_boxes[j][im_id][:, -1] >= image_thresh)[0]
all_boxes[j][im_id] = all_boxes[j][im_id][keep, :]
all_masks[j][im_id] = all_masks[j][im_id][keep, :]
if vis:
# TODO: finish vis code
if not os.path.isdir(visualization_path):
os.makedirs(visualization_path)
# visualize_dets(batch.data[0][i].asnumpy(),
# [[]] + [all_boxes[j][im_id] for j in range(1, self.num_classes)], im_info[i, 2],
# self.cfg.network.PIXEL_MEANS, self.class_names, threshold=0.5,
# save_path=os.path.join(visualization_path, '{}{}'.format(im_id, vis_ext)))
# cv2.imwrite(os.path.join(visualization_path, '{}{}'.format(im_id, vis_ext)), batch.data[0][i].asnumpy())
visualize_masks(batch.data[0][i].asnumpy(), [[]] + [
all_boxes[j][im_id]
for j in range(1, self.num_classes)
], [[]] + [
all_masks[j][im_id]
for j in range(1, self.num_classes)
],
im_info[i, 2],
self.cfg.network.PIXEL_MEANS,
self.class_names,
threshold=0.5,
save_path=os.path.join(
visualization_path,
'{}{}'.format(im_id, vis_ext)))
data_counter += self.test_iter.get_batch_size()
post_time += time.time() - stime
if self.verbose:
self.show_info(
'Tester: {}/{}, Detection: {:.4f}s, Post Processing: {:.4}s'
.format(min(data_counter, self.num_images),
self.num_images, detect_time / data_counter,
post_time / data_counter))
if self.thread_pool:
self.thread_pool.close()
return all_boxes, all_masks
def aggregateSingle(self,
scale_cls_dets,
scale_cls_masks,
vis=False,
cache_name='cache',
vis_path=None,
vis_name=None,
vis_ext='.png'):
n_scales = len(scale_cls_dets)
assert n_scales == len(
self.cfg.TEST.VALID_RANGES
), 'A valid range should be specified for each test scale'
all_boxes = [[[] for _ in range(self.num_images)]
for _ in range(self.num_classes)]
all_masks = [[[] for _ in range(self.num_images)]
for _ in range(self.num_classes)]
if len(scale_cls_dets) > 1:
self.show_info(
'Aggregating detections from multiple scales and applying NMS...'
)
else:
self.show_info('Performing NMS on detections...')
# TODO: change the hard code here, change it to soft_nms or mask_nms
nms = py_nms_wrapper(0.3)
# nms = gpu_nms_wrapper(0.3, 0)
# Apply ranges and store detections per category
for i in range(self.num_images):
for j in range(1, self.num_classes):
agg_dets = np.empty((0, 5), dtype=np.float32)
agg_masks = np.empty((0, 28, 28), dtype=np.float32)
for all_cls_dets, all_cls_masks, valid_range in zip(
scale_cls_dets, scale_cls_masks,
self.cfg.TEST.VALID_RANGES):
cls_dets = all_cls_dets[j][i]
cls_masks = all_cls_masks[j][i]
heights = cls_dets[:, 2] - cls_dets[:, 0]
widths = cls_dets[:, 3] - cls_dets[:, 1]
areas = widths * heights
lvalid_ids = np.where(areas > valid_range[0]*valid_range[0])[0] if valid_range[0] > 0 else \
np.arange(len(areas))
uvalid_ids = np.where(areas <= valid_range[1]*valid_range[1])[0] if valid_range[1] > 0 else \
np.arange(len(areas))
valid_ids = np.intersect1d(lvalid_ids, uvalid_ids)
cls_dets = cls_dets[
valid_ids, :] if len(valid_ids) > 0 else cls_dets
cls_masks = cls_masks[
valid_ids, :, :] if len(valid_ids) > 0 else cls_masks
# pdb.set_trace()
agg_dets = np.vstack(
(agg_dets, cls_dets.astype(np.float32)))
# pdb.set_trace()
agg_masks = np.concatenate((agg_masks, cls_masks), axis=0)
# start = timeit.default_timer()
keep = nms(agg_dets)
# stop = timeit.default_timer()
# print 'nms time: ', stop - start
all_boxes[j][i] = agg_dets[keep, :]
all_masks[j][i] = agg_masks[keep, :]
# parallel_nms_args[int(i/n_roi_per_pool)].append(agg_dets)
# Divide roidb and perform NMS in parallel to reduce the memory usage
# TODO: change to multi process later
# Limit number of detections to MAX_PER_IMAGE if requested and visualize if vis is True
for i in range(self.num_images):
if self.cfg.TEST.MAX_PER_IMAGE > 0:
image_scores = np.hstack([
all_boxes[j][i][:, -1] for j in range(1, self.num_classes)
])
if len(image_scores) > self.cfg.TEST.MAX_PER_IMAGE:
image_thresh = np.sort(
image_scores)[-self.cfg.TEST.MAX_PER_IMAGE]
for j in range(1, self.num_classes):
keep = np.where(
all_boxes[j][i][:, -1] >= image_thresh)[0]
all_boxes[j][i] = all_boxes[j][i][keep, :]
all_masks[j][i] = all_masks[j][i][keep, :]
if vis:
visualization_path = vis_path if vis_path else os.path.join(
self.cfg.TEST.VISUALIZATION_PATH, cache_name)
if not os.path.isdir(visualization_path):
os.makedirs(visualization_path)
import cv2
im = cv2.cvtColor(cv2.imread(self.roidb[i]['image']),
cv2.COLOR_BGR2RGB)
visualize_masks(
im, [[]] +
[all_boxes[j][i]
for j in range(1, self.num_classes)], [[]] +
[all_masks[j][i] for j in range(1, self.num_classes)],
1.0,
self.cfg.network.PIXEL_MEANS,
self.class_names,
threshold=0.5,
save_path=os.path.join(
visualization_path,
'{}{}'.format(vis_name if vis_name else i, vis_ext)),
transform=False)
if cache_name:
cache_path = os.path.join(self.result_path, cache_name)
if not os.path.isdir(cache_path):
os.makedirs(cache_path)
cache_path = os.path.join(cache_path, 'detections.pkl')
self.show_info(
'Done! Saving detections into: {}'.format(cache_path))
with open(cache_path, 'wb') as detfile:
cPickle.dump(all_boxes, detfile)
return all_boxes, all_masks