def inference(self, image_names): # get data data = self.load_data_and_get_predictor(image_names) # test for idx, im_name in enumerate(image_names): data_batch = mx.io.DataBatch(data=[data[idx]], label=[], pad=0, index=idx, provide_data=[ [(k, v.shape) for k, v in zip(self.data_names, data[idx])]], provide_label=[None]) scales = [data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data))] tic() scores, boxes, data_dict = im_detect( self.predictor, data_batch, self.data_names, scales, config) boxes = boxes[0].astype('f') scores = scores[0].astype('f') dets_nms = [] for j in range(1, scores.shape[1]): cls_scores = scores[:, j, np.newaxis] cls_boxes = boxes[:, 4:8] if config.CLASS_AGNOSTIC else boxes[:, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = self.nms(cls_dets) cls_dets = cls_dets[keep, :] cls_dets = cls_dets[cls_dets[:, -1] > 0.6, :] dets_nms.append(cls_dets) print 'testing {} {:.4f}s'.format(im_name, toc()) # visualize #im = cv2.imread(cur_path + '/../demo/' + im_name) im = cv2.imread(im_name) im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) self.show_save_boxes(im, dets_nms, self.classes, im_name, 1) #show_boxes(im, dets_nms, self.classes, 1) pass
def net_inference(self, msgs): assert isinstance(msgs, pb.ForwardMsgs) start = time.time() start_forward = time.time() assert len(msgs.msgs) == 1 arr = cPickle.loads(msgs.msgs[0].network_input_buf) data = [[mx.nd.array(arr['im_array']), mx.nd.array(arr['im_info'])]] data_batch = mx.io.DataBatch(data=data, label=[None], provide_data=arr['data_shapes'], provide_label=[None]) with self.lock: # https://github.com/ataraxialab/Deformable-ConvNets/blob/master/rfcn/core/tester.py#L124 scores, boxes, _ = im_detect(self.predictor, data_batch, ['data', 'im_info'], arr['im_scale'], config) end_forward = time.time() msgs_out = [] for _ in range(len(msgs.msgs)): output = { 'scores': scores[0], 'boxes': boxes[0], } msgs_out.append( pb.ForwardMsg(network_output_buf=cPickle.dumps( output, protocol=cPickle.HIGHEST_PROTOCOL))) log.info('{} use time {}, forward time {}, batch_size: {}/{}'.format( self.app_name, time.time() - start, end_forward - start_forward, len(msgs.msgs), self.batch_size)) return pb.ForwardMsgs(msgs=msgs_out)
def net_inference(model, reqs): CTX.logger.info("inference begin...") # datas = json.loads(args) predictor = model['predictor'] classes_dict = model['labels']['class'] # threshold uses for default threshold_dict = model['labels']['minMt'] # minModelThreshold rets = [] nms = py_nms_wrapper(config.TEST.NMS) box_voting = py_box_voting_wrapper(config.TEST.BOX_VOTING_IOU_THRESH, config.TEST.BOX_VOTING_SCORE_THRESH, with_nms=True) try: for data in reqs: try: im = load_image(data['data']['uri'], body=data['data']['body']) except ErrorBase as e: rets.append({"code":e.code, "message": e.message, "result": None}) continue # return [], 400, 'load image error' if im.shape[0] > im.shape[1]: long_side, short_side = im.shape[0], im.shape[1] else: long_side, short_side = im.shape[1], im.shape[0] if short_side > 0 and float(long_side)/float(short_side) > 50.0: msg = "aspect ration is too large, long_size:short_side should not larger than 50.0" # raise ErrorBase.__init__(400, msg) rets.append({"code": 400, "message": msg, "result": None}) continue data_batch, data_names, im_scale = generate_batch(im) scores, boxes, data_dict = im_detect(predictor, data_batch, data_names, im_scale, config) det_ret = [] # labels.csv file not include background for cls_index in sorted(classes_dict.keys()): cls_ind = cls_index cls_name = classes_dict.get(cls_ind) cls_boxes = boxes[0][:, 4:8] if config.CLASS_AGNOSTIC else boxes[0][:, 4 * cls_ind:4 *4 * (cls_ind + 1)] cls_scores = scores[0][:, cls_ind, np.newaxis] threshold = float(threshold_dict[cls_ind]) keep = np.where(cls_scores > threshold)[0] dets = np.hstack((cls_boxes, cls_scores)).astype(np.float32)[keep, :] keep = nms(dets) det_ret.extend(_build_result(det, cls_name, cls_ind, model['labels']) for det in dets[keep, :]) # get review value rets.append(dict(code=0,message='',result=dict(detections=det_ret))) except Exception as e: # print(traceback.format_exc()) CTX.logger.info("inference error:%s"%(traceback.format_exc())) return [], 599, str(e) return rets, 0, ''
def main(im_name, frame, score): data = [] # only resize input image to target size and return scale im, im_scale = resize(frame, target_size, max_size, stride=config.network.IMAGE_STRIDE) im_tensor = transform(im, config.network.PIXEL_MEANS) im_info = np.array([[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) data.append({'data': im_tensor, 'im_info': im_info}) data = [[mx.nd.array(data[i][name]) for name in data_names] for i in xrange(len(data))] data_batch = mx.io.DataBatch(data=[data[0]], label=[], pad=0, index=0, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[0])]], provide_label=[None]) scales = [data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data))] tic() scores, boxes, masks, data_dict = im_detect(predictor, data_batch, data_names, scales, config) im_shapes = [data_batch.data[i][0].shape[2:4] for i in xrange(len(data_batch.data))] masks = masks[0][:, 1:, :, :] im_height = np.round(im_shapes[0][0] / scales[0]).astype('int') im_width = np.round(im_shapes[0][1] / scales[0]).astype('int') # print (im_height, im_width) boxes = clip_boxes(boxes[0], (im_height, im_width)) result_masks, result_dets = cpu_mask_voting(masks, boxes, scores[0], num_classes, 100, im_width, im_height, config.TEST.NMS, config.TEST.MASK_MERGE_THRESH, config.BINARY_THRESH) dets = [result_dets[j] for j in range(1, num_classes)] masks = [result_masks[j][:, 0, :, :] for j in range(1, num_classes)] print 'testing {} {:.4f}s'.format(im_name, toc()) min_confidence = score # visualize for i in xrange(len(dets)): keep = np.where(dets[i][:,-1] > min_confidence) dets[i] = dets[i][keep] masks[i] = masks[i][keep] ''' dets: [ numpy.ndarray([[x1 y1 x2 y2 score]]) for j in classes ] [x1 y1] - upper left corner of object [x2 y2] - lower right corner of object masks: [ numpy.ndarray([21, 21]) for j in classes ] confidence that pixel belongs to object ''' for i in range(len(dets)): if len(dets[i]) > 0: for j in range(len(dets[i])): print('{name}: {score} ({loc})'.format(name = classes[i], score = dets[i][j][-1], loc = dets[i][j][:-1].tolist()))
def fcis_seg(image, classes, predictor, args): num_classes = len(classes) + 1 data = [] target_size = config.SCALES[0][0] max_size = config.SCALES[0][1] im, im_scale = resize(image, target_size, max_size, stride=config.network.IMAGE_STRIDE) im_tensor = transform(im, config.network.PIXEL_MEANS) start = time.time() im_info = np.array([[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) data.append({'data': im_tensor, 'im_info': im_info}) data_names = ['data', 'im_info'] data = [[mx.nd.array(data[i][name]) for name in data_names] for i in xrange(len(data))] data_batch = mx.io.DataBatch(data=[data[0]], label=[], pad=0, index=0, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[0])]], provide_label=[None]) scales = [data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data))] scores, boxes, masks, data_dict = im_detect(predictor, data_batch, data_names, scales, config) im_shapes = [data_batch.data[i][0].shape[2:4] for i in xrange(len(data_batch.data))] if not config.TEST.USE_MASK_MERGE: all_boxes = [[] for _ in xrange(num_classes)] all_masks = [[] for _ in xrange(num_classes)] nms = py_nms_wrapper(config.TEST.NMS) for j in range(1, num_classes): indexes = np.where(scores[0][:, j] > 0.7)[0] cls_scores = scores[0][indexes, j, np.newaxis] cls_masks = masks[0][indexes, 1, :, :] try: if config.CLASS_AGNOSTIC: cls_boxes = boxes[0][indexes, :] else: raise Exception() except: cls_boxes = boxes[0][indexes, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = nms(cls_dets) all_boxes[j] = cls_dets[keep, :] all_masks[j] = cls_masks[keep, :] dets = [all_boxes[j] for j in range(1, num_classes)] masks = [all_masks[j] for j in range(1, num_classes)] else: masks = masks[0][:, 1:, :, :] im_height = np.round(im_shapes[0][0] / scales[0]).astype('int') im_width = np.round(im_shapes[0][1] / scales[0]).astype('int') # print (im_height, im_width) boxes = clip_boxes(boxes[0], (im_height, im_width)) result_masks, result_dets = gpu_mask_voting(masks, boxes, scores[0], num_classes, 100, im_width, im_height, config.TEST.NMS, config.TEST.MASK_MERGE_THRESH, config.BINARY_THRESH, 0) dets = [result_dets[j] for j in range(1, num_classes)] masks = [result_masks[j][:, 0, :, :] for j in range(1, num_classes)] cods, bimsks, names = decode_mask(im, dets, masks, classes, config, args) return cods, bimsks, names
def inference(predictor, data_batch, data_names, num_classes, BINARY_THRESH = 0.4, CONF_THRESH=0.7, gpu_id=0): scales = [data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data))] im_shapes = [data_batch.data[i][0].shape[2:4] for i in xrange(len(data_batch.data))] scores, boxes, masks, data_dict = im_detect(predictor, data_batch, data_names, scales, config) if not config.TEST.USE_MASK_MERGE: all_boxes = [[] for _ in xrange(num_classes)] all_masks = [[] for _ in xrange(num_classes)] nms = py_nms_wrapper(config.TEST.NMS) for j in range(1, num_classes): indexes = np.where(scores[0][:, j] > CONF_THRESH)[0] cls_scores = scores[0][indexes, j, np.newaxis] cls_masks = masks[0][indexes, 1, :, :] # try: # if config.CLASS_AGNOSTIC: # cls_boxes = boxes[0][indexes, :] # else: # raise Exception() # except: if config.CLASS_AGNOSTIC: cls_boxes = boxes[0][indexes, :] else: cls_boxes = boxes[0][indexes, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = nms(cls_dets) all_boxes[j] = cls_dets[keep, :] all_masks[j] = cls_masks[keep, :] dets = [all_boxes[j] for j in range(1, num_classes)] masks = [all_masks[j] for j in range(1, num_classes)] else: masks = masks[0][:, 1:, :, :] im_height = np.round(im_shapes[0][0] / scales[0]).astype('int') im_width = np.round(im_shapes[0][1] / scales[0]).astype('int') # print (im_height, im_width) boxes = clip_boxes(boxes[0], (im_height, im_width)) result_masks, result_dets = gpu_mask_voting(masks, boxes, scores[0], num_classes, 100, im_width, im_height, config.TEST.NMS, config.TEST.MASK_MERGE_THRESH, BINARY_THRESH, gpu_id) dets = [result_dets[j] for j in range(1, num_classes)] masks = [result_masks[j][:, 0, :, :] for j in range(1, num_classes)] for i in xrange(len(dets)): keep = np.where(dets[i][:,-1] > CONF_THRESH) dets[i] = dets[i][keep] masks[i] = masks[i][keep] return dets, masks
def warm_up(self): image_names = ['COCO_test2015_000000000891.jpg', 'COCO_test2015_000000001669.jpg'] data = self.load_data_and_get_predictor(image_names) # warm up for j in xrange(2): data_batch = mx.io.DataBatch(data=[data[0]], label=[], pad=0, index=0, provide_data=[ [(k, v.shape) for k, v in zip(self.data_names, data[0])]], provide_label=[None]) scales = [data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data))] scores, boxes, data_dict = im_detect( self.predictor, data_batch, self.data_names, scales, config)
def predict(self, im): im, im_scale = resize(im, self.scales[0], self.scales[1], stride=config.network.IMAGE_STRIDE) im_tensor = transform(im, config.network.PIXEL_MEANS) im_info = np.array( [[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) data = {'data': im_tensor, 'im_info': im_info} data = [[ mx.nd.array(data[self.data_names[0]]), mx.nd.array(data[self.data_names[1]]) ]] data_batch = mx.io.DataBatch( data=data, label=[], pad=0, index=0, provide_data=[[(k, v.shape) for k, v in zip(self.data_names, data[0])]], provide_label=[None]) scores, boxes, data_dict = im_detect(self.predictor, data_batch, self.data_names, [im_scale], config) boxes = boxes[0].astype('f') scores = scores[0].astype('f') dets_nms = [] for j in range(1, scores.shape[1]): cls_scores = scores[:, j, np.newaxis] cls_boxes = boxes[:, 4:8] if config.CLASS_AGNOSTIC else boxes[:, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = self.nms(cls_dets) cls_dets = cls_dets[keep, :] cls_dets = cls_dets[cls_dets[:, -1] > 0.7, :] dets_nms.append(cls_dets) res = {} for idx, cls in enumerate(self.classes): res['{}'.format(cls)] = dets_nms[idx].tolist() logging.debug("Predictions: {}".format(res)) return res
def init_net(self): config = self.cfg # get symbol sym_instance = resnet_v1_101_fcis() sym = sym_instance.get_symbol(config, is_train=False) # key parameters data_names = ['data', 'im_info'] label_names = [] max_data_shape = [[('data', (1, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES])))]] target_size = config.SCALES[0][0] max_size = config.SCALES[0][1] self.data_batch_wr = DataBatchWrapper(target_size, max_size, image_stride=config.network.IMAGE_STRIDE, pixel_means=config.network.PIXEL_MEANS, data_names=data_names, label_names=label_names) im = np.zeros((target_size,max_size,3)) data_tensor_info = self.data_batch_wr.get_data_tensor_info(im) # get predictor arg_params, aux_params = load_param_file(self.model_path, process=True) print("\nLoaded model %s\n"%(self.model_path)) self.net = Predictor(sym, data_names, label_names, context=[mx.gpu(self.ctx_id[0])], max_data_shapes=max_data_shape, provide_data=[[(k, v.shape) for k, v in zip(data_names, data_tensor_info)]], provide_label=[None], arg_params=arg_params, aux_params=aux_params) self.data_names = data_names # # warm up predictor for i in xrange(2): data_batch = self.data_batch_wr.get_data_batch(im) scales = [data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data))] _, _, _, _ = im_detect(self.net, data_batch, data_names, scales, config)
def predict_url(cfg,predictor,image_set_index,nms,box_voting,all_boxes,jsonlistwithscore ,jsonlistwithoutscore, thresh, vis=False, use_box_voting=False,num_gpu=1): """ generate data_batch -> im_detect -> post process :param predictor: Predictor :param image_name: image name :param vis: will save as a new image if not visualized :return: None """ # image_file = image_path_from_index(index, dataset_path, image_set) for idx in range(0,len(image_set_index)/num_gpu): import urllib im = [] for j in range(0,num_gpu): image_file = image_set_index[num_gpu * idx + j] proxies = {'http': 'http://xsio.qiniu.io'} data = urllib.urlopen(image_file.strip(), proxies=proxies).read() nparr = np.fromstring(data, np.uint8) im.append(cv2.imdecode(nparr, 1)) # im = cv2.imread(image_file)/ box_list =[] data_batch, data_names, im_scale = generate_batch_V2(im,num_gpu) scores, boxes, data_dict = im_detect(predictor, data_batch, data_names, im_scale, config) for cls in CLASSES: cls_ind = CLASSES.index(cls) cls_boxes = boxes[0][:, 4:8] if config.CLASS_AGNOSTIC else boxes[0][:, 4 * cls_ind:4 * (cls_ind + 1)] cls_scores = scores[0][:, cls_ind, np.newaxis] keep = np.where(cls_scores >= thresh)[0] cls_dets = np.hstack((cls_boxes, cls_scores)).astype(np.float32)[keep, :] keep = nms(cls_dets) # apply box voting after nms if use_box_voting: nms_cls_dets = cls_dets[keep, :] all_boxes[cls_ind][idx] = box_voting(nms_cls_dets, cls_dets) else: all_boxes[cls_ind][idx] = cls_dets[keep, :] boxes_this_image = [[]] + [box_list[j] for j in xrange(1, len(CLASSES))] dets, dets_s = [], [] for j in xrange(1, len(CLASSES)): if CLASSES[j] == 'not terror': continue boxes = box_list[j] for box in boxes: det_score = box[4] if det_score > thresh: det, det_s = dict(), dict() xmin = float(box[0]) ymin = float(box[1]) xmax = float(box[2]) ymax = float(box[3]) det['pts'] = [[xmin, ymin], [xmax, ymin], [xmax, ymax], [xmin, ymax]] det['class'] = CLASSES[j] det_s['pts'] = [[xmin, ymin], [xmax, ymin], [xmax, ymax], [xmin, ymax]] det_s['class'] = CLASSES[j] det_s['score'] = float(det_score) # det['score'] = float(det_score) # det['bbox'] = [float(box[0]),float(box[1]),float(box[2]),float(box[3])] # det['class'] = CLASSES[j] dets.append(det) dets_s.append(det_s) # line = {} # line['detections'] = dets # line['img'] = image_file ress = { "url": image_file, "label": {"detect": {"general_d": {"bbox": dets}}}, "type": "image", "source_url": "", "ops": "download()" } ress_s = { "url": image_file, "label": {"detect": {"general_d": {"bbox": dets_s}}}, "type": "image", "source_url": "", "ops": "download()" } if vis: # vis_all_detection(data_dict['data'].oasnumpy(), boxes_this_image, CLASSES, im_scale) save_vis(image_file, data_dict, boxes_this_image, im_scale, cfg) jsonlistwithscore[idx] = json.dumps(ress_s) jsonlistwithoutscore[idx] = json.dumps(ress) return
def predict(self, images, feat_output, aggr_feat_output): model = self.model all_frame_interval = self.all_frame_interval feat_sym = self.feat_sym aggr_sym = self.aggr_sym num_classes = self.num_classes classes = self.classes max_per_image = self.max_per_image output_dir = cur_path + '/../demo/rfcn_fgfa_{}/'.format(self.index) self.index += 1 if not os.path.exists(output_dir): os.makedirs(output_dir) data = [] for im in images: target_size = cfg.SCALES[0][0] max_size = cfg.SCALES[0][1] im, im_scale = resize(im, target_size, max_size, stride=cfg.network.IMAGE_STRIDE) im_tensor = transform(im, cfg.network.PIXEL_MEANS) im_info = np.array( [[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) feat_stride = float(cfg.network.RCNN_FEAT_STRIDE) data.append({ 'data': im_tensor, 'im_info': im_info, 'data_cache': im_tensor, 'feat_cache': im_tensor }) # get predictor print 'get-predictor' data_names = ['data', 'im_info', 'data_cache', 'feat_cache'] label_names = [] t1 = time.time() data = [[mx.nd.array(data[i][name]) for name in data_names] for i in xrange(len(data))] max_data_shape = [[ ('data', (1, 3, max([v[0] for v in cfg.SCALES]), max([v[1] for v in cfg.SCALES]))), ('data_cache', (11, 3, max([v[0] for v in cfg.SCALES]), max([v[1] for v in cfg.SCALES]))), ('feat_cache', ((11, cfg.network.FGFA_FEAT_DIM, np.ceil(max([v[0] for v in cfg.SCALES]) / feat_stride).astype( np.int), np.ceil(max([v[1] for v in cfg.SCALES]) / feat_stride).astype( np.int)))) ]] provide_data = [[(k, v.shape) for k, v in zip(data_names, data[i])] for i in xrange(len(data))] provide_label = [None for _ in xrange(len(data))] arg_params, aux_params = load_param(cur_path + model, 0, process=True) feat_predictors = Predictor(feat_sym, data_names, label_names, context=[mx.gpu(0)], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) aggr_predictors = Predictor(aggr_sym, data_names, label_names, context=[mx.gpu(0)], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) nms = py_nms_wrapper(cfg.TEST.NMS) # First frame of the video idx = 0 data_batch = mx.io.DataBatch( data=[data[idx]], label=[], pad=0, index=idx, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[idx])]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] all_boxes = [[[] for _ in range(len(data))] for _ in range(num_classes)] data_list = deque(maxlen=all_frame_interval) feat_list = deque(maxlen=all_frame_interval) image, feat = get_resnet_output(feat_predictors, data_batch, data_names) # append cfg.TEST.KEY_FRAME_INTERVAL padding images in the front (first frame) while len(data_list) < cfg.TEST.KEY_FRAME_INTERVAL: data_list.append(image) feat_list.append(feat) vis = False file_idx = 0 thresh = 1e-3 for idx, element in enumerate(data): data_batch = mx.io.DataBatch( data=[element], label=[], pad=0, index=idx, provide_data=[[(k, v.shape) for k, v in zip(data_names, element)]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] if (idx != len(data) - 1): if len(data_list) < all_frame_interval - 1: image, feat = get_resnet_output(feat_predictors, data_batch, data_names) data_list.append(image) feat_list.append(feat) else: ################################################# # main part of the loop ################################################# image, feat = get_resnet_output(feat_predictors, data_batch, data_names) data_list.append(image) feat_list.append(feat) prepare_data(data_list, feat_list, data_batch) pred_result, aggr_feat = im_detect(aggr_predictors, data_batch, data_names, scales, cfg) assert len(aggr_feat) == 1 data_batch.data[0][-2] = None data_batch.provide_data[0][-2] = ('data_cache', None) data_batch.data[0][-1] = None data_batch.provide_data[0][-1] = ('feat_cache', None) out_im = process_pred_result( classes, pred_result, num_classes, thresh, cfg, nms, all_boxes, file_idx, max_per_image, vis, data_list[cfg.TEST.KEY_FRAME_INTERVAL].asnumpy(), scales) total_time = time.time() - t1 if (cfg.TEST.SEQ_NMS == False): save_image(output_dir, file_idx, out_im) print 'testing {} {:.4f}s'.format( str(file_idx) + '.JPEG', total_time / (file_idx + 1)) file_idx += 1 else: ################################################# # end part of a video # ################################################# end_counter = 0 image, feat = get_resnet_output(feat_predictors, data_batch, data_names) while end_counter < cfg.TEST.KEY_FRAME_INTERVAL + 1: data_list.append(image) feat_list.append(feat) prepare_data(data_list, feat_list, data_batch) pred_result, aggr_feat = im_detect(aggr_predictors, data_batch, data_names, scales, cfg) assert len(aggr_feat) == 1 out_im = process_pred_result( classes, pred_result, num_classes, thresh, cfg, nms, all_boxes, file_idx, max_per_image, vis, data_list[cfg.TEST.KEY_FRAME_INTERVAL].asnumpy(), scales) total_time = time.time() - t1 if (cfg.TEST.SEQ_NMS == False): save_image(output_dir, file_idx, out_im) print 'testing {} {:.4f}s'.format( str(file_idx) + '.JPEG', total_time / (file_idx + 1)) file_idx += 1 end_counter += 1 if (cfg.TEST.SEQ_NMS): video = [all_boxes[j][:] for j in range(1, num_classes)] dets_all = seq_nms(video) for cls_ind, dets_cls in enumerate(dets_all): for frame_ind, dets in enumerate(dets_cls): keep = nms(dets) all_boxes[cls_ind + 1][frame_ind] = dets[keep, :] for idx in range(len(data)): boxes_this_image = [[]] + [ all_boxes[j][idx] for j in range(1, num_classes) ] out_im = draw_all_detection(data[idx][0].asnumpy(), boxes_this_image, classes, scales[0], cfg) save_image(output_dir, idx, out_im) print 'done'
def main(): # get symbol pprint.pprint(config) config.symbol = 'resnet_v1_101_flownet_rfcn' model = '/../model/rfcn_dff_flownet_vid' sym_instance = eval(config.symbol + '.' + config.symbol)() key_sym = sym_instance.get_key_test_symbol(config) cur_sym = sym_instance.get_cur_test_symbol(config) # set up class names num_classes = 31 classes = [ 'airplane', 'antelope', 'bear', 'bicycle', 'bird', 'bus', 'car', 'cattle', 'dog', 'domestic_cat', 'elephant', 'fox', 'giant_panda', 'hamster', 'horse', 'lion', 'lizard', 'monkey', 'motorcycle', 'rabbit', 'red_panda', 'sheep', 'snake', 'squirrel', 'tiger', 'train', 'turtle', 'watercraft', 'whale', 'zebra' ] # load demo data image_names = glob.glob(cur_path + '/../demo/ILSVRC2015_val_00007010/*.JPEG') output_dir = cur_path + '/../demo/rfcn_dff/' if not os.path.exists(output_dir): os.makedirs(output_dir) key_frame_interval = 10 # data = [] key_im_tensor = None for idx, im_name in enumerate(image_names): assert os.path.exists(im_name), ('%s does not exist'.format(im_name)) im = cv2.imread(im_name, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION) target_size = config.SCALES[0][0] max_size = config.SCALES[0][1] im, im_scale = resize(im, target_size, max_size, stride=config.network.IMAGE_STRIDE) im_tensor = transform(im, config.network.PIXEL_MEANS) im_info = np.array( [[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) if idx % key_frame_interval == 0: key_im_tensor = im_tensor data.append({ 'data': im_tensor, 'im_info': im_info, 'data_key': key_im_tensor, 'feat_key': np.zeros((1, config.network.DFF_FEAT_DIM, 1, 1)) }) # get predictor data_names = ['data', 'im_info', 'data_key', 'feat_key'] label_names = [] data = [[mx.nd.array(data[i][name]) for name in data_names] for i in range(len(data))] max_data_shape = [[ ('data', (1, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES]))), ('data_key', (1, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES]))), ]] provide_data = [[(k, v.shape) for k, v in zip(data_names, data[i])] for i in range(len(data))] provide_label = [None for i in range(len(data))] arg_params, aux_params = load_param(cur_path + model, 0, process=True) key_predictor = Predictor(key_sym, data_names, label_names, context=[mx.gpu(0)], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) cur_predictor = Predictor(cur_sym, data_names, label_names, context=[mx.gpu(0)], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) nms = gpu_nms_wrapper(config.TEST.NMS, 0) # warm up for j in range(2): data_batch = mx.io.DataBatch(data=[data[j]], label=[], pad=0, index=0, provide_data=[[ (k, v.shape) for k, v in zip(data_names, data[j]) ]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in range(len(data_batch.data)) ] if j % key_frame_interval == 0: scores, boxes, data_dict, feat = im_detect(key_predictor, data_batch, data_names, scales, config) else: data_batch.data[0][-1] = feat data_batch.provide_data[0][-1] = ('feat_key', feat.shape) scores, boxes, data_dict, _ = im_detect(cur_predictor, data_batch, data_names, scales, config) print("warmup done") # test time = 0 count = 0 for idx, im_name in enumerate(image_names): data_batch = mx.io.DataBatch( data=[data[idx]], label=[], pad=0, index=idx, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[idx])]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in range(len(data_batch.data)) ] tic() if idx % key_frame_interval == 0: scores, boxes, data_dict, feat = im_detect(key_predictor, data_batch, data_names, scales, config) else: data_batch.data[0][-1] = feat data_batch.provide_data[0][-1] = ('feat_key', feat.shape) scores, boxes, data_dict, _ = im_detect(cur_predictor, data_batch, data_names, scales, config) time += toc() count += 1 print('testing {} {:.4f}s'.format(im_name, time / count)) boxes = boxes[0].astype('f') scores = scores[0].astype('f') dets_nms = [] for j in range(1, scores.shape[1]): cls_scores = scores[:, j, np.newaxis] cls_boxes = boxes[:, 4:8] if config.CLASS_AGNOSTIC else boxes[:, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = nms(cls_dets) cls_dets = cls_dets[keep, :] cls_dets = cls_dets[cls_dets[:, -1] > 0.7, :] dets_nms.append(cls_dets) # visualize im = cv2.imread(im_name) im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) # show_boxes(im, dets_nms, classes, 1) out_im = draw_boxes(im, dets_nms, classes, 1) _, filename = os.path.split(im_name) cv2.imwrite(output_dir + filename, out_im) print('done')
def main(): # get symbol pprint.pprint(config) config.symbol = 'resnet_v1_101_rfcn' model = '/../model/rfcn_vid' sym_instance = eval(config.symbol + '.' + config.symbol)() sym = sym_instance.get_test_symbol(config) # set up class names num_classes = 31 classes = ['airplane', 'antelope', 'bear', 'bicycle', 'bird', 'bus', 'car', 'cattle', 'dog', 'domestic_cat', 'elephant', 'fox', 'giant_panda', 'hamster', 'horse', 'lion', 'lizard', 'monkey', 'motorcycle', 'rabbit', 'red_panda', 'sheep', 'snake', 'squirrel', 'tiger', 'train', 'turtle', 'watercraft', 'whale', 'zebra'] # load demo data image_names = glob.glob(cur_path + '/../demo/ILSVRC2015_val_00007010/*.JPEG') output_dir = cur_path + '/../demo/rfcn/' if not os.path.exists(output_dir): os.makedirs(output_dir) # data = [] for im_name in image_names: assert os.path.exists(im_name), ('%s does not exist'.format(im_name)) im = cv2.imread(im_name, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION) target_size = config.SCALES[0][0] max_size = config.SCALES[0][1] im, im_scale = resize(im, target_size, max_size, stride=config.network.IMAGE_STRIDE) im_tensor = transform(im, config.network.PIXEL_MEANS) im_info = np.array([[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) data.append({'data': im_tensor, 'im_info': im_info}) # get predictor data_names = ['data', 'im_info'] label_names = [] data = [[mx.nd.array(data[i][name]) for name in data_names] for i in xrange(len(data))] max_data_shape = [[('data', (1, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES])))]] provide_data = [[(k, v.shape) for k, v in zip(data_names, data[i])] for i in xrange(len(data))] provide_label = [None for i in xrange(len(data))] arg_params, aux_params = load_param(cur_path + model, 0, process=True) predictor = Predictor(sym, data_names, label_names, context=[mx.gpu(0)], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) nms = gpu_nms_wrapper(config.TEST.NMS, 0) # warm up for j in xrange(2): data_batch = mx.io.DataBatch(data=[data[0]], label=[], pad=0, index=0, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[0])]], provide_label=[None]) scales = [data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data))] scores, boxes, data_dict = im_detect(predictor, data_batch, data_names, scales, config) # test time = 0 count = 0 for idx, im_name in enumerate(image_names): data_batch = mx.io.DataBatch(data=[data[idx]], label=[], pad=0, index=idx, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[idx])]], provide_label=[None]) scales = [data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data))] tic() scores, boxes, data_dict = im_detect(predictor, data_batch, data_names, scales, config) time += toc() count += 1 print 'testing {} {:.4f}s'.format(im_name, time/count) boxes = boxes[0].astype('f') scores = scores[0].astype('f') dets_nms = [] for j in range(1, scores.shape[1]): cls_scores = scores[:, j, np.newaxis] cls_boxes = boxes[:, 4:8] if config.CLASS_AGNOSTIC else boxes[:, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = nms(cls_dets) cls_dets = cls_dets[keep, :] cls_dets = cls_dets[cls_dets[:, -1] > 0.7, :] dets_nms.append(cls_dets) # visualize im = cv2.imread(im_name) im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) # show_boxes(im, dets_nms, classes, 1) out_im = draw_boxes(im, dets_nms, classes, 1) _, filename = os.path.split(im_name) cv2.imwrite(output_dir + filename,out_im) print 'done'
def main(): # get symbol ctx_id = [int(i) for i in config.gpus.split(',')] pprint.pprint(config) sym_instance = eval(config.symbol)() sym = sym_instance.get_symbol(config, is_train=False) # set up class names num_classes = 31 classes =['alphabet_of_the_magi', 'anglo-saxon_futhorc', 'arcadian', 'armenian', 'asomtavruli_(georgian)', 'balinese', 'bengali', 'blackfoot_(canadian_aboriginal_syllabics)', 'braille', 'burmese_(myanmar)', 'cyrillic', 'early_aramaic', 'futurama', 'grantha', 'greek', 'gujarati', 'hebrew', 'inuktitut_(canadian_aboriginal_syllabics)', 'japanese_(hiragana)', 'japanese_(katakana)', 'korean', 'latin', 'malay_(jawi_-_arabic)', 'mkhedruli_(georgian)', 'n_ko', 'ojibwe_(canadian_aboriginal_syllabics)', 'sanskrit', 'syriac_(estrangelo)', 'tagalog', 'tifinagh'] # load demo data image_names = ['train_000000.jpg', 'train_000001.jpg', 'train_000002.jpg', 'val_000000.jpg', 'val_000001.jpg', 'val_000002.jpg'] data = [] for im_name in image_names: assert os.path.exists(cur_path + '/../demo/' + im_name), ('%s does not exist'.format('../demo/' + im_name)) im = cv2.imread(cur_path + '/../demo/' + im_name, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION) target_size = config.SCALES[0][0] max_size = config.SCALES[0][1] im, im_scale = resize(im, target_size, max_size, stride=config.network.IMAGE_STRIDE) im_tensor = transform(im, config.network.PIXEL_MEANS) im_info = np.array([[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) data.append({'data': im_tensor, 'im_info': im_info}) # get predictor data_names = ['data', 'im_info'] label_names = [] data = [[mx.nd.array(data[i][name]) for name in data_names] for i in xrange(len(data))] max_data_shape = [[('data', (1, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES])))]] provide_data = [[(k, v.shape) for k, v in zip(data_names, data[i])] for i in xrange(len(data))] provide_label = [None for i in xrange(len(data))] arg_params, aux_params = load_param(cur_path + '/../model/fcis_voc', 0, process=True) predictor = Predictor(sym, data_names, label_names, context=[mx.gpu(ctx_id[0])], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) # warm up for i in xrange(2): data_batch = mx.io.DataBatch(data=[data[0]], label=[], pad=0, index=0, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[0])]], provide_label=[None]) scales = [data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data))] _, _, _, _ = im_detect(predictor, data_batch, data_names, scales, config) # test for idx, im_name in enumerate(image_names): data_batch = mx.io.DataBatch(data=[data[idx]], label=[], pad=0, index=idx, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[idx])]], provide_label=[None]) scales = [data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data))] tic() scores, boxes, masks, data_dict = im_detect(predictor, data_batch, data_names, scales, config) im_shapes = [data_batch.data[i][0].shape[2:4] for i in xrange(len(data_batch.data))] if not config.TEST.USE_MASK_MERGE: all_boxes = [[] for _ in xrange(num_classes)] all_masks = [[] for _ in xrange(num_classes)] nms = py_nms_wrapper(config.TEST.NMS) for j in range(1, num_classes): indexes = np.where(scores[0][:, j] > 0.7)[0] cls_scores = scores[0][indexes, j, np.newaxis] cls_masks = masks[0][indexes, 1, :, :] try: if config.CLASS_AGNOSTIC: cls_boxes = boxes[0][indexes, :] else: raise Exception() except: cls_boxes = boxes[0][indexes, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = nms(cls_dets) all_boxes[j] = cls_dets[keep, :] all_masks[j] = cls_masks[keep, :] dets = [all_boxes[j] for j in range(1, num_classes)] masks = [all_masks[j] for j in range(1, num_classes)] else: masks = masks[0][:, 1:, :, :] im_height = np.round(im_shapes[0][0] / scales[0]).astype('int') im_width = np.round(im_shapes[0][1] / scales[0]).astype('int') print (im_height, im_width) boxes = clip_boxes(boxes[0], (im_height, im_width)) result_masks, result_dets = gpu_mask_voting(masks, boxes, scores[0], num_classes, 100, im_width, im_height, config.TEST.NMS, config.TEST.MASK_MERGE_THRESH, config.BINARY_THRESH, ctx_id[0]) dets = [result_dets[j] for j in range(1, num_classes)] masks = [result_masks[j][:, 0, :, :] for j in range(1, num_classes)] print 'testing {} {:.4f}s'.format(im_name, toc()) # visualize for i in xrange(len(dets)): keep = np.where(dets[i][:,-1]>0.7) dets[i] = dets[i][keep] masks[i] = masks[i][keep] im = cv2.imread(cur_path + '/../demo/' + im_name) im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) show_masks(im, dets, masks, classes, config) print 'done'
def process_image_fun(imagesPath=None, fileOp=None, vis=None, model_params_list=None, count=0): # init rfcn dcn detect model (mxnet) # model_params_list = init_detect_model() # num_classes = RFCN_DCN_CONFIG['num_classes'] # 0 is background, classes = RFCN_DCN_CONFIG['num_classes_name_list'] min_threshold = min(list( RFCN_DCN_CONFIG['need_label_thresholds'].values())) im_name = imagesPath all_can_read_image = [] data = [] all_can_read_image.append(im_name) target_size = config.SCALES[0][0] max_size = config.SCALES[0][1] im, im_scale = resize(im_name, target_size, max_size, stride=config.network.IMAGE_STRIDE) im_tensor = transform(im, config.network.PIXEL_MEANS) im_info = np.array([[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) data.append({'data': im_tensor, 'im_info': im_info}) # get predictor data_names = ['data', 'im_info'] label_names = [] data = [[mx.nd.array(data[i][name]) for name in data_names] for i in xrange(len(data))] max_data_shape = [[('data', (1, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES])))]] provide_data = [[(k, v.shape) for k, v in zip(data_names, data[i])] for i in xrange(len(data))] provide_label = [None for i in xrange(len(data))] predictor = Predictor(model_params_list[0], data_names, label_names, context=[mx.gpu(1)], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=model_params_list[1], aux_params=model_params_list[2]) nms = gpu_nms_wrapper(config.TEST.NMS, 0) for idx, im_name in enumerate(all_can_read_image): data_batch = mx.io.DataBatch( data=[data[idx]], label=[], pad=0, index=idx, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[idx])]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] tic() scores, boxes, data_dict = im_detect(predictor, data_batch, data_names, scales, config) boxes = boxes[0].astype('f') scores = scores[0].astype('f') dets_nms = [] for j in range(1, scores.shape[1]): cls_scores = scores[:, j, np.newaxis] cls_boxes = boxes[:, 4:8] if config.CLASS_AGNOSTIC else boxes[:, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = nms(cls_dets) cls_dets = cls_dets[keep, :] cls_dets = cls_dets[cls_dets[:, -1] > min_threshold, :] dets_nms.append(cls_dets) #print('testing {} {:.4f}s'.format(im_name, toc())) im = show_boxes_write_rg(im=im_name, dets=dets_nms, classes=classes, scale=1, vis=vis, fileOp=fileOp, count=count) return im
def run_detection(im_root, result_root, conf_threshold): # get symbol pprint.pprint(config) config.symbol = 'resnet_v1_101_rfcn_dcn' if not args.rfcn_only else 'resnet_v1_101_rfcn' sym_instance = eval(config.symbol + '.' + config.symbol)() sym = sym_instance.get_symbol(config, is_train=False) # set up class names num_classes = 81 classes = [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush' ] print('detection in {}'.format(im_root)) im_names = sorted(os.listdir(im_root)) # get predictor data_names = ['data', 'im_info'] label_names = [] data = [] for idx, im_name in enumerate(im_names[:2]): im_file = os.path.join(im_root, im_name) im = cv2.imread(im_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION) target_size = config.SCALES[0][0] max_size = config.SCALES[0][1] im, im_scale = resize(im, target_size, max_size, stride=config.network.IMAGE_STRIDE) im_tensor = transform(im, config.network.PIXEL_MEANS) im_info = np.array( [[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) data.append({'data': im_tensor, 'im_info': im_info}) data = [[mx.nd.array(data[i][name]) for name in data_names] for i in xrange(len(data))] max_data_shape = [[('data', (1, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES])))]] provide_data = [[(k, v.shape) for k, v in zip(data_names, data[i])] for i in xrange(len(data))] provide_label = [None for i in xrange(len(data))] arg_params, aux_params = load_param( cur_path + '/../model/' + ('rfcn_dcn_coco' if not args.rfcn_only else 'rfcn_coco'), config.TEST.test_epoch, process=True) predictor = Predictor(sym, data_names, label_names, context=[mx.gpu(0)], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) # nms = gpu_nms_wrapper(config.TEST.NMS, 0) # nms = soft_nms_wrapper(config.TEST.NMS, method=2) nms = gpu_soft_nms_wrapper(config.TEST.NMS, method=2, device_id=0) nms_t = Timer() for idx, im_name in enumerate(im_names): im_file = os.path.join(im_root, im_name) im = cv2.imread(im_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION) origin_im = im.copy() target_size = config.SCALES[0][0] max_size = config.SCALES[0][1] im, im_scale = resize(im, target_size, max_size, stride=config.network.IMAGE_STRIDE) im_tensor = transform(im, config.network.PIXEL_MEANS) im_info = np.array( [[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) # input data = [mx.nd.array(im_tensor), mx.nd.array(im_info)] data_batch = mx.io.DataBatch(data=[data], label=[], pad=0, index=idx, provide_data=[[ (k, v.shape) for k, v in zip(data_names, data) ]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] tic() scores, boxes, data_dict = im_detect(predictor, data_batch, data_names, scales, config) boxes = boxes[0].astype('f') scores = scores[0].astype('f') dets_nms = [] for j in range(1, scores.shape[1]): cls_scores = scores[:, j, np.newaxis] cls_boxes = boxes[:, 4:8] if config.CLASS_AGNOSTIC else boxes[:, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) nms_t.tic() keep = nms(cls_dets) nms_t.toc() cls_dets = cls_dets[keep, :] cls_dets = cls_dets[cls_dets[:, -1] > 0.7, :] dets_nms.append(cls_dets) print 'testing {} {:.2f}ms'.format(im_name, toc() * 1000) print 'nms: {:.2f}ms'.format(nms_t.total_time * 1000) nms_t.clear() # save results person_dets = dets_nms[0] with open(os.path.join(result_root, '{:04d}.txt'.format(idx)), 'w') as f: f.write('{}\n'.format(len(person_dets))) for det in person_dets: x1, y1, x2, y2, s = det w = x2 - x1 h = y2 - y1 f.write('0 {} {} {} {} {}\n'.format(s, w, h, x1, y1)) # visualize im = origin_im # im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) im = show_boxes_cv2(im, dets_nms, classes, 1) cv2.imshow('det', im) cv2.waitKey(1)
def main(): # get symbol ctx_id = [int(i) for i in config.gpus.split(',')] pprint.pprint(config) sym_instance = eval(config.symbol)() sym = sym_instance.get_symbol(config, is_train=False) # set up class names # set up class names num_classes = 81 classes = [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush' ] # load demo data image_names = [] names_dirs = os.listdir(cur_path + '/../' + test_dir) for im_name in names_dirs: if im_name[-4:] == '.jpg' or im_name[-4:] == '.png': image_names.append(im_name) data = [] for im_name in image_names: assert os.path.exists(cur_path + '/../' + test_dir + im_name), ( '%s does not exist'.format('../demo/' + im_name)) im = cv2.imread(cur_path + '/../' + test_dir + im_name, cv2.IMREAD_COLOR | long(128)) target_size = config.SCALES[0][0] max_size = config.SCALES[0][1] #print "before scale: " #print im.shape im, im_scale = resize(im, target_size, max_size, stride=config.network.IMAGE_STRIDE) #print "after scale: " #print im.shape #im_scale = 1.0 #print "scale ratio: " #print im_scale im_tensor = transform(im, config.network.PIXEL_MEANS) #print im_tensor.shape im_info = np.array( [[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) data.append({'data': im_tensor, 'im_info': im_info}) # get predictor data_names = ['data', 'im_info'] label_names = [] data = [[mx.nd.array(data[i][name]) for name in data_names] for i in xrange(len(data))] max_data_shape = [[('data', (1, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES])))]] provide_data = [[(k, v.shape) for k, v in zip(data_names, data[i])] for i in xrange(len(data))] provide_label = [None for i in xrange(len(data))] arg_params, aux_params = load_param(cur_path + '/../' + model_dir, 0, process=True) predictor = Predictor(sym, data_names, label_names, context=[mx.gpu(ctx_id[0])], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) # warm up for i in xrange(2): data_batch = mx.io.DataBatch(data=[data[0]], label=[], pad=0, index=0, provide_data=[[ (k, v.shape) for k, v in zip(data_names, data[0]) ]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] _, _, _, _ = im_detect(predictor, data_batch, data_names, scales, config) # test for idx, im_name in enumerate(image_names): data_batch = mx.io.DataBatch( data=[data[idx]], label=[], pad=0, index=idx, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[idx])]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] tic() scores, boxes, masks, data_dict = im_detect(predictor, data_batch, data_names, [1.0], config) im_shapes = [ data_batch.data[i][0].shape[2:4] for i in xrange(len(data_batch.data)) ] #print im_shapes if not config.TEST.USE_MASK_MERGE: all_boxes = [[] for _ in xrange(num_classes)] all_masks = [[] for _ in xrange(num_classes)] nms = py_nms_wrapper(config.TEST.NMS) for j in range(1, num_classes): indexes = np.where(scores[0][:, j] > 0.7)[0] cls_scores = scores[0][indexes, j, np.newaxis] cls_masks = masks[0][indexes, 1, :, :] try: if config.CLASS_AGNOSTIC: cls_boxes = boxes[0][indexes, :] else: raise Exception() except: cls_boxes = boxes[0][indexes, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = nms(cls_dets) all_boxes[j] = cls_dets[keep, :] all_masks[j] = cls_masks[keep, :] dets = [all_boxes[j] for j in range(1, num_classes)] masks = [all_masks[j] for j in range(1, num_classes)] else: masks = masks[0][:, 1:, :, :] result_masks, result_dets = gpu_mask_voting( masks, boxes[0], scores[0], num_classes, 100, im_shapes[0][1], im_shapes[0][0], config.TEST.NMS, config.TEST.MASK_MERGE_THRESH, config.BINARY_THRESH, ctx_id[0]) dets = [result_dets[j] for j in range(1, num_classes)] masks = [ result_masks[j][:, 0, :, :] for j in range(1, num_classes) ] print 'testing {} {:.4f}s'.format(im_name, toc()) # visualize for i in xrange(len(dets)): keep = np.where(dets[i][:, -1] > 0.7) dets[i] = dets[i][keep] masks[i] = masks[i][keep] im = cv2.imread(cur_path + '/../' + test_dir + im_name) im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) show_masks(im, dets, masks, classes, config, 1.0 / scales[0], False) # Save img cv2.imwrite(cur_path + '/../' + result_dir + im_name, cv2.cvtColor(im, cv2.COLOR_BGR2RGB)) print 'done'
def process_one_batch_images_fun(isUrlFlag=False, one_batch_images_list=None, init_model_param=None, fileOp=None, vis=False): # init_model_param list : [sym, arg_params, aux_params] num_classes = 11 # 0 is background, # classes = ['tibetan flag', 'guns', 'knives', # 'not terror', 'islamic flag', 'isis flag'] classes = [ 'islamic flag', 'isis flag', 'tibetan flag', 'knives_true', 'guns_true', 'knives_false', 'knives_kitchen', 'guns_anime', 'guns_tools', 'not terror' ] image_names = one_batch_images_list if len(image_names) <= 0: return all_can_read_image = [] data = [] for im_name in image_names: #print("process : %s"%(im_name)) im = readImage_fun(isUrlFlag=isUrlFlag, imagePath=im_name) # 判断 这个图片是否可读 if np.shape(im) == (): print("ReadImageError : %s" % (im_name)) continue all_can_read_image.append(im_name) target_size = config.SCALES[0][0] max_size = config.SCALES[0][1] im, im_scale = resize(im, target_size, max_size, stride=config.network.IMAGE_STRIDE) im_tensor = transform(im, config.network.PIXEL_MEANS) im_info = np.array( [[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) data.append({'data': im_tensor, 'im_info': im_info}) # get predictor data_names = ['data', 'im_info'] label_names = [] data = [[mx.nd.array(data[i][name]) for name in data_names] for i in xrange(len(data))] max_data_shape = [[('data', (1, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES])))]] provide_data = [[(k, v.shape) for k, v in zip(data_names, data[i])] for i in xrange(len(data))] provide_label = [None for i in xrange(len(data))] predictor = Predictor(init_model_param[0], data_names, label_names, context=[mx.gpu(int(args.gpuId))], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=init_model_param[1], aux_params=init_model_param[2]) nms = gpu_nms_wrapper(config.TEST.NMS, 0) for idx, im_name in enumerate(all_can_read_image): data_batch = mx.io.DataBatch( data=[data[idx]], label=[], pad=0, index=idx, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[idx])]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] tic() scores, boxes, data_dict = im_detect(predictor, data_batch, data_names, scales, config) boxes = boxes[0].astype('f') scores = scores[0].astype('f') dets_nms = [] for j in range(1, scores.shape[1]): cls_scores = scores[:, j, np.newaxis] cls_boxes = boxes[:, 4:8] if config.CLASS_AGNOSTIC else boxes[:, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = nms(cls_dets) cls_dets = cls_dets[keep, :] cls_dets = cls_dets[cls_dets[:, -1] > args.threshold, :] dets_nms.append(cls_dets) print('testing {} {:.4f}s'.format(im_name, toc())) show_boxes(isUrlFlag=isUrlFlag, im_name=im_name, dets=dets_nms, classes=classes, scale=1, vis=vis, fileOp=fileOp, flag=args.outputFileFlag) print('process one batch images done') pass
def main(): # get symbol ctx_id = [int(i) for i in config.gpus.split(',')] pprint.pprint(config) sym_instance = eval(config.symbol)() sym = sym_instance.get_symbol(config, is_train=False) # set up class names num_classes = 81 classes = [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush' ] # load demo data image_names = [ 'COCO_test2015_000000000275.jpg', 'COCO_test2015_000000001412.jpg', 'COCO_test2015_000000073428.jpg', 'COCO_test2015_000000393281.jpg' ] data = [] for im_name in image_names: assert os.path.exists(cur_path + '/../demo/' + im_name), ( '%s does not exist'.format('../demo/' + im_name)) im = cv2.imread(cur_path + '/../demo/' + im_name, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION) target_size = config.SCALES[0][0] max_size = config.SCALES[0][1] im, im_scale = resize(im, target_size, max_size, stride=config.network.IMAGE_STRIDE) im_tensor = transform(im, config.network.PIXEL_MEANS) im_info = np.array( [[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) data.append({'data': im_tensor, 'im_info': im_info}) # get predictor data_names = ['data', 'im_info'] label_names = [] data = [[mx.nd.array(data[i][name]) for name in data_names] for i in xrange(len(data))] max_data_shape = [[('data', (1, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES])))]] provide_data = [[(k, v.shape) for k, v in zip(data_names, data[i])] for i in xrange(len(data))] provide_label = [None for i in xrange(len(data))] arg_params, aux_params = load_param(cur_path + '/../model/fcis_coco', convert=True, 0, process=True) predictor = Predictor(sym, data_names, label_names, context=[mx.cpu()], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) # test for idx, im_name in enumerate(image_names): data_batch = mx.io.DataBatch( data=[data[idx]], label=[], pad=0, index=idx, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[idx])]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] tic() scores, boxes, masks, data_dict = im_detect(predictor, data_batch, data_names, scales, config) im_shapes = [ data_batch.data[i][0].shape[2:4] for i in xrange(len(data_batch.data)) ] if not config.TEST.USE_MASK_MERGE: all_boxes = [[] for _ in xrange(num_classes)] all_masks = [[] for _ in xrange(num_classes)] nms = py_nms_wrapper(config.TEST.NMS) for j in range(1, num_classes): indexes = np.where(scores[0][:, j] > 0.7)[0] cls_scores = scores[0][indexes, j, np.newaxis] cls_masks = masks[0][indexes, 1, :, :] try: if config.CLASS_AGNOSTIC: cls_boxes = boxes[0][indexes, :] else: raise Exception() except: cls_boxes = boxes[0][indexes, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = nms(cls_dets) all_boxes[j] = cls_dets[keep, :] all_masks[j] = cls_masks[keep, :] dets = [all_boxes[j] for j in range(1, num_classes)] masks = [all_masks[j] for j in range(1, num_classes)] else: masks = masks[0][:, 1:, :, :] im_height = np.round(im_shapes[0][0] / scales[0]).astype('int') im_width = np.round(im_shapes[0][1] / scales[0]).astype('int') print(im_height, im_width) boxes = clip_boxes(boxes[0], (im_height, im_width)) result_masks, result_dets = cpu_mask_voting( masks, boxes, scores[0], num_classes, 100, im_width, im_height, config.TEST.NMS, config.TEST.MASK_MERGE_THRESH, config.BINARY_THRESH) dets = [result_dets[j] for j in range(1, num_classes)] masks = [ result_masks[j][:, 0, :, :] for j in range(1, num_classes) ] print 'testing {} {:.4f}s'.format(im_name, toc()) # visualize for i in xrange(len(dets)): keep = np.where(dets[i][:, -1] > 0.7) dets[i] = dets[i][keep] masks[i] = masks[i][keep] for i in range(len(dets)): if len(dets[i]) > 0: for j in range(len(dets[i])): print('{name}: {score} ({loc})'.format( name=classes[i], score=dets[i][j][-1], loc=dets[i][j][:-1].tolist())) print 'done'
def Seg(self): for i in xrange(2): data_batch = mx.io.DataBatch( data=[self.data], label=[], pad=0, index=0, provide_data=[[(k, v.shape) for k, v in zip(self.data_names, self.data)]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] _, _, _, _ = im_detect(self.predictor, data_batch, self.data_names, scales, config) data_batch = mx.io.DataBatch( data=[self.data], label=[], pad=0, index=0, provide_data=[[(k, v.shape) for k, v in zip(self.data_names, self.data)]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] tic() scores, boxes, masks, data_dict = im_detect(self.predictor, data_batch, self.data_names, scales, config) #print masks #right im_shapes = [ data_batch.data[i][0].shape[2:4] for i in xrange(len(data_batch.data)) ] if not config.TEST.USE_MASK_MERGE: all_boxes = [[] for _ in xrange(self.num_classes)] all_masks = [[] for _ in xrange(self.num_classes)] nms = py_nms_wrapper(config.TEST.NMS) for j in range(1, self.num_classes): indexes = np.where(scores[0][:, j] > 0.7)[0] cls_scores = scores[0][indexes, j, np.newaxis] cls_masks = masks[0][indexes, 1, :, :] try: if config.CLASS_AGNOSTIC: cls_boxes = boxes[0][indexes, :] else: raise Exception() except: cls_boxes = boxes[0][indexes, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = nms(cls_dets) all_boxes[j] = cls_dets[keep, :] all_masks[j] = cls_masks[keep, :] dets = [all_boxes[j] for j in range(1, self.num_classes)] masks = [all_masks[j] for j in range(1, self.num_classes)] else: masks = masks[0][:, 1:, :, :] im_height = np.round(im_shapes[0][0] / scales[0]).astype('int') im_width = np.round(im_shapes[0][1] / scales[0]).astype('int') print(im_height, im_width) boxes = clip_boxes(boxes[0], (im_height, im_width)) result_masks, result_dets = gpu_mask_voting( masks, boxes, scores[0], self.num_classes, 100, im_width, im_height, config.TEST.NMS, config.TEST.MASK_MERGE_THRESH, config.BINARY_THRESH, self.ctx_id[0]) dets = [result_dets[j] for j in range(1, self.num_classes)] masks = [ result_masks[j][:, 0, :, :] for j in range(1, self.num_classes) ] for i in xrange(1, len(dets)): keep = np.where(dets[i][:, -1] > 1) dets[i] = dets[i][keep] masks[i] = masks[i][keep] keep = np.where(dets[0][:, -1] > 0.8) dets[0] = dets[0][keep] masks[0] = masks[0][keep] newmask = show_masks(self.fg, dets, masks, self.classes, config) #!!!!!!!! wrong mask self.result = newmask return newmask
def demo_net(predictor, dataset, image_set, root_path, dataset_path, thresh, vis=False, vis_image_dir='vis', use_box_voting=False): """ generate data_batch -> im_detect -> post process :param predictor: Predictor :param image_name: image name :param vis: will save as a new image if not visualized :return: None """ # visualization nms = py_nms_wrapper(config.TEST.NMS) box_voting = py_box_voting_wrapper(config.TEST.BOX_VOTING_IOU_THRESH, config.TEST.BOX_VOTING_SCORE_THRESH, with_nms=True) image_set_index_file = os.path.join(dataset_path, 'DET', 'ImageSets', 'DET', image_set + '.txt') assert os.path.exists( image_set_index_file), image_set_index_file + ' not found' with open(image_set_index_file) as f: image_set_index = [x.strip().split(' ')[0] for x in f.readlines()] num_images = len(image_set_index) num_classes = len(CLASSES) all_boxes = [[[] for _ in xrange(num_images)] for _ in xrange(num_classes)] i = 0 for index in image_set_index: image_file = image_path_from_index(index, dataset_path, image_set) print("processing {}/{} image:{}".format(i, num_images, image_file)) im = cv2.imread(image_file) data_batch, data_names, im_scale = generate_batch(im) scores, boxes, data_dict = im_detect(predictor, data_batch, data_names, im_scale, config) for cls in CLASSES: cls_ind = CLASSES.index(cls) #print cls_ind, 4 * cls_ind, 4 * (cls_ind + 1), boxes[0], boxes[0][:, 4 * cls_ind:4 * (cls_ind + 1)] #cls_boxes = boxes[0][:, 4 * cls_ind:4 * (cls_ind + 1)] cls_boxes = boxes[0][:, 4:8] if config.CLASS_AGNOSTIC else boxes[ 0][:, 4 * cls_ind:4 * (cls_ind + 1)] cls_scores = scores[0][:, cls_ind, np.newaxis] keep = np.where(cls_scores >= thresh)[0] cls_dets = np.hstack( (cls_boxes, cls_scores)).astype(np.float32)[keep, :] keep = nms(cls_dets) # apply box voting after nms if use_box_voting: nms_cls_dets = cls_dets[keep, :] all_boxes[cls_ind][i] = box_voting(nms_cls_dets, cls_dets) else: all_boxes[cls_ind][i] = cls_dets[keep, :] boxes_this_image = [[]] + [ all_boxes[j][i] for j in xrange(1, len(CLASSES)) ] i += 1 if vis: #vis_all_detection(data_dict['data'].asnumpy(), boxes_this_image, CLASSES, im_scale) if not os.path.exists(vis_image_dir): os.mkdir(vis_image_dir) result_file = os.path.join( vis_image_dir, index.strip().split('/')[-1] + '_result' + '.JPEG') print('results saved to %s' % result_file) im = draw_all_detection(data_dict['data'].asnumpy(), boxes_this_image, CLASSES, im_scale) im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) cv2.imwrite(result_file, im) print("num of images: detection:{}, gt:{}".format(len(all_boxes[0]), num_images)) #assert len(all_boxes) == num_images, 'calculations not complete' # save results cache_folder = os.path.join(root_path, 'cache') if not os.path.exists(cache_folder): os.mkdir(cache_folder) cache_file = os.path.join(cache_folder, dataset + '_' + image_set + '_detections.pkl') with open(cache_file, 'wb') as f: cPickle.dump(all_boxes, f, cPickle.HIGHEST_PROTOCOL)
def main(tempFileList, fileOp): # get symbol pprint.pprint(config) config.symbol = 'resnet_v1_101_rfcn_dcn' sym_instance = eval(config.symbol + '.' + config.symbol)() sym = sym_instance.get_symbol(config, is_train=False) out_dir = os.path.join( cur_path, 'demo/output/terror-det-rg-data-output/terror-det-v0.9-test/JPEGImages' ) if not os.path.exists(out_dir): os.makedirs(out_dir) # set up class names num_classes = 7 classes = [ 'tibetan flag', 'guns', 'knives', 'not terror', 'islamic flag', 'isis flag' ] # load demo data image_names = tempFileList data = [] for im_name in image_names: im_file = im_name print(im_file) im = cv2.imread(im_file, cv2.IMREAD_COLOR) target_size = config.SCALES[0][0] max_size = config.SCALES[0][1] im, im_scale = resize(im, target_size, max_size, stride=config.network.IMAGE_STRIDE) im_tensor = transform(im, config.network.PIXEL_MEANS) im_info = np.array( [[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) data.append({'data': im_tensor, 'im_info': im_info}) # get predictor data_names = ['data', 'im_info'] label_names = [] data = [[mx.nd.array(data[i][name]) for name in data_names] for i in xrange(len(data))] max_data_shape = [[('data', (1, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES])))]] provide_data = [[(k, v.shape) for k, v in zip(data_names, data[i])] for i in xrange(len(data))] provide_label = [None for i in xrange(len(data))] arg_params, aux_params = load_param(cur_path + '/demo/models/' + ('rfcn_voc'), 10, process=True) #modify by zxt #mx.model.save_checkpoint('f1/final', 10, sym, arg_params, aux_params) predictor = Predictor(sym, data_names, label_names, context=[mx.gpu(0)], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) nms = gpu_nms_wrapper(config.TEST.NMS, 0) # warm up for j in xrange(2): data_batch = mx.io.DataBatch(data=[data[0]], label=[], pad=0, index=0, provide_data=[[ (k, v.shape) for k, v in zip(data_names, data[0]) ]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] scores, boxes, data_dict = im_detect(predictor, data_batch, data_names, scales, config) # test # fileOp = open(os.path.join(cur_path, 'terror-det-rg-test-result.txt'), 'w') fileOp = fileOp for idx, im_name in enumerate(image_names): print("begining process %s" % (im_name)) data_batch = mx.io.DataBatch( data=[data[idx]], label=[], pad=0, index=idx, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[idx])]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] tic() scores, boxes, data_dict = im_detect(predictor, data_batch, data_names, scales, config) boxes = boxes[0].astype('f') scores = scores[0].astype('f') dets_nms = [] for j in range(1, scores.shape[1]): cls_scores = scores[:, j, np.newaxis] cls_boxes = boxes[:, 4:8] if config.CLASS_AGNOSTIC else boxes[:, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = nms(cls_dets) cls_dets = cls_dets[keep, :] cls_dets = cls_dets[cls_dets[:, -1] > 0.7, :] dets_nms.append(cls_dets) print 'testing {} {:.4f}s'.format(im_name, toc()) # visualize im = cv2.imread(im_name) #im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) im_result = show_boxes(fileOp, im_name, im, dets_nms, classes, 1) cv2.imwrite(out_dir + im_name.split('/')[-1], im_result) print 'done'
def generate_detections(data, data_names, predictor, config, nms, image_list, detection_num): global classes ret_boxes = [] ret_scores = [] ret_classes = [] k = 0 conversion_dict = { 0: "11", 1: "12", 2: "13", 3: "15", 4: "17", 5: "18", 6: "19", 7: "20", 8: "21", 9: "23", 10: "24", 11: "25", 12: "26", 13: "27", 14: "28", 15: "29", 16: "32", 17: "33", 18: "34", 19: "35", 20: "36", 21: "37", 22: "38", 23: "40", 24: "41", 25: "42", 26: "44", 27: "45", 28: "47", 29: "49", 30: "50", 31: "51", 32: "52", 33: "53", 34: "54", 35: "55", 36: "56", 37: "57", 38: "59", 39: "60", 40: "61", 41: "62", 42: "63", 43: "64", 44: "65", 45: "66", 46: "71", 47: "72", 48: "73", 49: "74", 50: "76", 51: "77", 52: "79", 53: "83", 54: "84", 55: "86", 56: "89", 57: "91", 58: "93", 59: "94" } for idx, im in enumerate(image_list): data_batch = mx.io.DataBatch( data=[data[idx]], label=[], pad=0, index=idx, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[idx])]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] tic() scores, boxes, data_dict = im_detect(predictor, data_batch, data_names, scales, config) boxes = boxes[0].astype('f') scores = scores[0].astype('f') dets_nms = [] for j in range(1, scores.shape[1]): cls_scores = scores[:, j, np.newaxis] cls_boxes = boxes[:, 4:8] if config.CLASS_AGNOSTIC else boxes[:, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = nms(cls_dets) cls_dets = cls_dets[keep, :] cls_dets = cls_dets[cls_dets[:, -1] > 0.08, :] dets_nms.append(cls_dets) # dets_nms format: # a length 60 array, with each element being a dict representing each class # the coordinates are in (xmin, ymin, xmax, ymax, confidence) format. The coordinates are not normalized # one sample is: [290.09448 439.60617 333.31235 461.8115 0.94750994] # below iterates class by class image_detection_num = 0 for index_class in range(len(dets_nms)): # for each class single_class_nms = dets_nms[index_class] image_detection_num += len(single_class_nms) if len(single_class_nms) != 0: print("detecting", single_class_nms.size, "number of objects", len(single_class_nms)) # print(single_class_nms) for index_single_class_nms in range( min(len(single_class_nms), detection_num)): # print("index_class,index_single_class_nms", index_class,index_single_class_nms, ) ret_boxes.append( dets_nms[index_class][index_single_class_nms] [:-1]) #get all the element other than the last one ret_scores.append( dets_nms[index_class][index_single_class_nms] [-1]) #last element ret_classes.append(conversion_dict[index_class]) # pad zeros # print("1st: len(ret_boxes), image_detection_num", len(ret_boxes), image_detection_num) if image_detection_num <= detection_num: for index_element in range(int(detection_num - image_detection_num)): ret_boxes.append(np.zeros( (4, ), dtype=np.float32 )) #get all the element other than the last one ret_scores.append(0) #last element ret_classes.append(0) else: print("~~~~~ too many predictions ~~~~~~~~~~~~~~~") print("len(ret_boxes), image_detection_num", len(ret_boxes), image_detection_num) print('testing image {} {:.4f}s, detection number {}'.format( idx + 1, toc(), image_detection_num)) im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) # visualize # don't show the final images. # show_boxes(im, dets_nms, classes, 1, show_image = False, img_save_name = str(idx) + ".png") ret_boxes = np.squeeze(np.array(ret_boxes)) ret_scores = np.squeeze(np.array(ret_scores)) ret_classes = np.squeeze(np.array(ret_classes)) return ret_boxes, ret_scores, ret_classes
def process_video_frame(raw_frame_queue, bbox_frame_queue): # get symbol pprint.pprint(config) config.symbol = 'resnet_v1_101_rfcn' sym_instance = eval(config.symbol + '.' + config.symbol)() sym = sym_instance.get_symbol(config, is_train=False) arg_params, aux_params = load_param( './output/rfcn/road_obj/road_train_all/all/' + 'rfcn_road', 19, process=True) # set up class names; Don't count the background in, even we are treat the background as label '0' num_classes = 4 classes = ['vehicle', 'pedestrian', 'cyclist', 'traffic lights'] target_size = config.SCALES[0][1] max_size = config.SCALES[0][1] while True: tic() i = 0 data = [] frame_list = [] while len(data) < 15: frame = raw_frame_queue.get() if frame is None: continue if i < 2: i += 1 frame, im_scale = resize(frame, target_size, max_size, stride=config.network.IMAGE_STRIDE) bbox_frame_queue.put(frame) continue frame, im_scale = resize(frame, target_size, max_size, stride=config.network.IMAGE_STRIDE) im_tensor = transform(frame, config.network.PIXEL_MEANS) im_info = np.array( [[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) data.append({'data': im_tensor, 'im_info': im_info}) frame_list.append(frame) # get predictor data_names = ['data', 'im_info'] label_names = [] data = [[mx.nd.array(data[i][name]) for name in data_names] for i in xrange(len(data))] # print('Debug: [data] shape: {}, cont: {}'.format(type(data), data)) max_data_shape = [[('data', (1, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES])))]] # print('Debug: [max_data_shape] shape: {}, cont: {}'.format(type(max_data_shape), max_data_shape)) provide_data = [[(k, v.shape) for k, v in zip(data_names, data[i])] for i in xrange(len(data))] # print('Debug: [provide_data] shape: {}, cont: {}'.format(type(provide_data), provide_data)) provide_label = [None for i in xrange(len(data))] # print('Debug: [provide_label] shape: {}, cont: {}'.format(type(provide_label), provide_label)) predictor = Predictor(sym, data_names, label_names, context=[mx.gpu(0)], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) nms = gpu_nms_wrapper(config.TEST.NMS, 0) # Process video frame # image_names = ['frame'] # for idx, frame in enumerate(frame_list): data_batch = mx.io.DataBatch(data=data, label=[], pad=0, provide_data=provide_data, provide_label=provide_label) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] # print("length: {}".format(len(data_batch.data))) # print('Debug: [scales] cont: {}'.format(scales)) scores_all, boxes_all, data_dict_all = im_detect( predictor, data_batch, data_names, scales, config) # print('scores_all: Type: {}, Values: {}, Length: {}'.format(type(scores_all), scores_all, len(scores_all))) # print('boxes_all: Type: {}, Values: {}, Length: {}'.format(type(boxes_all), boxes_all, len(boxes_all))) # print('data_dict_all: Type: {}, Values: {}, length: {}'.format(type(data_dict_all), data_dict_all, len(data_dict_all))) # print('frame_list: Type: {}, Values: {}, Length: {}'.format(type(frame_list), frame_list, len(frame_list))) # print('scores_all: Type: {}, Length: {}, Values: {}'.format(type(scores_all[0]), len(scores_all[0]), scores_all[0])) # print(scores_all[0].shape) # print('boxes_all: Type: {}, Length: {}'.format(type(boxes_all), len(boxes_all))) # print(boxes_all[0].shape) # print('data_dict_all: Type: {}, length: {}'.format(type(data_dict_all), len(data_dict_all))) # print('frame_list: Type: {}, Length: {}'.format(type(frame_list), len(frame_list))) for idx, frame in enumerate(frame_list): # print('index: {}'.format(str(idx))) boxes = boxes_all[0].astype('f') scores = scores_all[0].astype('f') dets_nms = [] # print(scores.shape) for j in range(1, scores.shape[1]): cls_scores = scores[:, j, np.newaxis] cls_boxes = boxes[:, 4:8] if config.CLASS_AGNOSTIC else boxes[:, j * 4: (j + 1 ) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = nms(cls_dets) cls_dets = cls_dets[keep, :] cls_dets = cls_dets[cls_dets[:, -1] > 0.7, :] dets_nms.append(cls_dets) bbox_frame_queue.put( draw_bbox_on_frame(frame, dets_nms, classes, scale=scales[idx])) print(toc())
def main(): # get symbol pprint.pprint(cfg) cfg.symbol = 'resnet_v1_101_flownet_rfcn' model = '/../model/rfcn_fgfa_flownet_vid' all_frame_interval = cfg.TEST.KEY_FRAME_INTERVAL * 2 + 1 max_per_image = cfg.TEST.max_per_image feat_sym_instance = eval(cfg.symbol + '.' + cfg.symbol)() aggr_sym_instance = eval(cfg.symbol + '.' + cfg.symbol)() feat_sym = feat_sym_instance.get_feat_symbol(cfg) aggr_sym = aggr_sym_instance.get_aggregation_symbol(cfg) # set up class names num_classes = 31 classes = ['__background__','airplane', 'antelope', 'bear', 'bicycle', 'bird', 'bus', 'car', 'cattle', 'dog', 'domestic_cat', 'elephant', 'fox', 'giant_panda', 'hamster', 'horse', 'lion', 'lizard', 'monkey', 'motorcycle', 'rabbit', 'red_panda', 'sheep', 'snake', 'squirrel', 'tiger', 'train', 'turtle', 'watercraft', 'whale', 'zebra'] # load demo data image_names = glob.glob(cur_path + '/../demo/ILSVRC2015_val_00007010/*.JPEG') output_dir = cur_path + '/../demo/rfcn_fgfa/' if not os.path.exists(output_dir): os.makedirs(output_dir) data = [] for im_name in image_names: assert os.path.exists(im_name), ('%s does not exist'.format(im_name)) im = cv2.imread(im_name, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION) target_size = cfg.SCALES[0][0] max_size = cfg.SCALES[0][1] im, im_scale = resize(im, target_size, max_size, stride=cfg.network.IMAGE_STRIDE) im_tensor = transform(im, cfg.network.PIXEL_MEANS) im_info = np.array([[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) feat_stride = float(cfg.network.RCNN_FEAT_STRIDE) data.append({'data': im_tensor, 'im_info': im_info, 'data_cache': im_tensor, 'feat_cache': im_tensor}) # get predictor print 'get-predictor' data_names = ['data', 'im_info', 'data_cache', 'feat_cache'] label_names = [] t1 = time.time() data = [[mx.nd.array(data[i][name]) for name in data_names] for i in xrange(len(data))] max_data_shape = [[('data', (1, 3, max([v[0] for v in cfg.SCALES]), max([v[1] for v in cfg.SCALES]))), ('data_cache', (19, 3, max([v[0] for v in cfg.SCALES]), max([v[1] for v in cfg.SCALES]))), ('feat_cache', ((19, cfg.network.FGFA_FEAT_DIM, np.ceil(max([v[0] for v in cfg.SCALES]) / feat_stride).astype(np.int), np.ceil(max([v[1] for v in cfg.SCALES]) / feat_stride).astype(np.int))))]] provide_data = [[(k, v.shape) for k, v in zip(data_names, data[i])] for i in xrange(len(data))] provide_label = [None for _ in xrange(len(data))] arg_params, aux_params = load_param(cur_path + model, 0, process=True) feat_predictors = Predictor(feat_sym, data_names, label_names, context=[mx.gpu(0)], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) aggr_predictors = Predictor(aggr_sym, data_names, label_names, context=[mx.gpu(0)], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) nms = py_nms_wrapper(cfg.TEST.NMS) # First frame of the video idx = 0 data_batch = mx.io.DataBatch(data=[data[idx]], label=[], pad=0, index=idx, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[idx])]], provide_label=[None]) scales = [data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data))] all_boxes = [[[] for _ in range(len(data))] for _ in range(num_classes)] data_list = deque(maxlen=all_frame_interval) feat_list = deque(maxlen=all_frame_interval) image, feat = get_resnet_output(feat_predictors, data_batch, data_names) # append cfg.TEST.KEY_FRAME_INTERVAL padding images in the front (first frame) while len(data_list) < cfg.TEST.KEY_FRAME_INTERVAL: data_list.append(image) feat_list.append(feat) vis = False file_idx = 0 thresh = 1e-3 for idx, element in enumerate(data): data_batch = mx.io.DataBatch(data=[element], label=[], pad=0, index=idx, provide_data=[[(k, v.shape) for k, v in zip(data_names, element)]], provide_label=[None]) scales = [data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data))] if(idx != len(data)-1): if len(data_list) < all_frame_interval - 1: image, feat = get_resnet_output(feat_predictors, data_batch, data_names) data_list.append(image) feat_list.append(feat) else: ################################################# # main part of the loop ################################################# image, feat = get_resnet_output(feat_predictors, data_batch, data_names) data_list.append(image) feat_list.append(feat) prepare_data(data_list, feat_list, data_batch) pred_result = im_detect(aggr_predictors, data_batch, data_names, scales, cfg) data_batch.data[0][-2] = None data_batch.provide_data[0][-2] = ('data_cache', None) data_batch.data[0][-1] = None data_batch.provide_data[0][-1] = ('feat_cache', None) out_im = process_pred_result(classes, pred_result, num_classes, thresh, cfg, nms, all_boxes, file_idx, max_per_image, vis, data_list[cfg.TEST.KEY_FRAME_INTERVAL].asnumpy(), scales) total_time = time.time()-t1 if (cfg.TEST.SEQ_NMS==False): save_image(output_dir, file_idx, out_im) print 'testing {} {:.4f}s'.format(str(file_idx)+'.JPEG', total_time /(file_idx+1)) file_idx += 1 else: ################################################# # end part of a video # ################################################# end_counter = 0 image, feat = get_resnet_output(feat_predictors, data_batch, data_names) while end_counter < cfg.TEST.KEY_FRAME_INTERVAL + 1: data_list.append(image) feat_list.append(feat) prepare_data(data_list, feat_list, data_batch) pred_result = im_detect(aggr_predictors, data_batch, data_names, scales, cfg) out_im = process_pred_result(classes, pred_result, num_classes, thresh, cfg, nms, all_boxes, file_idx, max_per_image, vis, data_list[cfg.TEST.KEY_FRAME_INTERVAL].asnumpy(), scales) total_time = time.time() - t1 if (cfg.TEST.SEQ_NMS == False): save_image(output_dir, file_idx, out_im) print 'testing {} {:.4f}s'.format(str(file_idx)+'.JPEG', total_time / (file_idx+1)) file_idx += 1 end_counter+=1 if(cfg.TEST.SEQ_NMS): video = [all_boxes[j][:] for j in range(1, num_classes)] dets_all = seq_nms(video) for cls_ind, dets_cls in enumerate(dets_all): for frame_ind, dets in enumerate(dets_cls): keep = nms(dets) all_boxes[cls_ind + 1][frame_ind] = dets[keep, :] for idx in range(len(data)): boxes_this_image = [[]] + [all_boxes[j][idx] for j in range(1, num_classes)] out_im = draw_all_detection(data[idx][0].asnumpy(), boxes_this_image, classes, scales[0], cfg) save_image(output_dir, idx, out_im) print 'done'
def main(): # get symbol pprint.pprint(config) config.symbol = 'resnet_v1_101_rfcn_dcn' if not args.rfcn_only else 'resnet_v1_101_rfcn' sym_instance = eval(config.symbol + '.' + config.symbol)() sym = sym_instance.get_symbol(config, is_train=False) # set up class names num_classes = 81 classes = [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush' ] # load demo data image_names = [ 'COCO_test2015_000000000891.jpg', 'COCO_test2015_000000001669.jpg' ] data = [] for im_name in image_names: assert os.path.exists(cur_path + '/../demo/' + im_name), ( '%s does not exist'.format('../demo/' + im_name)) im = cv2.imread(cur_path + '/../demo/' + im_name, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION) target_size = config.SCALES[0][0] max_size = config.SCALES[0][1] im, im_scale = resize(im, target_size, max_size, stride=config.network.IMAGE_STRIDE) im_tensor = transform(im, config.network.PIXEL_MEANS) im_info = np.array( [[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) data.append({'data': im_tensor, 'im_info': im_info}) # get predictor data_names = ['data', 'im_info'] label_names = [] data = [[mx.nd.array(data[i][name]) for name in data_names] for i in range(len(data))] max_data_shape = [[('data', (1, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES])))]] provide_data = [[(k, v.shape) for k, v in zip(data_names, data[i])] for i in range(len(data))] provide_label = [None for i in range(len(data))] arg_params, aux_params = load_param( cur_path + '/../model/' + ('rfcn_dcn_coco' if not args.rfcn_only else 'rfcn_coco'), 0, process=True) predictor = Predictor(sym, data_names, label_names, context=[mx.gpu(0)], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) nms = gpu_nms_wrapper(config.TEST.NMS, 0) # warm up for j in range(2): data_batch = mx.io.DataBatch(data=[data[0]], label=[], pad=0, index=0, provide_data=[[ (k, v.shape) for k, v in zip(data_names, data[0]) ]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in range(len(data_batch.data)) ] scores, boxes, data_dict = im_detect(predictor, data_batch, data_names, scales, config) # test for idx, im_name in enumerate(image_names): data_batch = mx.io.DataBatch( data=[data[idx]], label=[], pad=0, index=idx, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[idx])]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in range(len(data_batch.data)) ] tic() scores, boxes, data_dict = im_detect(predictor, data_batch, data_names, scales, config) boxes = boxes[0].astype('f') scores = scores[0].astype('f') dets_nms = [] for j in range(1, scores.shape[1]): cls_scores = scores[:, j, np.newaxis] cls_boxes = boxes[:, 4:8] if config.CLASS_AGNOSTIC else boxes[:, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = nms(cls_dets) cls_dets = cls_dets[keep, :] cls_dets = cls_dets[cls_dets[:, -1] > 0.7, :] dets_nms.append(cls_dets) print('testing {} {:.4f}s'.format(im_name, toc())) # visualize im = cv2.imread(cur_path + '/../demo/' + im_name) im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) show_boxes(im, dets_nms, classes, 1) print('done')
def main(): # get symbol pprint.pprint(config) config.symbol = 'resnet_v1_101_rfcn_dcn' if not args.rfcn_only else 'resnet_v1_101_rfcn' sym_instance = eval(config.symbol + '.' + config.symbol)() sym = sym_instance.get_symbol(config, is_train=False) # set up class names num_classes = 81 classes = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'] # load demo data image_names = ['COCO_test2015_000000000891.jpg', 'COCO_test2015_000000001669.jpg'] data = [] for im_name in image_names: assert os.path.exists(cur_path + '/../demo/' + im_name), ('%s does not exist'.format('../demo/' + im_name)) im = cv2.imread(cur_path + '/../demo/' + im_name, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION) target_size = config.SCALES[0][0] max_size = config.SCALES[0][1] im, im_scale = resize(im, target_size, max_size, stride=config.network.IMAGE_STRIDE) im_tensor = transform(im, config.network.PIXEL_MEANS) im_info = np.array([[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) data.append({'data': im_tensor, 'im_info': im_info}) # get predictor data_names = ['data', 'im_info'] label_names = [] data = [[mx.nd.array(data[i][name]) for name in data_names] for i in xrange(len(data))] max_data_shape = [[('data', (1, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES])))]] provide_data = [[(k, v.shape) for k, v in zip(data_names, data[i])] for i in xrange(len(data))] provide_label = [None for i in xrange(len(data))] arg_params, aux_params = load_param(cur_path + '/../model/' + ('rfcn_dcn_coco' if not args.rfcn_only else 'rfcn_coco'), 0, process=True) predictor = Predictor(sym, data_names, label_names, context=[mx.gpu(0)], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) nms = gpu_nms_wrapper(config.TEST.NMS, 0) # warm up for j in xrange(2): data_batch = mx.io.DataBatch(data=[data[0]], label=[], pad=0, index=0, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[0])]], provide_label=[None]) scales = [data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data))] scores, boxes, data_dict = im_detect(predictor, data_batch, data_names, scales, config) # test for idx, im_name in enumerate(image_names): data_batch = mx.io.DataBatch(data=[data[idx]], label=[], pad=0, index=idx, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[idx])]], provide_label=[None]) scales = [data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data))] tic() scores, boxes, data_dict = im_detect(predictor, data_batch, data_names, scales, config) boxes = boxes[0].astype('f') scores = scores[0].astype('f') dets_nms = [] for j in range(1, scores.shape[1]): cls_scores = scores[:, j, np.newaxis] cls_boxes = boxes[:, 4:8] if config.CLASS_AGNOSTIC else boxes[:, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = nms(cls_dets) cls_dets = cls_dets[keep, :] cls_dets = cls_dets[cls_dets[:, -1] > 0.7, :] dets_nms.append(cls_dets) print 'testing {} {:.4f}s'.format(im_name, toc()) # visualize im = cv2.imread(cur_path + '/../demo/' + im_name) im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) show_boxes(im, dets_nms, classes, 1) print 'done'
def main(): # get symbol pprint.pprint(cfg) cfg.symbol = 'resnet_v1_101_flownet_rfcn' model = '/data2/output/fgfa_rfcn/jrdb/resnet_v1_101_flownet_jrdb/VID_train_15frames/fgfa_rfcn_vid' all_frame_interval = cfg.TEST.KEY_FRAME_INTERVAL * 2 + 1 max_per_image = cfg.TEST.max_per_image feat_sym_instance = eval(cfg.symbol + '.' + cfg.symbol)() aggr_sym_instance = eval(cfg.symbol + '.' + cfg.symbol)() feat_sym = feat_sym_instance.get_feat_symbol(cfg) aggr_sym = aggr_sym_instance.get_aggregation_symbol(cfg) # set up class names classes = ['__background__', 'p'] num_classes = len(classes) data_loader = DataLoader(args.input + '/*', cfg) #line = '%s %s %s %s\n' % (dire.replace(join(join(data_dir, ANNOTATIONS, VID)) + "/", ""), '1', str(i), str(num_of_images)) output_dir_tmp = '%s_%s_epoc_%d' % ( args.dataset, os.path.basename(args.input), args.epoc) output_dir = os.path.join("/data2", "output", output_dir_tmp) print(output_dir) if not os.path.exists(output_dir): os.makedirs(output_dir) data_names = ['data', 'im_info', 'data_cache', 'feat_cache'] feat_stride = float(cfg.network.RCNN_FEAT_STRIDE) # for im_name in image_names: # im_info, im_tensor = read_data(im_name) # im_info, im_tensor = read_data(im_name) # data.append({'data': im_tensor, 'im_info': im_info, 'data_cache': im_tensor, 'feat_cache': im_tensor}) # ### data (1,3, 562,1000) #data = [[mx.nd.array(data[i][name]) for name in data_names] for i in xrange(len(data))] # list of data images # get predictor print 'get-predictor' label_names = [] t1 = time.time() max_data_shape = [[ ('data', (1, 3, max([v[0] for v in cfg.SCALES]), max([v[1] for v in cfg.SCALES]))), ('data_cache', (19, 3, max([v[0] for v in cfg.SCALES]), max([v[1] for v in cfg.SCALES]))), ('feat_cache', ((19, cfg.network.FGFA_FEAT_DIM, np.ceil(max([v[0] for v in cfg.SCALES]) / feat_stride).astype(np.int), np.ceil(max([v[1] for v in cfg.SCALES]) / feat_stride).astype(np.int)))) ]] provide_data, provide_label = data_loader.get_provided_data_and_label() arg_params, aux_params = load_param(model, args.epoc, process=True) feat_predictors = Predictor(feat_sym, data_names, label_names, context=[mx.gpu(0)], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) aggr_predictors = Predictor(aggr_sym, data_names, label_names, context=[mx.gpu(0)], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) nms = py_nms_wrapper(cfg.TEST.NMS) # First frame of the video idx = 0 data_idx_ = data_loader.get_data_by_index(idx) len_data = data_loader.get_len_data() data_batch = mx.io.DataBatch(data=[data_idx_], label=[], pad=0, index=idx, provide_data=[[ (k, v.shape) for k, v in zip(data_names, data_idx_) ]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] all_boxes = [[[] for _ in range(len_data)] for _ in range(num_classes)] data_list = deque(maxlen=all_frame_interval) feat_list = deque(maxlen=all_frame_interval) image, feat = get_resnet_output(feat_predictors, data_batch, data_names) # append cfg.TEST.KEY_FRAME_INTERVAL padding images in the front (first frame) while len(data_list) < cfg.TEST.KEY_FRAME_INTERVAL: data_list.append(image) feat_list.append(feat) vis = False file_idx = 0 thresh = args.threshold for idx, element in enumerate(data_loader): # loop through list of images! data_batch = mx.io.DataBatch(data=[element], label=[], pad=0, index=idx, provide_data=[[ (k, v.shape) for k, v in zip(data_names, element) ]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] if (idx != len_data - 1): image, feat = get_resnet_output(feat_predictors, data_batch, data_names) data_list.append(image) feat_list.append(feat) if len(data_list) >= all_frame_interval - 1: ################################################# # main part of the loop ################################################# image, feat = get_resnet_output(feat_predictors, data_batch, data_names) data_list.append(image) feat_list.append(feat) prepare_data(data_list, feat_list, data_batch) pred_result = im_detect(aggr_predictors, data_batch, data_names, scales, cfg) data_batch.data[0][-2] = None data_batch.provide_data[0][-2] = ('data_cache', None) data_batch.data[0][-1] = None data_batch.provide_data[0][-1] = ('feat_cache', None) out_im = process_pred_result( classes, pred_result, num_classes, thresh, cfg, nms, all_boxes, file_idx, max_per_image, vis, data_list[cfg.TEST.KEY_FRAME_INTERVAL].asnumpy(), scales) total_time = time.time() - t1 if (cfg.TEST.SEQ_NMS == False): save_image(output_dir, file_idx, out_im) print 'testing {} {:.4f}s'.format( str(file_idx) + '.JPEG', total_time / (file_idx + 1)) file_idx += 1 else: ################################################# # end part of a video # ################################################# end_counter = 0 image, feat = get_resnet_output(feat_predictors, data_batch, data_names) while end_counter < cfg.TEST.KEY_FRAME_INTERVAL + 1: data_list.append(image) feat_list.append(feat) prepare_data(data_list, feat_list, data_batch) pred_result = im_detect(aggr_predictors, data_batch, data_names, scales, cfg) out_im = process_pred_result( classes, pred_result, num_classes, thresh, cfg, nms, all_boxes, file_idx, max_per_image, vis, data_list[cfg.TEST.KEY_FRAME_INTERVAL].asnumpy(), scales) total_time = time.time() - t1 if (cfg.TEST.SEQ_NMS == False): save_image(output_dir, file_idx, out_im) print 'testing {} {:.4f}s'.format( str(file_idx) + '.JPEG', total_time / (file_idx + 1)) file_idx += 1 end_counter += 1 if (cfg.TEST.SEQ_NMS): video = [all_boxes[j][:] for j in range(1, num_classes)] dets_all = seq_nms(video) for cls_ind, dets_cls in enumerate(dets_all): for frame_ind, dets in enumerate(dets_cls): keep = nms(dets) all_boxes[cls_ind + 1][frame_ind] = dets[keep, :] for idx in range(len_data): boxes_this_image = [[]] + [ all_boxes[j][idx] for j in range(1, num_classes) ] out_im = draw_all_detection(data[idx][0].asnumpy(), boxes_this_image, classes, scales[0], cfg) save_image(output_dir, idx, out_im) print 'done'
for j in xrange(2): data_batch = mx.io.DataBatch(data=[data[j]], label=[], pad=0, index=0, provide_data=[[ (k, v.shape) for k, v in zip(data_names, data[j]) ]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] if j % key_frame_interval == 0: scores, boxes, data_dict, feat = im_detect(key_predictor, data_batch, data_names, scales, config) else: data_batch.data[0][-1] = feat data_batch.provide_data[0][-1] = ('feat_key', feat.shape) scores, boxes, data_dict, _ = im_detect(cur_predictor, data_batch, data_names, scales, config) #%% mx.nd.save('feat_key_list', [feat]) feat_np = feat.asnumpy() print feat_np.shape feat_np.tofile('feat_key.binary') h5f = h5py.File('feat_key.h5', 'w') print h5f.create_dataset('dataset_1', data=feat_np) h5f.close() # exit(0)
def pred_eval(predictor, test_data, imdb, cfg, vis=False, thresh=1e-3, logger=None, ignore_cache=True): """ wrapper for calculating offline validation for faster data analysis in this example, all threshold are set by hand :param predictor: Predictor :param test_data: data iterator, must be non-shuffle :param imdb: image database :param vis: controls visualization :param thresh: valid detection threshold :return: """ det_file = os.path.join(imdb.result_path, imdb.name + '_detections.pkl') if os.path.exists(det_file) and not ignore_cache: with open(det_file, 'rb') as fid: cache_res = cPickle.load(fid) all_boxes = cache_res['all_boxes'] all_keypoints = cache_res.get('all_keypoints') info_str = imdb.evaluate_detections(all_boxes, all_keypoints=all_keypoints) if logger: logger.info('evaluate detections: \n{}'.format(info_str)) return assert vis or not test_data.shuffle data_names = [k[0] for k in test_data.provide_data] if not isinstance(test_data, PrefetchingIter): test_data = PrefetchingIter(test_data) nms = py_nms_wrapper(cfg.TEST.NMS) # limit detections to max_per_image over all classes max_per_image = cfg.TEST.max_per_image num_images = imdb.num_images # all detections are collected into: # all_boxes[cls][image] = N x 5 array of detections in # (x1, y1, x2, y2, score) all_boxes = [[np.array([]) for _ in range(num_images)] for _ in range(imdb.num_classes)] all_keypoints = None if cfg.network.PREDICT_KEYPOINTS: all_keypoints = [[np.array([]) for _ in range(num_images)] for _ in range(imdb.num_classes)] idx = 0 data_time, net_time, post_time = 0.0, 0.0, 0.0 t = time.time() for data_batch in test_data: t1 = time.time() - t t = time.time() rets = im_detect(predictor, data_batch, data_names, cfg) scores_all = rets[0] boxes_all = rets[1] if cfg.network.PREDICT_KEYPOINTS: pred_kps_all = rets[2] t2 = time.time() - t t = time.time() for delta, (scores, boxes) in enumerate(zip(scores_all, boxes_all)): if idx + delta >= num_images: break for j in range(1, imdb.num_classes): indexes = np.where(scores[:, j] > thresh)[0] cls_scores = scores[indexes, j, np.newaxis] cls_boxes = boxes[indexes, 4:8] if cfg.CLASS_AGNOSTIC else boxes[ indexes, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = nms(cls_dets) all_boxes[j][idx + delta] = cls_dets[keep, :] if cfg.network.PREDICT_KEYPOINTS: all_keypoints[j][idx + delta] = pred_kps_all[delta][ indexes, :][keep, :] if max_per_image > 0: image_scores = np.hstack([ all_boxes[j][idx + delta][:, -1] for j in range(1, imdb.num_classes) ]) if len(image_scores) > max_per_image: image_thresh = np.sort(image_scores)[-max_per_image] for j in range(1, imdb.num_classes): keep = np.where( all_boxes[j][idx + delta][:, -1] >= image_thresh)[0] all_boxes[j][idx + delta] = all_boxes[j][idx + delta][keep, :] if cfg.network.PREDICT_KEYPOINTS: all_keypoints[j][idx + delta] = all_keypoints[j][ idx + delta][keep, :] if vis: boxes_this_image = [[]] + [ all_boxes[j][idx + delta] for j in range(1, imdb.num_classes) ] vis_all_detection(data_dict['data'].asnumpy(), boxes_this_image, imdb.classes, scales[delta], cfg) idx += test_data.batch_size t3 = time.time() - t t = time.time() msg = 'testing {}/{} data {:.4f}s net {:.4f}s post {:.4f}s'.format( idx, imdb.num_images, t1, t2, t3) print msg if logger: logger.info(msg) with open(det_file, 'wb') as f: cPickle.dump({ 'all_boxes': all_boxes, 'all_keypoints': all_keypoints }, f, protocol=cPickle.HIGHEST_PROTOCOL) info_str = imdb.evaluate_detections(all_boxes, all_keypoints=all_keypoints) if logger: logger.info('evaluate detections: \n{}'.format(info_str))
def main(): # get symbol pprint.pprint(config) config.symbol = 'resnet_v1_101_rfcn' sym_instance = eval(config.symbol + '.' + config.symbol)() sym = sym_instance.get_symbol(config, is_train=False) # set up class names; Don't count the background in, even we are treat the background as label '0' num_classes = 4 classes = ['vehicle', 'pedestrian', 'cyclist', 'traffic lights'] # load demo data image_path = './data/RoadImages/test/' image_names = glob.glob(image_path + '*.jpg') print("Image amount {}".format(len(image_names))) data = [] for im_name in image_names: assert os.path.exists(im_name), ('%s does not exist'.format(im_name)) im = cv2.imread(im_name, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION) target_size = config.SCALES[0][1] max_size = config.SCALES[0][1] im, im_scale = resize(im, target_size, max_size, stride=config.network.IMAGE_STRIDE) im_tensor = transform(im, config.network.PIXEL_MEANS) im_info = np.array( [[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) data.append({'data': im_tensor, 'im_info': im_info}) # get predictor data_names = ['data', 'im_info'] label_names = [] data = [[mx.nd.array(data[i][name]) for name in data_names] for i in xrange(len(data))] max_data_shape = [[('data', (1, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES])))]] provide_data = [[(k, v.shape) for k, v in zip(data_names, data[i])] for i in xrange(len(data))] provide_label = [None for i in xrange(len(data))] arg_params, aux_params = load_param( './output/rfcn/road_obj/road_train_all/all/' + 'rfcn_road', 19, process=True) predictor = Predictor(sym, data_names, label_names, context=[mx.gpu(0)], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) nms = gpu_nms_wrapper(config.TEST.NMS, 0) # test notation_dict = {} for idx, im_name in enumerate(image_names): data_batch = mx.io.DataBatch( data=[data[idx]], label=[], pad=0, index=idx, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[idx])]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] tic() scores, boxes, data_dict = im_detect(predictor, data_batch, data_names, scales, config) boxes = boxes[0].astype('f') scores = scores[0].astype('f') dets_nms = [] for j in range(1, scores.shape[1]): cls_scores = scores[:, j, np.newaxis] cls_boxes = boxes[:, 4:8] if config.CLASS_AGNOSTIC else boxes[:, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = nms(cls_dets) cls_dets = cls_dets[keep, :] cls_dets = cls_dets[cls_dets[:, -1] > 0.7, :] dets_nms.append(cls_dets) print 'testing {} {:.4f}s'.format(im_name, toc()) # notation_list.append(get_notation(im_name, dets_nms, classes, scale=1.0, gen_bbox_pic=True)) notation_dict.update( get_notation(im_name, dets_nms, classes, scale=1.0, gen_bbox_pic=True)) save_notation_file(notation_dict) print 'done'
def main(): # get symbol ctx_id = [int(i) for i in config.gpus.split(',')] pprint.pprint(config) sym_instance = eval(config.symbol)() sym = sym_instance.get_symbol(config, is_train=False) # set up class names num_classes = 6 classes = [ '__background__', # always index 0 '1', '2', '3', '4', '5' ] # classes = ['__background__', # always index 0 # 'CubeBody', 'CubeFace', 'CylinderBody', 'CylinderFace', 'Grip'] # load demo data image_names = [] names_dirs = os.listdir(cur_path + '/../' + test_dir) for im_name in names_dirs: if im_name[-4:] == '.jpg': image_names.append(im_name) data = [] # next [3001: len(image_names)] for im_name in image_names[9001:len(image_names)]: assert os.path.exists(cur_path + '/../' + test_dir + im_name), ( '%s does not exist'.format('../demo/' + im_name)) im = cv2.imread(cur_path + '/../' + test_dir + im_name, cv2.IMREAD_COLOR | long(128)) target_size = config.SCALES[0][0] max_size = config.SCALES[0][1] # print "before scale: " # print im.shape im, im_scale = resize(im, target_size, max_size, stride=config.network.IMAGE_STRIDE) # print "after scale: " # print im.shape # im_scale = 1.0 # print "scale ratio: " # print im_scale im_tensor = transform(im, config.network.PIXEL_MEANS) # print im_tensor.shape im_info = np.array( [[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32) data.append({'data': im_tensor, 'im_info': im_info}) # get predictor data_names = ['data', 'im_info'] label_names = [] data = [[mx.nd.array(data[i][name]) for name in data_names] for i in xrange(len(data))] max_data_shape = [[('data', (1, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES])))]] provide_data = [[(k, v.shape) for k, v in zip(data_names, data[i])] for i in xrange(len(data))] provide_label = [None for i in xrange(len(data))] model_path = cur_path + '/../' + config.output_path + '/' + config.config_name + '/' + config.dataset.image_set + '/' + config.TRAIN.model_prefix arg_params, aux_params = load_param(model_path, config.TEST.test_epoch, process=True) predictor = Predictor(sym, data_names, label_names, context=[mx.gpu(ctx_id[0])], max_data_shapes=max_data_shape, provide_data=provide_data, provide_label=provide_label, arg_params=arg_params, aux_params=aux_params) # warm up for i in xrange(2): data_batch = mx.io.DataBatch(data=[data[0]], label=[], pad=0, index=0, provide_data=[[ (k, v.shape) for k, v in zip(data_names, data[0]) ]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] _, _, _, _ = im_detect(predictor, data_batch, data_names, scales, config) # 12342_mask Log # LogTxt = open('LogRecorder.txt', 'w') # test for idx, im_name in enumerate(image_names[9001:len(image_names)]): data_batch = mx.io.DataBatch( data=[data[idx]], label=[], pad=0, index=idx, provide_data=[[(k, v.shape) for k, v in zip(data_names, data[idx])]], provide_label=[None]) scales = [ data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data)) ] tic() scores, boxes, masks, data_dict = im_detect(predictor, data_batch, data_names, [1.0], config) im_shapes = [ data_batch.data[i][0].shape[2:4] for i in xrange(len(data_batch.data)) ] # print im_shapes if not config.TEST.USE_MASK_MERGE: all_boxes = [[] for _ in xrange(num_classes)] all_masks = [[] for _ in xrange(num_classes)] nms = py_nms_wrapper(config.TEST.NMS) for j in range(1, num_classes): indexes = np.where(scores[0][:, j] > 0.7)[0] cls_scores = scores[0][indexes, j, np.newaxis] cls_masks = masks[0][indexes, 1, :, :] try: if config.CLASS_AGNOSTIC: cls_boxes = boxes[0][indexes, :] else: raise Exception() except: cls_boxes = boxes[0][indexes, j * 4:(j + 1) * 4] cls_dets = np.hstack((cls_boxes, cls_scores)) keep = nms(cls_dets) all_boxes[j] = cls_dets[keep, :] all_masks[j] = cls_masks[keep, :] dets = [all_boxes[j] for j in range(1, num_classes)] masks = [all_masks[j] for j in range(1, num_classes)] else: masks = masks[0][:, 1:, :, :] result_masks, result_dets = gpu_mask_voting( masks, boxes[0], scores[0], num_classes, 100, im_shapes[0][1], im_shapes[0][0], config.TEST.NMS, config.TEST.MASK_MERGE_THRESH, config.BINARY_THRESH, ctx_id[0]) dets = [result_dets[j] for j in range(1, num_classes)] masks = [ result_masks[j][:, 0, :, :] for j in range(1, num_classes) ] print 'testing {} {:.4f}s'.format(im_name, toc()) # visualize for i in xrange(len(dets)): keep = np.where(dets[i][:, -1] > 0.7) dets[i] = dets[i][keep] masks[i] = masks[i][keep] im = cv2.imread(cur_path + '/../' + test_dir + im_name) # im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) ims = show_masks(im, dets, masks, classes, config, 1.0 / scales[0], False) # label = cv2.imread(cur_path + '/../' + label_dir + im_name[0:len(im_name)-4] + '.png') # label = cv2.cvtColor(label, cv2.COLOR_BGR2RGB) # label_flag = checkLabel(label) if dets[0].size: cv2.imwrite( cur_path + '/../' + result_dir + im_name[0:-4] + '_CuBo.png', ims[0]) if dets[1].size: cv2.imwrite( cur_path + '/../' + result_dir + im_name[0:-4] + '_CuFa.png', ims[1]) if dets[2].size: cv2.imwrite( cur_path + '/../' + result_dir + im_name[0:-4] + '_CyBo.png', ims[2]) if dets[3].size: cv2.imwrite( cur_path + '/../' + result_dir + im_name[0:-4] + '_CyFa.png', ims[3]) if dets[4].size: cv2.imwrite( cur_path + '/../' + result_dir + im_name[0:-4] + '_Grip.png', ims[4]) # if np.unique(ims[0]).shape[0] >= 2 and (1 in label_flag): # cv2.imwrite(cur_path + '/../' + result_dir + im_name[0:-4] + '_CuBo.png', ims[0]) # if np.unique(ims[1]).shape[0] >= 2 and (2 in label_flag): # cv2.imwrite(cur_path + '/../' + result_dir + im_name[0:-4] + '_CuFa.png', ims[1]) # if np.unique(ims[2]).shape[0] >= 2 and (3 in label_flag): # cv2.imwrite(cur_path + '/../' + result_dir + im_name[0:-4] + '_CyBo.png', ims[2]) # if np.unique(ims[3]).shape[0] >= 2 and (4 in label_flag): # cv2.imwrite(cur_path + '/../' + result_dir + im_name[0:-4] + '_CyFa.png', ims[3]) # if np.unique(ims[4]).shape[0] >= 2 and (5 in label_flag): # cv2.imwrite(cur_path + '/../' + result_dir + im_name[0:-4] + '_Grip.png', ims[4]) print 'done'