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