def eval_with_plac(det_net, args): # 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) if cfgs.NET_NAME in ['resnet152_v1d', 'resnet101_v1d', 'resnet50_v1d']: img_batch = (img_batch / 255 - tf.constant( cfgs.PIXEL_MEAN_)) / tf.constant(cfgs.PIXEL_STD) else: img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN) img_batch = tf.expand_dims(img_batch, axis=0) detection_scores, detection_category, detection_boxes_angle = det_net.build_whole_detection_network( input_img_batch=img_batch, gtboxes_batch_h=None, gtboxes_batch_r=None, gt_encode_label=None) 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 = [] img_short_side_len_list = cfgs.IMG_SHORT_SIDE_LEN if isinstance( cfgs.IMG_SHORT_SIDE_LEN, list) else [cfgs.IMG_SHORT_SIDE_LEN] img_short_side_len_list = [ img_short_side_len_list[0] ] if not args.multi_scale else img_short_side_len_list imgs = os.listdir(args.img_dir) pbar = tqdm(imgs) for a_img_name in pbar: a_img_name = a_img_name.split(args.image_ext)[0] raw_img = cv2.imread( os.path.join(args.img_dir, a_img_name + args.image_ext)) raw_h, raw_w = raw_img.shape[0], raw_img.shape[1] box_res_rotate = [] label_res_rotate = [] score_res_rotate = [] for short_size in img_short_side_len_list: max_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(raw_img, (new_w, new_h)) resized_img, det_boxes_r_, det_scores_r_, det_category_r_ = \ sess.run( [img_batch, detection_boxes_angle, detection_scores, detection_category], feed_dict={img_plac: img_resize[:, :, ::-1]} ) resized_h, resized_w = resized_img.shape[1], resized_img.shape[ 2] if len(det_boxes_r_) > 0: det_boxes_r_ = forward_convert(det_boxes_r_, False) det_boxes_r_[:, 0::2] *= (raw_w / resized_w) det_boxes_r_[:, 1::2] *= (raw_h / resized_h) for ii in range(len(det_boxes_r_)): box_rotate = det_boxes_r_[ii] box_res_rotate.append(box_rotate) label_res_rotate.append(det_category_r_[ii]) score_res_rotate.append(det_scores_r_[ii]) box_res_rotate = np.array(box_res_rotate) label_res_rotate = np.array(label_res_rotate) score_res_rotate = np.array(score_res_rotate) box_res_rotate_ = [] label_res_rotate_ = [] score_res_rotate_ = [] threshold = {'car': 0.2, 'plane': 0.3} for sub_class in range(1, cfgs.CLASS_NUM + 1): index = np.where(label_res_rotate == sub_class)[0] if len(index) == 0: continue tmp_boxes_r = box_res_rotate[index] tmp_label_r = label_res_rotate[index] tmp_score_r = score_res_rotate[index] 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=threshold[LABEL_NAME_MAP[sub_class]], max_output_size=150) 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(threshold[LABEL_NAME_MAP[sub_class]]), 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]) box_res_rotate_ = np.array(box_res_rotate_) score_res_rotate_ = np.array(score_res_rotate_) label_res_rotate_ = np.array(label_res_rotate_) if args.draw_imgs: detected_indices = score_res_rotate_ >= cfgs.VIS_SCORE detected_scores = score_res_rotate_[detected_indices] detected_boxes = box_res_rotate_[detected_indices] detected_boxes = backward_convert(detected_boxes, with_label=False) detected_categories = label_res_rotate_[detected_indices] det_detections_r = draw_box_in_img.draw_boxes_with_label_and_scores( np.array(raw_img, np.float32), boxes=detected_boxes, labels=detected_categories, scores=detected_scores, method=1, in_graph=False, is_csl=True) save_dir = os.path.join('test_ucas_aod', cfgs.VERSION, 'ucas_aod_img_vis') tools.mkdir(save_dir) cv2.imwrite(save_dir + '/{}.jpg'.format(a_img_name), det_detections_r[:, :, ::-1]) if box_res_rotate_.shape[0] != 0: box_res_rotate_ = backward_convert(box_res_rotate_, False) x_c, y_c, w, h, theta = box_res_rotate_[:, 0], box_res_rotate_[:, 1], box_res_rotate_[:, 2], \ box_res_rotate_[:, 3], box_res_rotate_[:, 4] boxes_r = np.transpose(np.stack([x_c, y_c, w, h, theta])) dets_r = np.hstack((label_res_rotate_.reshape(-1, 1), score_res_rotate_.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 detect(det_net, inference_save_path, real_test_imgname_list): # 1. preprocess img img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3]) # is RGB. not GBR img_batch = tf.cast(img_plac, tf.float32) img_batch = short_side_resize_for_inference_data( img_tensor=img_batch, target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN, length_limitation=cfgs.IMG_MAX_LENGTH) img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN) img_batch = tf.expand_dims(img_batch, axis=0) # [1, None, None, 3] detection_boxes, detection_scores, detection_category = det_net.build_whole_detection_network( input_img_batch=img_batch, gtboxes_batch=None) 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') for i, a_img_name in enumerate(real_test_imgname_list): raw_img = cv2.imread(a_img_name) start = time.time() resized_img, detected_boxes, detected_scores, detected_categories = \ sess.run( [img_batch, detection_boxes, detection_scores, detection_category], feed_dict={img_plac: raw_img} ) end = time.time() # print("{} cost time : {} ".format(img_name, (end - start))) raw_h, raw_w = raw_img.shape[0], raw_img.shape[1] xmin, ymin, xmax, ymax = detected_boxes[:, 0], detected_boxes[:, 1], \ detected_boxes[:, 2], detected_boxes[:, 3] resized_h, resized_w = resized_img.shape[1], resized_img.shape[2] xmin = xmin * raw_w / resized_w xmax = xmax * raw_w / resized_w ymin = ymin * raw_h / resized_h ymax = ymax * raw_h / resized_h detected_boxes = np.transpose(np.stack([xmin, ymin, xmax, ymax])) show_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD show_scores = detected_scores[show_indices] show_boxes = detected_boxes[show_indices] show_categories = detected_categories[show_indices] final_detections = draw_box_in_img.draw_boxes_with_label_and_scores( raw_img - np.array(cfgs.PIXEL_MEAN), boxes=show_boxes, labels=show_categories, scores=show_scores) nake_name = a_img_name.split('/')[-1] # print (inference_save_path + '/' + nake_name) cv2.imwrite(inference_save_path + '/' + nake_name, final_detections[:, :, ::-1]) tools.view_bar( '{} image cost {}s'.format(a_img_name, (end - start)), i + 1, len(real_test_imgname_list))
def eval_with_plac(det_net, real_test_imgname_list, img_root, draw_imgs=False): # 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) img_batch = short_side_resize_for_inference_data(img_tensor=img_batch, target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN, length_limitation=cfgs.IMG_MAX_LENGTH) if cfgs.NET_NAME in ['resnet101_v1d', 'resnet50_v1d']: img_batch = (img_batch / 255 - tf.constant(cfgs.PIXEL_MEAN_)) / tf.constant(cfgs.PIXEL_STD) else: img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN) # img_batch = (img_batch - tf.constant(cfgs.PIXEL_MEAN)) / (tf.constant(cfgs.PIXEL_STD)*255) 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, gtboxes_batch=None) 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') for i, a_img_name in enumerate(real_test_imgname_list): raw_img = cv2.imread(os.path.join(img_root, a_img_name)) raw_h, raw_w = raw_img.shape[0], raw_img.shape[1] start = time.time() resized_img, detected_boxes, detected_scores, detected_categories = \ sess.run( [img_batch, detection_boxes, detection_scores, detection_category], feed_dict={img_plac: raw_img[:, :, ::-1]} # cv is BGR. But need RGB ) end = time.time() # print("{} cost time : {} ".format(img_name, (end - start))) if draw_imgs: show_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD show_scores = detected_scores[show_indices] show_boxes = detected_boxes[show_indices] show_categories = detected_categories[show_indices] draw_img = np.squeeze(resized_img, 0) if cfgs.NET_NAME in ['resnet101_v1d', 'resnet50_v1d']: draw_img = (draw_img * np.array(cfgs.PIXEL_STD) + np.array(cfgs.PIXEL_MEAN_)) * 255 else: draw_img = draw_img + np.array(cfgs.PIXEL_MEAN) # draw_img = draw_img * (np.array(cfgs.PIXEL_STD)*255) + np.array(cfgs.PIXEL_MEAN) final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(draw_img, boxes=show_boxes, labels=show_categories, scores=show_scores, in_graph=False) if not os.path.exists(cfgs.TEST_SAVE_PATH): os.makedirs(cfgs.TEST_SAVE_PATH) cv2.imwrite(cfgs.TEST_SAVE_PATH + '/' + a_img_name, final_detections[:, :, ::-1]) xmin, ymin, xmax, ymax = detected_boxes[:, 0], detected_boxes[:, 1], \ detected_boxes[:, 2], detected_boxes[:, 3] resized_h, resized_w = resized_img.shape[1], resized_img.shape[2] xmin = xmin * raw_w / resized_w xmax = xmax * raw_w / resized_w ymin = ymin * raw_h / resized_h ymax = ymax * raw_h / resized_h boxes = np.transpose(np.stack([xmin, ymin, xmax, ymax])) dets = np.hstack((detected_categories.reshape(-1, 1), detected_scores.reshape(-1, 1), boxes)) # all_boxes.append(dets) # eval txt CLASS_VOC = NAME_LABEL_MAP.keys() write_handle = {} txt_dir = os.path.join('voc2012_eval', cfgs.VERSION, 'results', 'VOC2012', 'Main') tools.mkdir(txt_dir) for sub_class in CLASS_VOC: if sub_class == 'back_ground': continue write_handle[sub_class] = open(os.path.join(txt_dir, 'comp3_det_test_%s.txt' % sub_class), 'a+') for det in dets: command = '%s %.6f %.6f %.6f %.6f %.6f\n' % (a_img_name.split('/')[-1].split('.')[0], det[1], det[2], det[3], det[4], det[5]) write_handle[LABEl_NAME_MAP[det[0]]].write(command) for sub_class in CLASS_VOC: if sub_class == 'back_ground': continue write_handle[sub_class].close() tools.view_bar('%s image cost %.3fs' % (a_img_name, (end - start)), i + 1, len(real_test_imgname_list))
def test_icdar2015(det_net, real_test_img_list, gpu_ids, show_box, txt_name): save_path = os.path.join('./test_icdar2015', cfgs.VERSION) tools.mkdir(save_path) nr_records = len(real_test_img_list) pbar = tqdm(total=nr_records) gpu_num = len(gpu_ids.strip().split(',')) nr_image = math.ceil(nr_records / gpu_num) result_queue = Queue(500) procs = [] for i, gpu_id in enumerate(gpu_ids.strip().split(',')): start = i * nr_image end = min(start + nr_image, nr_records) split_records = real_test_img_list[start:end] proc = Process(target=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 res['boxes'].shape[0] == 0: fw_txt_dt = open( os.path.join( save_path, '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 show_box: boxes = backward_convert(boxes, False) nake_name = res['image_id'].split('/')[-1] draw_path = os.path.join(save_path, nake_name) draw_img = np.array(cv2.imread(res['image_id']), np.float32) final_detections = draw_box_in_img.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, '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 test_dota(det_net, real_test_img_list, args, txt_name): save_path = os.path.join('./test_dota', cfgs.VERSION) nr_records = len(real_test_img_list) pbar = tqdm(total=nr_records) gpu_num = len(args.gpus.strip().split(',')) nr_image = math.ceil(nr_records / gpu_num) result_queue = Queue(500) procs = [] for i, gpu_id in enumerate(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=worker, args=(int(gpu_id), split_records, det_net, args, result_queue)) print('process:%d, start:%d, end:%d' % (i, start, end)) proc.start() procs.append(proc) log_dir = './csl_log/{}'.format(cfgs.VERSION) tools.mkdir(log_dir) fw_tsv = open(os.path.join(log_dir, 'csl_meta.tsv'), 'w') # fw_tsv.write("Label\n") final_logits = [] for i in range(nr_records): res = result_queue.get() if args.show_box: nake_name = res['image_id'].split('/')[-1] tools.mkdir(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'] >= cfgs.VIS_SCORE detected_scores = res['scores'][detected_indices] detected_boxes = detected_boxes[detected_indices] detected_categories = res['labels'][detected_indices] final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(draw_img, boxes=detected_boxes, labels=detected_categories, scores=detected_scores, method=1, head=np.ones_like(detected_scores) * -1, is_csl=True, in_graph=False) cv2.imwrite(draw_path, final_detections) else: detected_indices = res['scores'] >= cfgs.VIS_SCORE res['scores'] = res['scores'][detected_indices] res['boxes'] = res['boxes'][detected_indices] res['labels'] = res['labels'][detected_indices] rboxes = backward_convert(res['boxes'], with_label=False) rboxes = coordinate_present_convert(rboxes, -1, False) rlogits = res['logits'][detected_indices] for ii, rb in enumerate(rboxes): fw_tsv.write("%d\n" % (int(rb[-1]))) final_logits.append(rlogits[ii]) 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() fw_tsv.close() final_logits = np.array(final_logits) np.save(os.path.join(log_dir, "final_logits.npy"), final_logits)
def detect(det_net, inference_save_path, real_test_imgname_list): # 1. preprocess img img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3]) # is RGB. not GBR img_batch = tf.cast(img_plac, tf.float32) img_batch = short_side_resize_for_inference_data( img_tensor=img_batch, target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN, length_limitation=cfgs.IMG_MAX_LENGTH) if cfgs.NET_NAME in ['resnet152_v1d', 'resnet101_v1d', 'resnet50_v1d']: img_batch = (img_batch / 255 - tf.constant( cfgs.PIXEL_MEAN_)) / tf.constant(cfgs.PIXEL_STD) else: img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN) img_batch = tf.expand_dims(img_batch, axis=0) # [1, None, None, 3] detection_boxes, detection_scores, detection_category = det_net.build_whole_detection_network( input_img_batch=img_batch, gtboxes_batch_h=None, gtboxes_batch_r=None) 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') for i, a_img_name in enumerate(real_test_imgname_list): raw_img = cv2.imread(a_img_name) start = time.time() resized_img, detected_boxes, detected_scores, detected_categories = \ sess.run( [img_batch, detection_boxes, detection_scores, detection_category], feed_dict={img_plac: raw_img[:, :, ::-1]} # cv is BGR. But need RGB ) end = time.time() # print("{} cost time : {} ".format(img_name, (end - start))) show_indices = detected_scores >= cfgs.VIS_SCORE show_scores = detected_scores[show_indices] show_boxes = detected_boxes[show_indices] show_categories = detected_categories[show_indices] draw_img = np.squeeze(resized_img, 0) if cfgs.NET_NAME in [ 'resnet152_v1d', 'resnet101_v1d', 'resnet50_v1d' ]: draw_img = (draw_img * np.array(cfgs.PIXEL_STD) + np.array(cfgs.PIXEL_MEAN_)) * 255 else: draw_img = draw_img + np.array(cfgs.PIXEL_MEAN) final_detections = draw_box_in_img.draw_boxes_with_label_and_scores( draw_img, boxes=show_boxes, labels=show_categories, scores=show_scores, method=1, in_graph=False) nake_name = a_img_name.split('/')[-1] # print (inference_save_path + '/' + nake_name) cv2.imwrite(inference_save_path + '/' + nake_name, final_detections[:, :, ::-1]) tools.view_bar( '{} image cost {}s'.format(nake_name, (end - start)), i + 1, len(real_test_imgname_list))
def test_dota(det_net, real_test_img_list, args, txt_name): save_path = os.path.join('./test_dota', cfgs.VERSION) nr_records = len(real_test_img_list) pbar = tqdm(total=nr_records) gpu_num = len(args.gpus.strip().split(',')) nr_image = math.ceil(nr_records / gpu_num) result_queue = Queue(500) procs = [] for i, gpu_id in enumerate(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=worker, args=(int(gpu_id), split_records, det_net, args, 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 args.show_box: nake_name = res['image_id'].split('/')[-1] tools.mkdir(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_indices = res['scores'] >= cfgs.VIS_SCORE detected_scores = res['scores'][detected_indices] detected_boxes = res['boxes'][detected_indices] detected_categories = res['labels'][detected_indices] final_detections = draw_box_in_img.draw_boxes_with_label_and_scores( draw_img, boxes=detected_boxes, labels=detected_categories, scores=detected_scores, method=1, in_graph=False) cv2.imwrite(draw_path, final_detections) else: CLASS_DOTA = NAME_LABEL_MAP.keys() write_handle = {} tools.mkdir(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+') rboxes = forward_convert(res['boxes'], with_label=False) for i, rbox in enumerate(rboxes): 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[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 detect(det_net, src_dir, res_dir, draw_imgs): # 1. preprocess img img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3]) # is RGB. not GBR img_batch = tf.cast(img_plac, tf.float32) img_batch = short_side_resize_for_inference_data( img_tensor=img_batch, target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN, length_limitation=cfgs.IMG_MAX_LENGTH) img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN) img_batch = tf.expand_dims(img_batch, axis=0) # [1, None, None, 3] result_dict = det_net.build_whole_detection_network( input_img_batch=img_batch, gtboxes_batch=None) 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') sub_folders = os.listdir(src_dir) for sub_folder in sub_folders: folder_dir = os.path.join(src_dir, sub_folder) real_test_imgname_list = [ os.path.join(folder_dir, img_name) for img_name in os.listdir(folder_dir) ] tools.mkdir(os.path.join(res_dir, sub_folder)) for i, a_img_name in enumerate(real_test_imgname_list): raw_img = cv2.imread(a_img_name) raw_h, raw_w = raw_img.shape[0], raw_img.shape[1] start = time.time() resized_img, result_dict_ = \ sess.run( [img_batch, result_dict], feed_dict={img_plac: raw_img[:, :, ::-1]} # cv is BGR. But need RGB ) end = time.time() detected_boxes, detected_scores, detected_categories = merge_result( result_dict_) nake_name = a_img_name.split('/')[-1] xmin, ymin, xmax, ymax = detected_boxes[:, 0], detected_boxes[:, 1], \ detected_boxes[:, 2], detected_boxes[:, 3] resized_h, resized_w = resized_img.shape[1], resized_img.shape[ 2] xmin = xmin * raw_w / resized_w xmax = xmax * raw_w / resized_w ymin = ymin * raw_h / resized_h ymax = ymax * raw_h / resized_h boxes = np.transpose(np.stack([xmin, ymin, xmax, ymax])) dets = np.hstack((detected_categories.reshape(-1, 1), detected_scores.reshape(-1, 1), boxes)) show_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD show_scores = detected_scores[show_indices] show_boxes = boxes[show_indices] show_categories = detected_categories[show_indices] f = open( os.path.join(res_dir, sub_folder) + '/' + nake_name.split('.')[0] + '.txt', 'w') f.write('{}\n'.format(nake_name.split('.')[0])) # f.write('{}\n'.format(dets.shape[0])) # for inx in range(dets.shape[0]): # # f.write('%d %d %d %d %.3f\n' % (int(dets[inx][2]), # int(dets[inx][3]), # int(dets[inx][4]) - int(dets[inx][2]), # int(dets[inx][5]) - int(dets[inx][3]), # dets[inx][1])) f.write('{}\n'.format(show_boxes.shape[0])) for inx in range(show_boxes.shape[0]): f.write('%d %d %d %d %.3f\n' % (int(show_boxes[inx][0]), int(show_boxes[inx][1]), int(show_boxes[inx][2]) - int(show_boxes[inx][0]), int(show_boxes[inx][3]) - int(show_boxes[inx][1]), show_scores[inx])) if draw_imgs: final_detections = draw_box_in_img.draw_boxes_with_label_and_scores( raw_img - np.array(cfgs.PIXEL_MEAN), boxes=show_boxes, labels=show_categories, scores=show_scores) tools.mkdir(cfgs.TEST_SAVE_PATH) cv2.imwrite(cfgs.TEST_SAVE_PATH + '/' + nake_name, final_detections) tools.view_bar( '{} image cost {}s'.format(a_img_name, (end - start)), i + 1, len(real_test_imgname_list))
def eval_with_plac(det_net, real_test_imgname_list, img_root, draw_imgs=False): # 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) img_batch = short_side_resize_for_inference_data(img_tensor=img_batch, target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN, length_limitation=cfgs.IMG_MAX_LENGTH) img_batch = img_batch - tf.constant(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, gtboxes_batch=None) 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 = [] for i, a_img_name in enumerate(real_test_imgname_list): raw_img = cv2.imread(os.path.join(img_root, a_img_name)) raw_h, raw_w = raw_img.shape[0], raw_img.shape[1] start = time.time() resized_img, detected_boxes, detected_scores, detected_categories = \ sess.run( [img_batch, detection_boxes, detection_scores, detection_category], feed_dict={img_plac: raw_img[:, :, ::-1]} # cv is BGR. But need RGB ) end = time.time() # print("{} cost time : {} ".format(img_name, (end - start))) if draw_imgs: show_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD show_scores = detected_scores[show_indices] show_boxes = detected_boxes[show_indices] show_categories = detected_categories[show_indices] final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(np.squeeze(resized_img, 0), boxes=show_boxes, labels=show_categories, scores=show_scores) if not os.path.exists(cfgs.TEST_SAVE_PATH): os.makedirs(cfgs.TEST_SAVE_PATH) cv2.imwrite(cfgs.TEST_SAVE_PATH + '/' + a_img_name, final_detections[:, :, ::-1]) xmin, ymin, xmax, ymax = detected_boxes[:, 0], detected_boxes[:, 1], \ detected_boxes[:, 2], detected_boxes[:, 3] resized_h, resized_w = resized_img.shape[1], resized_img.shape[2] xmin = xmin * raw_w / resized_w xmax = xmax * raw_w / resized_w ymin = ymin * raw_h / resized_h ymax = ymax * raw_h / resized_h boxes = np.transpose(np.stack([xmin, ymin, xmax, ymax])) dets = np.hstack((detected_categories.reshape(-1, 1), detected_scores.reshape(-1, 1), boxes)) all_boxes.append(dets) tools.view_bar('{} image cost {}s'.format(a_img_name, (end - start)), i + 1, len(real_test_imgname_list)) save_dir = os.path.join(cfgs.EVALUATE_DIR, cfgs.VERSION) if not os.path.exists(save_dir): os.makedirs(save_dir) fw1 = open(os.path.join(save_dir, 'detections.pkl'), 'wb') pickle.dump(all_boxes, fw1)
def eval_coco(det_net, real_test_img_list, draw_imgs=False): # 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) img_batch = short_side_resize_for_inference_data( img_tensor=img_batch, target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN, length_limitation=cfgs.IMG_MAX_LENGTH) if cfgs.NET_NAME in ['resnet101_v1d', 'resnet50_v1d']: img_batch = (img_batch / 255 - tf.constant( cfgs.PIXEL_MEAN_)) / tf.constant(cfgs.PIXEL_STD) else: img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN) # img_batch = (img_batch - tf.constant(cfgs.PIXEL_MEAN)) / (tf.constant(cfgs.PIXEL_STD)*255) 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, gtboxes_batch=None) 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') save_path = os.path.join('./eval_coco', cfgs.VERSION) tools.mkdir(save_path) fw_json_dt = open(os.path.join(save_path, 'coco_minival.json'), 'w') coco_det = [] for i, a_img in enumerate(real_test_img_list): record = json.loads(a_img) img_path = os.path.join('/data/COCO/val2017', record['fpath'].split('_')[-1]) raw_img = cv2.imread(img_path) # raw_img = cv2.imread(record['fpath']) raw_h, raw_w = raw_img.shape[0], raw_img.shape[1] start = time.time() resized_img, detected_boxes, detected_scores, detected_categories = \ sess.run( [img_batch, detection_boxes, detection_scores, detection_category], feed_dict={img_plac: raw_img[:, :, ::-1]} # cv is BGR. But need RGB ) end = time.time() eval_indices = detected_scores >= 0.01 detected_scores = detected_scores[eval_indices] detected_boxes = detected_boxes[eval_indices] detected_categories = detected_categories[eval_indices] # print("{} cost time : {} ".format(img_name, (end - start))) if draw_imgs: show_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD show_scores = detected_scores[show_indices] show_boxes = detected_boxes[show_indices] show_categories = detected_categories[show_indices] draw_img = np.squeeze(resized_img, 0) # draw_img = draw_img + np.array(cfgs.PIXEL_MEAN) # draw_img = draw_img * (np.array(cfgs.PIXEL_STD)*255) + np.array(cfgs.PIXEL_MEAN) final_detections = draw_box_in_img.draw_boxes_with_label_and_scores( draw_img, boxes=show_boxes, labels=show_categories, scores=show_scores, in_graph=False) if not os.path.exists(cfgs.TEST_SAVE_PATH): os.makedirs(cfgs.TEST_SAVE_PATH) cv2.imwrite(cfgs.TEST_SAVE_PATH + '/' + record['ID'], final_detections[:, :, ::-1]) xmin, ymin, xmax, ymax = detected_boxes[:, 0], detected_boxes[:, 1], \ detected_boxes[:, 2], detected_boxes[:, 3] resized_h, resized_w = resized_img.shape[1], resized_img.shape[2] xmin = xmin * raw_w / resized_w xmax = xmax * raw_w / resized_w ymin = ymin * raw_h / resized_h ymax = ymax * raw_h / resized_h boxes = np.transpose( np.stack([xmin, ymin, xmax - xmin, ymax - ymin])) # cost much time for j, box in enumerate(boxes): coco_det.append({ 'bbox': [ float(box[0]), float(box[1]), float(box[2]), float(box[3]) ], 'score': float(detected_scores[j]), 'image_id': int(record['ID'].split('.jpg')[0].split('_000000')[-1]), 'category_id': int(classes_originID[LABEl_NAME_MAP[ detected_categories[j]]]) }) tools.view_bar( '%s image cost %.3fs' % (record['ID'], (end - start)), i + 1, len(real_test_img_list)) json.dump(coco_det, fw_json_dt) fw_json_dt.close() return os.path.join(save_path, 'coco_minival.json')
img_name = xml_name.replace('.xml', '.jpg') img = cv2.imread(os.path.join(root_path, 'train', img_name)) img_height, img_width, gtbox_label = read_xml_gtbox_and_label( os.path.join(root_path, 'Annotations', xml_name)) gtbox_and_label_list = np.array(gtbox_label, dtype=np.int32) if gtbox_and_label_list.shape[0] == 0: return next_img(step + 1) else: return img_name, img[:, :, ::-1], gtbox_and_label_list if __name__ == '__main__': imgid, img, gtbox = next_img(3234) print("::") from libs.box_utils.draw_box_in_img import draw_boxes_with_label_and_scores img = draw_boxes_with_label_and_scores( img_array=img, boxes=gtbox[:, :-1], labels=gtbox[:, -1], scores=np.ones(shape=(len(gtbox), ))) print("_----") cv2.imshow("test", img) cv2.waitKey(0)
def eval_with_plac(det_net, real_test_imgname_list, img_root, draw_imgs=False): # 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) img_batch = short_side_resize_for_inference_data(img_tensor=img_batch, target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN, length_limitation=cfgs.IMG_MAX_LENGTH) if cfgs.NET_NAME in ['resnet152_v1d', 'resnet101_v1d', 'resnet50_v1d']: img_batch = (img_batch / 255 - tf.constant(cfgs.PIXEL_MEAN_)) / tf.constant(cfgs.PIXEL_STD) else: img_batch = img_batch - tf.constant(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, gtboxes_batch=None) 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 = [] pbar = tqdm(real_test_imgname_list) for a_img_name in pbar: raw_img = cv2.imread(os.path.join(img_root, a_img_name)) raw_h, raw_w = raw_img.shape[0], raw_img.shape[1] resized_img, detected_boxes, detected_scores, detected_categories = \ sess.run( [img_batch, detection_boxes, detection_scores, detection_category], feed_dict={img_plac: raw_img[:, :, ::-1]} # cv is BGR. But need RGB ) if draw_imgs: show_indices = detected_scores >= cfgs.VIS_SCORE show_scores = detected_scores[show_indices] show_boxes = detected_boxes[show_indices] show_categories = detected_categories[show_indices] draw_img = np.squeeze(resized_img, 0) if cfgs.NET_NAME in ['resnet152_v1d', 'resnet101_v1d', 'resnet50_v1d']: draw_img = (draw_img * np.array(cfgs.PIXEL_STD) + np.array(cfgs.PIXEL_MEAN_)) * 255 else: draw_img = draw_img + np.array(cfgs.PIXEL_MEAN) final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(draw_img, boxes=show_boxes, labels=show_categories, scores=show_scores, in_graph=False) if not os.path.exists(cfgs.TEST_SAVE_PATH): os.makedirs(cfgs.TEST_SAVE_PATH) cv2.imwrite(cfgs.TEST_SAVE_PATH + '/' + a_img_name + '.jpg', final_detections[:, :, ::-1]) xmin, ymin, xmax, ymax = detected_boxes[:, 0], detected_boxes[:, 1], \ detected_boxes[:, 2], detected_boxes[:, 3] resized_h, resized_w = resized_img.shape[1], resized_img.shape[2] xmin = xmin * raw_w / resized_w xmax = xmax * raw_w / resized_w ymin = ymin * raw_h / resized_h ymax = ymax * raw_h / resized_h boxes = np.transpose(np.stack([xmin, ymin, xmax, ymax])) dets = np.hstack((detected_categories.reshape(-1, 1), detected_scores.reshape(-1, 1), boxes)) all_boxes.append(dets) pbar.set_description("Eval image %s" % a_img_name) # save_dir = os.path.join(cfgs.EVALUATE_DIR, cfgs.VERSION) # if not os.path.exists(save_dir): # os.makedirs(save_dir) # fw1 = open(os.path.join(save_dir, 'detections.pkl'), 'w') # pickle.dump(all_boxes, fw1) return all_boxes
def exucute_detect(self, image_path, save_path): """ execute object detect :param detect_net: :param image_path: :return: """ input_image = tf.placeholder(dtype=tf.uint8, shape=(None, None, 3), name='inputs_images') resize_img = self.image_process(input_image) # expend dimension image_batch = tf.expand_dims(input=resize_img, axis=0) # (1, None, None, 3) self.detect_net.images_batch = image_batch # img_shape = tf.shape(inputs_img) # load detect network detection_boxes, detection_scores, detection_category = self.detect_net.inference( ) # restore pretrain weight restorer, restore_ckpt = self.detect_net.get_restorer() # config gpu to growth train config = tf.ConfigProto() config.gpu_options.allow_growth = True init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) with tf.Session(config=config) as sess: sess.run(init_op) if restorer is not None: restorer.restore(sess, save_path=restore_ckpt) print('*' * 80 + '\nSuccessful restore model from {0}\n'.format( restore_ckpt) + '*' * 80) # construct image path list format_list = ('.jpg', '.png', '.jpeg', '.tif', '.tiff') if os.path.isfile(image_path): image_name_list = [image_path] else: image_name_list = [ img_name for img_name in os.listdir(image_path) if img_name.endswith(format_list) and os.path.isfile(os.path.join(image_path, img_name)) ] assert len(image_name_list) != 0 print( "test_dir has no imgs there. Note that, we only support img format of {0}" .format(format_list)) #+++++++++++++++++++++++++++++++++++++start detect+++++++++++++++++++++++++++++++++++++++++++++++++++++=++ makedir(save_path) fw = open(os.path.join(save_path, 'detect_bbox.txt'), 'w') for index, img_name in enumerate(image_name_list): detect_dict = {} bgr_img = cv.imread(os.path.join(image_path, img_name)) rgb_img = cv.cvtColor( bgr_img, cv.COLOR_BGR2RGB ) # convert channel from BGR to RGB (cv is BGR) start_time = time.perf_counter() # image resize and white process # construct feed_dict feed_dict = {input_image: rgb_img} resized_img, detected_boxes, detected_scores, detected_categories = \ sess.run([resize_img, detection_boxes, detection_scores, detection_category], feed_dict=feed_dict) end_time = time.perf_counter() # select object according to threshold object_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD object_scores = detected_scores[object_indices] object_boxes = detected_boxes[object_indices] object_categories = detected_categories[object_indices] final_detections_img = draw_box_in_img.draw_boxes_with_label_and_scores( resized_img, boxes=object_boxes, labels=object_categories, scores=object_scores) final_detections_img = cv.cvtColor(final_detections_img, cv.COLOR_RGB2BGR) cv.imwrite(os.path.join(save_path, img_name), final_detections_img) # resize boxes and image according to raw input image raw_h, raw_w = rgb_img.shape[0], rgb_img.shape[1] resized_h, resized_w = resized_img.shape[1], resized_img.shape[ 2] x_min, y_min, x_max, y_max = object_boxes[:, 0], object_boxes[:, 1], object_boxes[:, 2], \ object_boxes[:, 3] x_min = x_min * raw_w / resized_w y_min = y_min * raw_h / resized_h x_max = x_max * raw_w / resized_w y_max = y_max * raw_h / resized_h object_boxes = np.stack([x_min, y_min, x_max, y_max], axis=1) # final_detections= cv.resize(final_detections[:, :, ::-1], (raw_w, raw_h)) # recover to raw size detect_dict['score'] = object_scores detect_dict['boxes'] = object_boxes detect_dict['categories'] = object_categories # convert from RGB to BG fw.write(f'\n{img_name}') for score, boxes, categories in zip(object_scores, object_boxes, object_categories): fw.write('\n\tscore:' + str(score)) fw.write('\tbboxes:' + str(boxes)) fw.write('\tcategories:' + str(categories)) view_bar( '{} image cost {} second'.format(img_name, (end_time - start_time)), index + 1, len(image_name_list)) fw.close()
def eval_with_plac(img_dir, det_net, num_imgs, image_ext, draw_imgs=False): # 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) img_batch = short_side_resize_for_inference_data(img_tensor=img_batch, target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN, length_limitation=cfgs.IMG_MAX_LENGTH) if cfgs.NET_NAME in ['resnet152_v1d', 'resnet101_v1d', 'resnet50_v1d']: img_batch = (img_batch / 255 - tf.constant(cfgs.PIXEL_MEAN_)) / tf.constant(cfgs.PIXEL_STD) else: img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN) img_batch = tf.expand_dims(img_batch, axis=0) detection_scores, detection_category, detection_boxes_angle = det_net.build_whole_detection_network( input_img_batch=img_batch, gtboxes_batch_h=None, gtboxes_batch_r=None, gt_encode_label=None) 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] resized_img, det_boxes_r_, det_scores_r_, det_category_r_ = \ sess.run( [img_batch, detection_boxes_angle, detection_scores, detection_category], feed_dict={img_plac: raw_img[:, :, ::-1]} ) if draw_imgs: detected_indices = det_scores_r_ >= cfgs.VIS_SCORE detected_scores = det_scores_r_[detected_indices] detected_boxes = det_boxes_r_[detected_indices] detected_categories = det_category_r_[detected_indices] det_detections_r = draw_box_in_img.draw_boxes_with_label_and_scores(np.squeeze(resized_img, 0), boxes=detected_boxes, labels=detected_categories, scores=detected_scores, method=1, in_graph=True) save_dir = os.path.join('test_hrsc', cfgs.VERSION, 'hrsc2016_img_vis') tools.mkdir(save_dir) cv2.imwrite(save_dir + '/{}.jpg'.format(a_img_name), det_detections_r[:, :, ::-1]) if det_boxes_r_.shape[0] != 0: resized_h, resized_w = resized_img.shape[1], resized_img.shape[2] det_boxes_r_ = forward_convert(det_boxes_r_, False) det_boxes_r_[:, 0::2] *= (raw_w / resized_w) det_boxes_r_[:, 1::2] *= (raw_h / resized_h) 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 exucute_detect(self, img_dir, img_name_list): """ execute object detect :param detect_net: detect network :param img_list: the image dir of detect :return: """ # config gpu to growth train config = tf.ConfigProto() config.gpu_options.allow_growth = True init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) with tf.Session() as sess: sess.run(init_op) # restore pretrain weight restorer, restore_ckpt = self.detect_net.get_restore( pretrained_model_dir=self.pretrain_model_dir, restore_from_rpn=False, is_pretrained=True) if not restorer is None: restorer.restore(sess, save_path=restore_ckpt) print('Successful restore model from {0}'.format(restore_ckpt)) # +++++++++++++++++++++++++++++++++++++start detect+++++++++++++++++++++++++++++++++++++++++++++++++++++=++ all_boxes = [] img_path_list = [ os.path.join(img_dir, img_name) for img_name in img_name_list ] for index, img_name in enumerate(img_path_list): bgr_img = cv.imread(img_name) raw_img = cv.cvtColor(bgr_img, cv.COLOR_BGR2RGB) resized_img = self.image_process(raw_img) # expend dimension image_batch = tf.expand_dims(input=resized_img, axis=0) # (1, None, None, 3) start_time = time.time() feed_dict = self.detect_net.fill_feed_dict( image_feed=image_batch.eval()) resized_img, (detected_boxes, detected_scores, detected_categories) = \ sess.run(fetches=[resized_img, self.detect_net.inference], feed_dict=feed_dict) # convert channel from BGR to RGB (cv is BGR) end_time = time.time() print("{} cost time : {} ".format(img_name, (end_time - start_time))) # draw object image if self.draw_img: # select object according to threshold object_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD object_scores = detected_scores[object_indices] object_boxes = detected_boxes[object_indices] object_categories = detected_categories[object_indices] final_detections = draw_box_in_img.draw_boxes_with_label_and_scores( img_array=resized_img, boxes=object_boxes, labels=object_categories, scores=object_scores) final_detections = cv.cvtColor(final_detections, cv.COLOR_RGB2BGR) object_boxes = self.bbox_resize( bbox=object_boxes, inputs_shape=resized_img.shape[1:3], target_shape=raw_img.shape[1:3]) # resize boxes and image shape size to raw input image detected_boxes = self.bbox_resize( bbox=detected_boxes, inputs_shape=(resized_img.shape[0], resized_img.shape[1]), target_shape=(raw_img.shape[0], resized_img.shape[1])) # construct detect array for evaluation detect_bbox_label = np.hstack( (detected_categories.reshape(-1, 1).astype(np.int32), detected_scores.reshape(-1, 1), detected_boxes)) all_boxes.append(detect_bbox_label) # dump bbox to local makedir(self.object_bbox_save_path) with open( os.path.join(self.object_bbox_save_path, 'detections.pkl'), 'wb') as fw: pickle.dump(all_boxes, fw)
def inference(det_net, file_paths, des_folder, h_len, w_len, h_overlap, w_overlap, save_res=False): if save_res: assert cfgs.SHOW_SCORE_THRSHOLD >= 0.5, \ 'please set score threshold (example: SHOW_SCORE_THRSHOLD = 0.5) in cfgs.py' else: assert cfgs.SHOW_SCORE_THRSHOLD <= 0.005, \ 'please set score threshold (example: SHOW_SCORE_THRSHOLD = 0.00) in cfgs.py' # 1. preprocess img img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3]) img_batch = tf.cast(img_plac, tf.float32) if cfgs.NET_NAME in ['resnet101_v1d']: img_batch = (img_batch / 255 - tf.constant( cfgs.PIXEL_MEAN_)) / tf.constant(cfgs.PIXEL_STD) else: img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN) img_batch = tf.expand_dims(img_batch, axis=0) img_batch = short_side_resize_for_inference_data( img_tensor=img_batch, target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN, is_resize=False) det_boxes_h, det_scores_h, det_category_h = det_net.build_whole_detection_network( input_img_batch=img_batch, gtboxes_batch=None) 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') if not os.path.exists('./tmp.txt'): fw = open('./tmp.txt', 'w') fw.close() fr = open('./tmp.txt', 'r') pass_img = fr.readlines() fr.close() for count, img_path in enumerate(file_paths): fw = open('./tmp.txt', 'a+') if img_path + '\n' in pass_img: continue start = timer() img = cv2.imread(img_path) box_res = [] label_res = [] score_res = [] imgH = img.shape[0] imgW = img.shape[1] if imgH < h_len: temp = np.zeros([h_len, imgW, 3], np.float32) temp[0:imgH, :, :] = img img = temp imgH = h_len if imgW < w_len: temp = np.zeros([imgH, w_len, 3], np.float32) temp[:, 0:imgW, :] = img img = temp imgW = w_len for hh in range(0, imgH, h_len - h_overlap): if imgH - hh - 1 < h_len: hh_ = imgH - h_len else: hh_ = hh for ww in range(0, imgW, w_len - w_overlap): if imgW - ww - 1 < w_len: ww_ = imgW - w_len else: ww_ = ww src_img = img[hh_:(hh_ + h_len), ww_:(ww_ + w_len), :] det_boxes_h_, det_scores_h_, det_category_h_ = \ sess.run( [det_boxes_h, det_scores_h, det_category_h], feed_dict={img_plac: src_img[:, :, ::-1]} ) if len(det_boxes_h_) > 0: for ii in range(len(det_boxes_h_)): box = det_boxes_h_[ii] box[0] = box[0] + ww_ box[1] = box[1] + hh_ box[2] = box[2] + ww_ box[3] = box[3] + hh_ box_res.append(box) label_res.append(det_category_h_[ii]) score_res.append(det_scores_h_[ii]) box_res = np.array(box_res) label_res = np.array(label_res) score_res = np.array(score_res) box_res_, label_res_, score_res_ = [], [], [] h_threshold = { 'roundabout': 0.35, 'tennis-court': 0.35, 'swimming-pool': 0.4, 'storage-tank': 0.3, 'soccer-ball-field': 0.3, 'small-vehicle': 0.4, 'ship': 0.35, 'plane': 0.35, 'large-vehicle': 0.4, 'helicopter': 0.4, 'harbor': 0.3, 'ground-track-field': 0.4, 'bridge': 0.3, 'basketball-court': 0.4, 'baseball-diamond': 0.3 } for sub_class in range(1, cfgs.CLASS_NUM + 1): index = np.where(label_res == sub_class)[0] if len(index) == 0: continue tmp_boxes_h = box_res[index] tmp_label_h = label_res[index] tmp_score_h = score_res[index] tmp_boxes_h = np.array(tmp_boxes_h) tmp = np.zeros( [tmp_boxes_h.shape[0], tmp_boxes_h.shape[1] + 1]) tmp[:, 0:-1] = tmp_boxes_h tmp[:, -1] = np.array(tmp_score_h) inx = nms.py_cpu_nms( dets=np.array(tmp, np.float32), thresh=h_threshold[LABEl_NAME_MAP[sub_class]], max_output_size=500) box_res_.extend(np.array(tmp_boxes_h)[inx]) score_res_.extend(np.array(tmp_score_h)[inx]) label_res_.extend(np.array(tmp_label_h)[inx]) time_elapsed = timer() - start if save_res: scores = np.array(score_res_) labels = np.array(label_res_) boxes = np.array(box_res_) valid_show = scores > cfgs.SHOW_SCORE_THRSHOLD scores = scores[valid_show] boxes = boxes[valid_show] labels = labels[valid_show] det_detections_h = draw_box_in_img.draw_boxes_with_label_and_scores( np.array(img, np.float32), boxes=np.array(boxes), labels=np.array(labels), scores=np.array(scores), in_graph=False) save_dir = os.path.join(des_folder, cfgs.VERSION) tools.mkdir(save_dir) cv2.imwrite( save_dir + '/' + img_path.split('/')[-1].split('.')[0] + '_h.jpg', det_detections_h) view_bar( '{} cost {}s'.format( img_path.split('/')[-1].split('.')[0], time_elapsed), count + 1, len(file_paths)) else: # eval txt CLASS_DOTA = NAME_LABEL_MAP.keys() # Task2 write_handle_h = {} txt_dir_h = os.path.join('txt_output', cfgs.VERSION + '_h') tools.mkdir(txt_dir_h) for sub_class in CLASS_DOTA: if sub_class == 'back_ground': continue write_handle_h[sub_class] = open( os.path.join(txt_dir_h, 'Task2_%s.txt' % sub_class), 'a+') for i, hbox in enumerate(box_res_): command = '%s %.3f %.1f %.1f %.1f %.1f\n' % ( img_path.split('/')[-1].split('.')[0], score_res_[i], hbox[0], hbox[1], hbox[2], hbox[3]) write_handle_h[LABEl_NAME_MAP[label_res_[i]]].write( command) for sub_class in CLASS_DOTA: if sub_class == 'back_ground': continue write_handle_h[sub_class].close() view_bar( '%s cost %.3fs' % (img_path.split('/')[-1].split('.')[0], time_elapsed), count + 1, len(file_paths)) fw.write('{}\n'.format(img_path)) fw.close() os.remove('./tmp.txt')