def test_pb(self, frozen_graph_path, test_dir): graph = self.load_graph(frozen_graph_path) print("we are testing ====>>>>", frozen_graph_path) img = graph.get_tensor_by_name("input_img:0") dets = graph.get_tensor_by_name("DetResults:0") with tf.Session(graph=graph) as sess: for img_path in os.listdir(test_dir): print(img_path) a_img = cv2.imread(os.path.join(test_dir, img_path))[:, :, ::-1] raw_h, raw_w = a_img.shape[0], a_img.shape[1] short_size, max_len = self.cfgs.IMG_SHORT_SIDE_LEN, cfgs.IMG_MAX_LENGTH if raw_h < raw_w: new_h, new_w = short_size, min(int(short_size * float(raw_w) / raw_h), max_len) else: new_h, new_w = min(int(short_size * float(raw_h) / raw_w), max_len), short_size img_resize = cv2.resize(a_img, (new_w, new_h)) dets_val = sess.run(dets, feed_dict={img: img_resize[:, :, ::-1]}) bbox_pred, cls_prob, proposal = dets_val[:, :5], dets_val[:, 5:(5+self.cfgs.CLASS_NUM)], \ dets_val[:, (5+self.cfgs.CLASS_NUM):] detected_boxes, detected_scores, detected_categories = self.postprocess_detctions(bbox_pred, cls_prob, proposal) if True: # detected_indices = det_scores_r_ >= self.cfgs.VIS_SCORE # detected_scores = det_scores_r_[detected_indices] # detected_boxes = det_boxes_r_[detected_indices] # detected_categories = det_category_r_[detected_indices] drawer = DrawBox(self.cfgs) det_detections_r = drawer.draw_boxes_with_label_and_scores(img_resize[:, :, ::-1], boxes=detected_boxes, labels=detected_categories, scores=detected_scores, method=1, in_graph=True) save_dir = os.path.join('test_pb', self.cfgs.VERSION, 'pb_img_vis') tools.makedirs(save_dir) cv2.imwrite(save_dir + '/{}'.format(img_path), det_detections_r[:, :, ::-1])
def test_dota(self, det_net, real_test_img_list, txt_name): save_path = os.path.join('./test_dota', self.cfgs.VERSION) nr_records = len(real_test_img_list) pbar = tqdm(total=nr_records) gpu_num = len(self.args.gpus.strip().split(',')) nr_image = math.ceil(nr_records / gpu_num) result_queue = Queue(500) procs = [] for i, gpu_id in enumerate(self.args.gpus.strip().split(',')): start = i * nr_image end = min(start + nr_image, nr_records) split_records = real_test_img_list[start:end] proc = Process(target=self.worker, args=(int(gpu_id), split_records, det_net, result_queue)) print('process:%d, start:%d, end:%d' % (i, start, end)) proc.start() procs.append(proc) for i in range(nr_records): res = result_queue.get() if self.args.show_box: nake_name = res['image_id'].split('/')[-1] tools.makedirs(os.path.join(save_path, 'dota_img_vis')) draw_path = os.path.join(save_path, 'dota_img_vis', nake_name) draw_img = np.array(cv2.imread(res['image_id']), np.float32) detected_boxes = backward_convert(res['boxes'], with_label=False) detected_indices = res['scores'] >= self.cfgs.VIS_SCORE detected_scores = res['scores'][detected_indices] detected_boxes = detected_boxes[detected_indices] detected_categories = res['labels'][detected_indices] drawer = DrawBox(self.cfgs) final_detections = drawer.draw_boxes_with_label_and_scores( draw_img, boxes=detected_boxes, labels=detected_categories, scores=detected_scores, method=1, is_csl=True, in_graph=False) cv2.imwrite(draw_path, final_detections) else: CLASS_DOTA = self.name_label_map.keys() write_handle = {} tools.makedirs(os.path.join(save_path, 'dota_res')) for sub_class in CLASS_DOTA: if sub_class == 'back_ground': continue write_handle[sub_class] = open( os.path.join(save_path, 'dota_res', 'Task1_%s.txt' % sub_class), 'a+') for i, rbox in enumerate(res['boxes']): command = '%s %.3f %.1f %.1f %.1f %.1f %.1f %.1f %.1f %.1f\n' % ( res['image_id'].split('/')[-1].split('.')[0], res['scores'][i], rbox[0], rbox[1], rbox[2], rbox[3], rbox[4], rbox[5], rbox[6], rbox[7], ) write_handle[self.label_name_map[res['labels'][i]]].write( command) for sub_class in CLASS_DOTA: if sub_class == 'back_ground': continue write_handle[sub_class].close() fw = open(txt_name, 'a+') fw.write('{}\n'.format(res['image_id'].split('/')[-1])) fw.close() pbar.set_description("Test image %s" % res['image_id'].split('/')[-1]) pbar.update(1) for p in procs: p.join()
def eval_with_plac(self, img_dir, det_net, image_ext): os.environ["CUDA_VISIBLE_DEVICES"] = self.args.gpu # 1. preprocess img img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3]) # is RGB. not BGR img_batch = tf.cast(img_plac, tf.float32) pretrain_zoo = PretrainModelZoo() if self.cfgs.NET_NAME in pretrain_zoo.pth_zoo or self.cfgs.NET_NAME in pretrain_zoo.mxnet_zoo: img_batch = (img_batch / 255 - tf.constant( self.cfgs.PIXEL_MEAN_)) / tf.constant(self.cfgs.PIXEL_STD) else: img_batch = img_batch - tf.constant(self.cfgs.PIXEL_MEAN) img_batch = tf.expand_dims(img_batch, axis=0) detection_boxes, detection_scores, detection_category = det_net.build_whole_detection_network( input_img_batch=img_batch) init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) restorer, restore_ckpt = det_net.get_restorer() config = tf.ConfigProto() config.gpu_options.allow_growth = True with tf.Session(config=config) as sess: sess.run(init_op) if not restorer is None: restorer.restore(sess, restore_ckpt) print('restore model') all_boxes_r = [] imgs = os.listdir(img_dir) pbar = tqdm(imgs) for a_img_name in pbar: a_img_name = a_img_name.split(image_ext)[0] raw_img = cv2.imread( os.path.join(img_dir, a_img_name + image_ext)) raw_h, raw_w = raw_img.shape[0], raw_img.shape[1] det_boxes_r_all, det_scores_r_all, det_category_r_all = [], [], [] img_short_side_len_list = self.cfgs.IMG_SHORT_SIDE_LEN if isinstance( self.cfgs.IMG_SHORT_SIDE_LEN, list) else [self.cfgs.IMG_SHORT_SIDE_LEN] img_short_side_len_list = [ img_short_side_len_list[0] ] if not self.args.multi_scale else img_short_side_len_list for short_size in img_short_side_len_list: max_len = self.cfgs.IMG_MAX_LENGTH if raw_h < raw_w: new_h, new_w = short_size, min( int(short_size * float(raw_w) / raw_h), max_len) else: new_h, new_w = min( int(short_size * float(raw_h) / raw_w), max_len), short_size img_resize = cv2.resize(raw_img, (new_w, new_h)) resized_img, detected_boxes, detected_scores, detected_categories = \ sess.run( [img_batch, detection_boxes, detection_scores, detection_category], feed_dict={img_plac: img_resize[:, :, ::-1]} ) if detected_boxes.shape[0] == 0: continue resized_h, resized_w = resized_img.shape[ 1], resized_img.shape[2] detected_boxes = forward_convert(detected_boxes, False) detected_boxes[:, 0::2] *= (raw_w / resized_w) detected_boxes[:, 1::2] *= (raw_h / resized_h) det_boxes_r_all.extend(detected_boxes) det_scores_r_all.extend(detected_scores) det_category_r_all.extend(detected_categories) det_boxes_r_all = np.array(det_boxes_r_all) det_scores_r_all = np.array(det_scores_r_all) det_category_r_all = np.array(det_category_r_all) box_res_rotate_ = [] label_res_rotate_ = [] score_res_rotate_ = [] if det_scores_r_all.shape[0] != 0: for sub_class in range(1, self.cfgs.CLASS_NUM + 1): index = np.where(det_category_r_all == sub_class)[0] if len(index) == 0: continue tmp_boxes_r = det_boxes_r_all[index] tmp_label_r = det_category_r_all[index] tmp_score_r = det_scores_r_all[index] if self.args.multi_scale: tmp_boxes_r_ = backward_convert(tmp_boxes_r, False) # try: # inx = nms_rotate.nms_rotate_cpu(boxes=np.array(tmp_boxes_r_), # scores=np.array(tmp_score_r), # iou_threshold=self.cfgs.NMS_IOU_THRESHOLD, # max_output_size=5000) # except: tmp_boxes_r_ = np.array(tmp_boxes_r_) tmp = np.zeros([ tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1 ]) tmp[:, 0:-1] = tmp_boxes_r_ tmp[:, -1] = np.array(tmp_score_r) # Note: the IoU of two same rectangles is 0, which is calculated by rotate_gpu_nms jitter = np.zeros([ tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1 ]) jitter[:, 0] += np.random.rand( tmp_boxes_r_.shape[0], ) / 1000 inx = rotate_gpu_nms( np.array(tmp, np.float32) + np.array(jitter, np.float32), float(self.cfgs.NMS_IOU_THRESHOLD), 0) else: inx = np.arange(0, tmp_score_r.shape[0]) box_res_rotate_.extend(np.array(tmp_boxes_r)[inx]) score_res_rotate_.extend(np.array(tmp_score_r)[inx]) label_res_rotate_.extend(np.array(tmp_label_r)[inx]) if len(box_res_rotate_) == 0: all_boxes_r.append(np.array([])) continue det_boxes_r_ = np.array(box_res_rotate_) det_scores_r_ = np.array(score_res_rotate_) det_category_r_ = np.array(label_res_rotate_) if self.args.draw_imgs: detected_indices = det_scores_r_ >= self.cfgs.VIS_SCORE detected_scores = det_scores_r_[detected_indices] detected_boxes = det_boxes_r_[detected_indices] detected_categories = det_category_r_[detected_indices] detected_boxes = backward_convert(detected_boxes, False) drawer = DrawBox(self.cfgs) det_detections_r = drawer.draw_boxes_with_label_and_scores( raw_img[:, :, ::-1], boxes=detected_boxes, labels=detected_categories, scores=detected_scores, method=1, in_graph=True) save_dir = os.path.join('test_hrsc', self.cfgs.VERSION, 'hrsc2016_img_vis') tools.makedirs(save_dir) cv2.imwrite(save_dir + '/{}.jpg'.format(a_img_name), det_detections_r[:, :, ::-1]) det_boxes_r_ = backward_convert(det_boxes_r_, False) x_c, y_c, w, h, theta = det_boxes_r_[:, 0], det_boxes_r_[:, 1], det_boxes_r_[:, 2], \ det_boxes_r_[:, 3], det_boxes_r_[:, 4] boxes_r = np.transpose(np.stack([x_c, y_c, w, h, theta])) dets_r = np.hstack((det_category_r_.reshape(-1, 1), det_scores_r_.reshape(-1, 1), boxes_r)) all_boxes_r.append(dets_r) pbar.set_description("Eval image %s" % a_img_name) # fw1 = open(cfgs.VERSION + '_detections_r.pkl', 'wb') # pickle.dump(all_boxes_r, fw1) return all_boxes_r
def test_icdar2015(self, det_net, real_test_img_list, txt_name): save_path = os.path.join('./test_icdar2015', self.cfgs.VERSION) tools.makedirs(save_path) nr_records = len(real_test_img_list) pbar = tqdm(total=nr_records) gpu_num = len(self.args.gpus.strip().split(',')) nr_image = math.ceil(nr_records / gpu_num) result_queue = Queue(500) procs = [] for i, gpu_id in enumerate(self.args.gpus.strip().split(',')): start = i * nr_image end = min(start + nr_image, nr_records) split_records = real_test_img_list[start:end] proc = Process(target=self.worker, args=(int(gpu_id), split_records, det_net, result_queue)) print('process:%d, start:%d, end:%d' % (i, start, end)) proc.start() procs.append(proc) for i in range(nr_records): res = result_queue.get() tools.makedirs(os.path.join(save_path, 'icdar2015_res')) if res['boxes'].shape[0] == 0: fw_txt_dt = open(os.path.join(save_path, 'icdar2015_res', 'res_{}.txt'.format(res['image_id'].split('/')[-1].split('.')[0])), 'w') fw_txt_dt.close() pbar.update(1) fw = open(txt_name, 'a+') fw.write('{}\n'.format(res['image_id'].split('/')[-1])) fw.close() continue x1, y1, x2, y2, x3, y3, x4, y4 = res['boxes'][:, 0], res['boxes'][:, 1], res['boxes'][:, 2], res['boxes'][:, 3],\ res['boxes'][:, 4], res['boxes'][:, 5], res['boxes'][:, 6], res['boxes'][:, 7] x1, y1 = x1 * res['scales'][0], y1 * res['scales'][1] x2, y2 = x2 * res['scales'][0], y2 * res['scales'][1] x3, y3 = x3 * res['scales'][0], y3 * res['scales'][1] x4, y4 = x4 * res['scales'][0], y4 * res['scales'][1] boxes = np.transpose(np.stack([x1, y1, x2, y2, x3, y3, x4, y4])) if self.args.show_box: boxes = backward_convert(boxes, False) nake_name = res['image_id'].split('/')[-1] tools.makedirs(os.path.join(save_path, 'icdar2015_img_vis')) draw_path = os.path.join(save_path, 'icdar2015_img_vis', nake_name) draw_img = np.array(cv2.imread(res['image_id']), np.float32) drawer = DrawBox(self.cfgs) final_detections = drawer.draw_boxes_with_label_and_scores(draw_img, boxes=boxes, labels=res['labels'], scores=res['scores'], method=1, in_graph=False) cv2.imwrite(draw_path, final_detections) else: fw_txt_dt = open(os.path.join(save_path, 'icdar2015_res', 'res_{}.txt'.format(res['image_id'].split('/')[-1].split('.')[0])), 'w') for box in boxes: line = '%d,%d,%d,%d,%d,%d,%d,%d\n' % (box[0], box[1], box[2], box[3], box[4], box[5], box[6], box[7]) fw_txt_dt.write(line) fw_txt_dt.close() fw = open(txt_name, 'a+') fw.write('{}\n'.format(res['image_id'].split('/')[-1])) fw.close() pbar.set_description("Test image %s" % res['image_id'].split('/')[-1]) pbar.update(1) for p in procs: p.join()
def __init__(self, cfgs): self.cfgs = cfgs self.drawer = DrawBox(cfgs)
def test_pb(self, frozen_graph_path, test_dir): graph = self.load_graph(frozen_graph_path) print("we are testing ====>>>>", frozen_graph_path) img = graph.get_tensor_by_name("input_img:0") dets = graph.get_tensor_by_name("DetResults:0") with tf.Session(graph=graph) as sess: for img_path in os.listdir(test_dir): print(img_path) a_img = cv2.imread(os.path.join(test_dir, img_path))[:, :, ::-1] raw_h, raw_w = a_img.shape[0], a_img.shape[1] short_size, max_len = self.cfgs.IMG_SHORT_SIDE_LEN, cfgs.IMG_MAX_LENGTH if raw_h < raw_w: new_h, new_w = short_size, min( int(short_size * float(raw_w) / raw_h), max_len) else: new_h, new_w = min(int(short_size * float(raw_h) / raw_w), max_len), short_size img_resize = cv2.resize(a_img, (new_w, new_h)) dets_val = sess.run(dets, feed_dict={img: img_resize[:, :, ::-1]}) bbox_pred, cls_prob = dets_val[:, :5], dets_val[:, 5:( 5 + self.cfgs.CLASS_NUM)] anchor = GenerateAnchors(self.cfgs, 'H') h1, w1 = math.ceil(new_h / 2), math.ceil(new_w / 2) h2, w2 = math.ceil(h1 / 2), math.ceil(w1 / 2) h3, w3 = math.ceil(h2 / 2), math.ceil(w2 / 2) h4, w4 = math.ceil(h3 / 2), math.ceil(w3 / 2) h5, w5 = math.ceil(h4 / 2), math.ceil(w4 / 2) h6, w6 = math.ceil(h5 / 2), math.ceil(w5 / 2) h7, w7 = math.ceil(h6 / 2), math.ceil(w6 / 2) h_dict = {'P3': h3, 'P4': h4, 'P5': h5, 'P6': h6, 'P7': h7} w_dict = {'P3': w3, 'P4': w4, 'P5': w5, 'P6': w6, 'P7': w7} anchors = anchor.generate_all_anchor_pb(h_dict, w_dict) anchors = np.concatenate(anchors, axis=0) x_c = (anchors[:, 2] + anchors[:, 0]) / 2 y_c = (anchors[:, 3] + anchors[:, 1]) / 2 h = anchors[:, 2] - anchors[:, 0] + 1 w = anchors[:, 3] - anchors[:, 1] + 1 theta = -90 * np.ones_like(x_c) anchors = np.transpose(np.stack([x_c, y_c, w, h, theta])) detected_boxes, detected_scores, detected_categories = self.postprocess_detctions( bbox_pred, cls_prob, anchors) if True: # detected_indices = det_scores_r_ >= self.cfgs.VIS_SCORE # detected_scores = det_scores_r_[detected_indices] # detected_boxes = det_boxes_r_[detected_indices] # detected_categories = det_category_r_[detected_indices] drawer = DrawBox(self.cfgs) det_detections_r = drawer.draw_boxes_with_label_and_scores( img_resize[:, :, ::-1], boxes=detected_boxes, labels=detected_categories, scores=detected_scores, method=1, in_graph=True) save_dir = os.path.join('test_pb', self.cfgs.VERSION, 'pb_img_vis') tools.makedirs(save_dir) cv2.imwrite(save_dir + '/{}'.format(img_path), det_detections_r[:, :, ::-1])
def test(self, frozen_graph_path, test_dir): graph = self.load_graph(frozen_graph_path) print("we are testing ====>>>>", frozen_graph_path) img = graph.get_tensor_by_name("input_img:0") dets = graph.get_tensor_by_name("DetResults:0") with tf.Session(graph=graph) as sess: for img_path in os.listdir(test_dir): print(img_path) a_img = cv2.imread(os.path.join(test_dir, img_path))[:, :, ::-1] raw_h, raw_w = a_img.shape[0], a_img.shape[1] short_size, max_len = self.cfgs.IMG_SHORT_SIDE_LEN, cfgs.IMG_MAX_LENGTH if raw_h < raw_w: new_h, new_w = short_size, min( int(short_size * float(raw_w) / raw_h), max_len) else: new_h, new_w = min(int(short_size * float(raw_h) / raw_w), max_len), short_size img_resize = cv2.resize(a_img, (new_w, new_h)) dets_val = sess.run(dets, feed_dict={img: img_resize[:, :, ::-1]}) box_res_rotate_ = [] label_res_rotate_ = [] score_res_rotate_ = [] if dets_val.shape[0] != 0: for sub_class in range(1, self.cfgs.CLASS_NUM + 1): index = np.where(dets_val[:, 0] == sub_class)[0] if len(index) == 0: continue tmp_boxes_r = dets_val[:, 2:][index] tmp_label_r = dets_val[:, 0][index] tmp_score_r = dets_val[:, 1][index] # try: inx = nms_rotate_cpu( boxes=np.array(tmp_boxes_r), scores=np.array(tmp_score_r), iou_threshold=self.cfgs.NMS_IOU_THRESHOLD, max_output_size=20) # except: # tmp_boxes_r_ = np.array(tmp_boxes_r) # tmp = np.zeros([tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1]) # tmp[:, 0:-1] = tmp_boxes_r_ # tmp[:, -1] = np.array(tmp_score_r) # # Note: the IoU of two same rectangles is 0, which is calculated by rotate_gpu_nms # jitter = np.zeros([tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1]) # jitter[:, 0] += np.random.rand(tmp_boxes_r_.shape[0], ) / 1000 # inx = rotate_gpu_nms(np.array(tmp, np.float32) + np.array(jitter, np.float32), # float(self.cfgs.NMS_IOU_THRESHOLD), 0) box_res_rotate_.extend(np.array(tmp_boxes_r)[inx]) score_res_rotate_.extend(np.array(tmp_score_r)[inx]) label_res_rotate_.extend(np.array(tmp_label_r)[inx]) det_boxes_r_ = np.array(box_res_rotate_) det_scores_r_ = np.array(score_res_rotate_) det_category_r_ = np.array(label_res_rotate_) if True: detected_indices = det_scores_r_ >= self.cfgs.VIS_SCORE detected_scores = det_scores_r_[detected_indices] detected_boxes = det_boxes_r_[detected_indices] detected_categories = det_category_r_[detected_indices] drawer = DrawBox(self.cfgs) det_detections_r = drawer.draw_boxes_with_label_and_scores( img_resize[:, :, ::-1], boxes=detected_boxes, labels=detected_categories, scores=detected_scores, method=1, in_graph=True) save_dir = os.path.join('test_pb', self.cfgs.VERSION, 'pb_img_vis') tools.makedirs(save_dir) cv2.imwrite(save_dir + '/{}'.format(img_path), det_detections_r[:, :, ::-1])