def load_network(self, detect_net, detect_model, classify_net,
classify_model):
config = tf.ConfigProto(allow_soft_placement=True)
config.gpu_options.allocator_type = 'BFC'
config.gpu_options.per_process_gpu_memory_fraction = 0.8
# 构建图
detect_graph = tf.Graph()
classify_graph = tf.Graph()
# 为构建的图分别建立会话
detect_sess = tf.Session(graph=detect_graph)
classify_sess = tf.Session(graph=classify_graph)
with detect_graph.as_default():
detect_network = get_network(detect_net)
print('Loading networks {:s}... '.format(detect_net))
saver = tf.train.Saver()
ckpt = tf.train.get_checkpoint_state(detect_model)
print('Restoring from {}...'.format(ckpt.model_checkpoint_path), )
saver.restore(detect_sess, ckpt.model_checkpoint_path)
print(' done.')
with classify_graph.as_default():
classify_network = get_network(classify_net)
print('Loading networks {:s}... '.format(classify_net))
saver = tf.train.Saver()
ckpt = tf.train.get_checkpoint_state(classify_model)
print('Restoring from {}...'.format(ckpt.model_checkpoint_path), )
saver.restore(classify_sess, ckpt.model_checkpoint_path)
print(' done.')
return detect_sess, classify_sess, detect_network, classify_network
def detectorload():
cfg_from_file('TextDetection/ctpn/text.yml')
# init session
config = tf.ConfigProto(allow_soft_placement=True)
sess = tf.Session(config=config)
# load network
net = get_network("VGGnet_test")
# load model
print(('Loading network {:s}... '.format("VGGnet_test")), end=' ')
saver = tf.train.Saver()
try:
ckpt = tf.train.get_checkpoint_state('TextDetection/checkpoints/')
print(
'Restoring from {}...'.format(ckpt.model_checkpoint_path), end=' ')
saver.restore(sess, ckpt.model_checkpoint_path)
print('done')
except:
raise 'Check your pretrained {:s}'.format(ckpt.model_checkpoint_path)
im = 128 * np.ones((300, 300, 3), dtype=np.uint8)
for i in range(2):
_, _ = test_ctpn(sess, net, im)
return net, sess
net, sess = detectorload()
ctpn(sess, net, im_name)
def ctpnSource(im_name):
if os.path.exists("data/results/"):
shutil.rmtree("data/results/") #目录下存在多级子目录，递归删除
os.makedirs("data/results/")
cfg_from_file('ctpn/text.yml')
# init session
config = tf.ConfigProto(allow_soft_placement=True)
sess = tf.Session(config=config)
# load network
net = get_network("VGGnet_test")
# load model
# print(('Loading network {:s}... '.format("VGGnet_test")), end=' ')
saver = tf.train.Saver()
try:
ckpt = tf.train.get_checkpoint_state(cfg.TEST.checkpoints_path)
# print('Restoring from {}...'.format(ckpt.model_checkpoint_path), end=' ')
saver.restore(sess, ckpt.model_checkpoint_path)
# print('done')
except:
raise 'Check your pretrained {:s}'.format(ckpt.model_checkpoint_path)
im = 128 * np.ones((300, 300, 3), dtype=np.uint8)
for i in range(2):
_, _ = test_ctpn(sess, net, im)
# im_names = glob.glob(os.path.join(cfg.DATA_DIR, 'demo', '*.png')) + \
# glob.glob(os.path.join(cfg.DATA_DIR, 'demo', '*.jpg'))
# for im_name in im_names:
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(('Demo for {:s}'.format(im_name)))
ctpn(sess, net, im_name)
def load_tf_model():
# load config file
cfg.TEST.checkpoints_path = '/Users/hjinfeng/Documents/cigna/ctpn-sagemaker/ctpn/checkpoints'
# for device in ['/device:gpu:1','/device:gpu:2']:
# with tf.device(d)
# init session
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=1.0)
config = tf.ConfigProto(allow_soft_placement=True, gpu_options=gpu_options)
sess = tf.Session(config=config)
# load network
net = get_network("VGGnet_test")
# load model
print('Loading network {:s}... '.format("VGGnet_test"))
saver = tf.train.Saver()
try:
print(cfg)
ckpt = tf.train.get_checkpoint_state(cfg.TEST.checkpoints_path)
print(ckpt)
print('Restoring from {}...'.format(ckpt.model_checkpoint_path))
saver.restore(sess, ckpt.model_checkpoint_path)
print('done')
except Exception as e:
print(e)
raise 'Check your pretrained {:s}'.format(ckpt.model_checkpoint_path)
return sess, net
def load_tf_model():
# load config file
cfg.TEST.checkpoints_path = './ctpn/checkpoints'
# init session
#gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.3, allow_growth = True)
#config = tf.ConfigProto(allow_soft_placement=True, gpu_options=gpu_options)
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
tf_config = tf.ConfigProto(allow_soft_placement=True)
tf_config.gpu_options.per_process_gpu_memory_fraction = 1.0
tf_config.gpu_options.allow_growth = True
sess = tf.Session(config=tf_config)
#ktf.set_session(sess)
# load network
#print('/gpu:{}'.format(1))
#with tf.device('/gpu:{}'.format(1)):
net = get_network("VGGnet_test")
# load model
print('Loading network {:s}... '.format("VGGnet_test"))
saver = tf.train.Saver()
try:
ckpt = tf.train.get_checkpoint_state(cfg.TEST.checkpoints_path)
print('Restoring from {}...'.format(ckpt.model_checkpoint_path))
saver.restore(sess, ckpt.model_checkpoint_path)
print('done')
except:
raise 'Check your pretrained {:s}'.format(ckpt.model_checkpoint_path)
return sess, net
def load_tf_model():
# load config file
cfg.TEST.checkpoints_path = './ctpn/checkpoints'
# init session
# gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.33)
gpu_options = tf.GPUOptions()
gpu_options.allow_growth = True
config = tf.ConfigProto(allow_soft_placement=True, gpu_options=gpu_options)
sess = tf.Session(config=config)
# load network
net = get_network("VGGnet_test")
# load model
print('Loading network {:s}... '.format("VGGnet_test"))
saver = tf.train.Saver()
try:
ckpt = tf.train.get_checkpoint_state(cfg.TEST.checkpoints_path)
print('Restoring from {}...'.format(ckpt.model_checkpoint_path))
saver.restore(sess, ckpt.model_checkpoint_path)
print('done')
except:
raise 'Check your pretrained {:s}'.format(ckpt.model_checkpoint_path)
return sess, net
def __init__(self,
init_score_thresh,
init_nms_thresh,
Meta_path=None,
Model_path=None):
cfg.TEST.HAS_RPN = True # Use RPN for proposals
self.score_thresh = init_score_thresh
self.nms_thresh = init_nms_thresh
#args = parse_args()
if Model_path == ' ' or not os.path.exists(Model_path):
print('current path is ' + os.path.abspath(__file__))
raise IOError(('Error: Model not found.\n'))
self.anchors = np.array([[-83., -39., 100., 56.],
[-175., -87., 192., 104.],
[-359., -183., 376., 200.],
[-55., -55., 72., 72.],
[-119., -119., 136., 136.],
[-247., -247., 264., 264.],
[-35., -79., 52., 96.],
[-79., -167., 96., 184.],
[-167., -343., 184., 360.]])
# init session
# load network
#print "loading kekekekekkeeeeeeeeeeeeeeeeeeeeeeeeeeee"
self.graph = tf.Graph()
self.sess = tf.Session(graph=self.graph)
with self.graph.as_default():
self.net = get_network('VGGnet_test')
self.sess.run(tf.global_variables_initializer())
self.saver = tf.train.Saver(
) #write_version = tf.train.SaverDef.V2)
self.saver.restore(self.sess, Model_path)
'''
def train_main(data_dir, model_dir, train_steps, input_yaml):
cfg_from_file(input_yaml)
print('Using config:')
pprint.pprint(cfg)
imdb = get_imdb('voc_2007_trainval')
print('Loaded dataset `{:s}` for training'.format(imdb.name))
roidb = get_training_roidb(imdb)
output_dir = get_output_dir(imdb, None)
log_dir = get_log_dir(imdb)
print('Output will be saved to `{:s}`'.format(output_dir))
print('Logs will be saved to `{:s}`'.format(log_dir))
device_name = '/gpu:0'
print(device_name)
network = get_network('VGGnet_train')
train_net(network,
imdb,
roidb,
output_dir=output_dir,
log_dir=log_dir,
pretrained_model='data/pretrain_model/VGG_imagenet.npy',
max_iters=int(cfg.TRAIN.max_steps),
restore=bool(int(cfg.TRAIN.restore)))
def load_tf_model():
cfg.TEST.HAS_RPN = True # Use RPN for proposals
# init session
config = tf.ConfigProto(allow_soft_placement=True)
net = get_network("VGGnet_test")
# load model
saver = tf.train.Saver()
# sess = tf.Session(config=config)
sess = tf.Session()
ckpt = tf.train.get_checkpoint_state(
# './checkpoints/')
'../ModelSet/Chinese-OCR/ctpn/checkpoints/')
# print("load ckpt=[", ckpt, "]")
reader = tf.train.NewCheckpointReader(ckpt.model_checkpoint_path)
var_to_shape_map = reader.get_variable_to_shape_map()
# for key in var_to_shape_map:
# print("Tensor_name is : ", key)
# print(reader.get_tensor(key))
saver.restore(sess, ckpt.model_checkpoint_path)
# print("load vggnet done")
return sess, saver, net
def get_model(cls):
if cls.sess == None:
cfg_from_file('ctpn/text.yml')
# init session
config = tf.ConfigProto(allow_soft_placement=True)
cls.sess = tf.Session(config=config)
# load network
cls.net = get_network("VGGnet_test")
# load model
print(('Loading network {:s}... '.format("VGGnet_test")), end=' ')
saver = tf.train.Saver()
try:
ckpt = tf.train.get_checkpoint_state(cfg.TEST.checkpoints_path)
print('Restoring from {}...'.format(
ckpt.model_checkpoint_path),
end=' ')
saver.restore(cls.sess, ckpt.model_checkpoint_path)
print('done')
except:
raise 'Check your pretrained {:s}'.format(
ckpt.model_checkpoint_path)
im = 128 * np.ones((300, 300, 3), dtype=np.uint8)
for i in range(2):
_, _ = test_ctpn(cls.sess, cls.net, im)
def __init__(self):
# ctpn model
self.ctpn_model_path = 'trained_models/trained_ctpn_models'
# path to needed detect image
self.detect_data_path = os.path.join('detect_data/detect_images')
# crnn model
self.crnn_model_path = 'trained_models/trained_crnn_models/latestModel.pth'
# recognition results
self.txt_results_path = 'detect_data/txt_results/text_info_results.txt'
# a character dictionary
self.alphabet = alphabets.alphabet
# total class
self.nclass = len(self.alphabet) + 1
# text proposals
self.boxes = []
# tensorflow configures
self.config = tf.ConfigProto(log_device_placement=True)
self.config.gpu_options.allow_growth = True
self.config.gpu_options.per_process_gpu_memory_fraction = 0.75
# tensorflow session
self.sess = tf.Session(config=self.config)
# load network
self.net = get_network("VGGnet_test")
# load model
self.saver = tf.train.Saver()
# crnn network
self.model = crnn.CRNN(32, 1, self.nclass, 256)
def detect(src,
dst,
draw_img=False,
show_area=False,
area_min=-0.0,
area_max=1.1):
#if os.path.exists("data/results/"):
# shutil.rmtree("data/results/")
if not os.path.exists(dst):
os.makedirs(dst)
cfg_from_file('ctpn/text.yml')
# init session
config = tf.ConfigProto(allow_soft_placement=True)
sess = tf.Session(config=config)
# load network
net = get_network("VGGnet_test")
# load model
print(('Loading network {:s}... '.format("VGGnet_test")), end=' ')
saver = tf.train.Saver()
try:
ckpt = tf.train.get_checkpoint_state(cfg.TEST.checkpoints_path)
print('Restoring from {}...'.format(ckpt.model_checkpoint_path),
end=' ')
saver.restore(sess, ckpt.model_checkpoint_path)
print('done')
except:
raise 'Check your pretrained {:s}'.format(ckpt.model_checkpoint_path)
im = 128 * np.ones((300, 300, 3), dtype=np.uint8)
for i in range(2):
_, _ = test_ctpn(sess, net, im)
'''
im_names = glob.glob(os.path.join(cfg.DATA_DIR, 'demo', '*.png')) + \
glob.glob(os.path.join(cfg.DATA_DIR, 'demo', '*.jpg')) + \
glob.glob(os.path.join(cfg.DATA_DIR, 'demo', '*.jpeg'))
'''
im_names = glob.glob(os.path.join(src, '*.png')) \
+ glob.glob(os.path.join(src, '*.PNG')) \
+ glob.glob(os.path.join(src, '*.jpg')) \
+ glob.glob(os.path.join(src, '*.JPG')) \
+ glob.glob(os.path.join(src, '*.jpeg')) \
+ glob.glob(os.path.join(src, '*.JPEG')) \
+ glob.glob(os.path.join(src, '*.bmp')) \
+ glob.glob(os.path.join(src, '*.BMP'))
print("images:{}".format(len(im_names)))
for im_name in im_names:
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(('Demo for {:s}'.format(im_name)))
ctpn(sess,
net,
im_name,
dst,
draw_img=draw_img,
show_area=show_area,
area_min=area_min,
area_max=area_max)
def demo(sess, net, image_name):
# def parse_args():
# """Parse input arguments."""
# parser = argparse.ArgumentParser(description='Faster R-CNN demo')
# parser.add_argument('--gpu', dest='gpu_id', help='GPU device id to use [0]',
# default=0, type=int)
# parser.add_argument('--cpu', dest='cpu_mode',
# help='Use CPU mode (overrides --gpu)',
# action='store_true')
# parser.add_argument('--net', dest='demo_net', help='Network to use [vgg16]',
# default='VGGnet_test')
# # parser.add_argument('--model', dest='model', help='Model path',
# # default='model/VGGnet_fast_rcnn_iter_60000.ckpt')
#
# args = parser.parse_args()
#
# return args
if __name__ == '__main__':
cfg.TEST.HAS_RPN = True # Use RPN for proposals
# args = parse_args()
# if args.model == ' ' or not os.path.exists(args.model):
# print ('current path is ' + os.path.abspath(__file__))
# raise IOError(('Error: Model not found.\n'))
# model = os.path.join(this_dir,'output','default','voc_2007_trainval','VGGnet_fast_rcnn_iter_100000.ckpt')
model= os.path.join(this_dir,'model-6.16.ckpt')
# init session
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
# load network
net = get_network('VGGnet_test')
# load model
# print ('Loading network {:s}... '.format(args.demo_net)),
saver = tf.train.Saver()
saver.restore(sess, model)
# print (' done.')
# model = os.path.join(this_dir, 'model.ckpt')
# saver.restore(sess, model)
# Warmup on a dummy image
im = 128 * np.ones((300, 300, 3), dtype=np.uint8)
for i in range(2):
_, _ = im_detect(sess, net, im)
im_names = glob.glob(os.path.join(cfg.DATA_DIR, 'demo42', '*.*'))
for im_name in im_names:
print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
print 'Demo for {:s}'.format(im_name)
demo(sess, net, im_name)
plt.show()
def load_network(self, network_name, model_root):
net = get_network(network_name)
# load model
print('Loading networks {:s}... '.format(network_name))
saver = tf.train.Saver()
ckpt = tf.train.get_checkpoint_state(model_root)
print('Restoring from {}...'.format(ckpt.model_checkpoint_path), )
saver.restore(self.sess, ckpt.model_checkpoint_path)
print(' done.')
return net
def export():
cfg_from_file(os.path.join(dir_path, 'text.yml'))
config = tf.ConfigProto(allow_soft_placement=True)
sess = tf.Session(config=config)
# with gfile.FastGFile('../data/ctpn.pb', 'rb') as f:
# graph_def = tf.GraphDef()
# graph_def.ParseFromString(f.read())
# sess.graph.as_default()
# tf.import_graph_def(graph_def, name='')
# sess.run(tf.global_variables_initializer())
net = get_network("VGGnet_test")
print(('Loading network {:s}... '.format("VGGnet_test")), end=' ')
saver = tf.train.Saver()
print(saver)
try:
ckpt_path = os.path.abspath(
os.path.join(dir_path, cfg.TEST.checkpoints_path))
print('check_path, ', ckpt_path)
ckpt = tf.train.get_checkpoint_state(ckpt_path)
print('Restoring from {}...'.format(ckpt.model_checkpoint_path),
end=' ')
saver.restore(sess, ckpt.model_checkpoint_path)
print('done')
except:
raise 'Check your pretrained {:s}'.format(ckpt.model_checkpoint_path)
print(' done.')
input_img = sess.graph.get_tensor_by_name('Placeholder:0')
output_cls_prob = sess.graph.get_tensor_by_name('Reshape_2:0')
output_box_pred = sess.graph.get_tensor_by_name(
'rpn_bbox_pred/Reshape_1:0')
builder = tf.saved_model.builder.SavedModelBuilder('./export/1')
imageplaceholder_info = tf.saved_model.utils.build_tensor_info(input_img)
cls_prob_info = tf.saved_model.utils.build_tensor_info(output_cls_prob)
box_pred_info = tf.saved_model.utils.build_tensor_info(output_box_pred)
print('predict_method_name,',
tf.saved_model.signature_constants.PREDICT_METHOD_NAME)
prediction_signature = (
tf.saved_model.signature_def_utils.build_signature_def(
inputs={'image': imageplaceholder_info},
outputs={
'output_cls_prob': cls_prob_info,
'output_box_pred': box_pred_info
},
method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME)
)
init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
builder.add_meta_graph_and_variables(
sess, [tf.saved_model.tag_constants.SERVING],
signature_def_map={'ctpn_recs_predict': prediction_signature},
legacy_init_op=init_op)
builder.save()
def load(self):
logging.info('Creating networks and loading parameters')
cfg_from_file(os.path.join(os.path.dirname(__file__),'text.yml'))
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=self.gpu_fraction)
config = tf.ConfigProto(allow_soft_placement=True,gpu_options=gpu_options)
self.session = tf.Session(config=config)
self.net = get_network("VGGnet_test")
self.textdetector = TextDetector()
saver = tf.train.Saver()
ckpt = tf.train.get_checkpoint_state(self.model_path)
saver.restore(self.session, ckpt.model_checkpoint_path)
def __init__(self, ckpt_path):
graph = tf.Graph()
with graph.as_default():
self.sess = tf.Session()
with self.sess.as_default():
self.net = get_network("VGGnet_test")
self.saver = tf.train.Saver()
self.ckpt = tf.train.get_checkpoint_state(ckpt_path)
self.saver.restore(self.sess, self.ckpt.model_checkpoint_path)
self.scale, self.max_scale = CTPNConfig.SCALE, CTPNConfig.MAX_SCALE
self.textdetector = TextDetector()
def load_tf_model():
cfg.TEST.HAS_RPN = True # Use RPN for proposals
# init session
config = tf.ConfigProto(allow_soft_placement=True)
sess = tf.Session(config=config)
# load network
net = get_network("VGGnet_test")
# load model
saver = tf.train.Saver()
ckpt = tf.train.get_checkpoint_state('ctpn/models/')
saver.restore(sess, ckpt.model_checkpoint_path)
return sess,saver,net
def load_tf_model():
cfg.TEST.HAS_RPN = True # Use RPN for proposals
# init session
config = tf.ConfigProto(allow_soft_placement=True)
sess = tf.Session(config=config)
# load network
net = get_network("VGGnet_test")
# load model
saver = tf.train.Saver()
ckpt = tf.train.get_checkpoint_state('ctpn/models/')
saver.restore(sess, ckpt.model_checkpoint_path)
return sess, saver, net
def detect(self,
img,
ret=dict(),
net='VGGnet_test',
model=os.path.join(local_dir,
'models/VGGnet_fast_rcnn_iter_150000.ckpt')):
"""Detect object classes in an image using pre-computed object proposals."""
self.sess = tf.Session(config=tf.ConfigProto(
allow_soft_placement=True))
# load network
self.net = get_network(net)
#load model
print('Loading network {:s}... '.format(net)),
self.saver = tf.train.Saver()
self.saver.restore(self.sess, model)
print(' done.')
starttime = time.time()
# Load the demo image
im = cv2.imread(img)
print im.shape
im = cv2.resize(im, (int((400.0 / im.shape[0]) * im.shape[1]), 400))
print im.shape
# Detect all object classes and regress object bounds
timer = Timer()
timer.tic()
scores, boxes = im_detect(self.sess, self.net, im)
timer.toc()
print('Detection took {:.3f}s for '
'{:d} object proposals').format(timer.total_time, boxes.shape[0])
# Visualize detections for each class
# im = im[:, :, (2, 1, 0)]
# fig, ax = plt.subplots(figsize=(12, 12))
# ax.imshow(im, aspect='equal')
CONF_THRESH = 0.8
NMS_THRESH = 0.3
for cls_ind, cls in enumerate(CLASSES[1:]):
cls_ind += 1 # because we skipped background
cls_boxes = boxes[:, 4 * cls_ind:4 * (cls_ind + 1)]
cls_scores = scores[:, cls_ind]
dets = np.hstack(
(cls_boxes, cls_scores[:, np.newaxis])).astype(np.float32)
keep = nms(dets, NMS_THRESH)
dets = dets[keep, :]
im = self.vis_detections(im, cls, dets, CONF_THRESH)
endtime = time.time()
ret['result'] = (True, "%.3f" % (endtime - starttime))
ret['drawImg'] = im
print "finish detecting"
return im
def __init__(self, cfg, ckt):
cfg_from_file(cfg)
config = tf.ConfigProto(allow_soft_placement=True)
self.sess = tf.Session(config=config)
self.net = get_network('VGGnet_test')
saver = tf.train.Saver()
try:
ckpt = tf.train.get_checkpoint_state(ckt)
saver.restore(self.sess, ckpt.model_checkpoint_path)
except:
raise 'Check your pretrained {:s}'.format(
ckpt.model_checkpoint_path)
def __init__(self, model_name, demo_net, demo_model, classes):
MLAlgorithm.__init__(self)
cfg.TEST.HAS_RPN = True
g = tf.Graph()
with g.as_default():
# init session
self.sess = tf.Session(
config=tf.ConfigProto(allow_soft_placement=True)
)
self.net = get_network(demo_net, len(classes))
saver = tf.train.Saver()
saver.restore(self.sess, demo_model)
self.classes = classes
print 'Loaded network {:s}'.format(demo_model)
def run_ctpn(img):
cfg_from_file('ctpn/text.yml')
config = tf.ConfigProto(allow_soft_placement=True)
# load network 构建网络模型
net = get_network("VGGnet_test")
# load model
print(('Loading network {:s}... '.format("VGGnet_test")), end=' ')
# 将图像进行resize并返回其缩放大小
img_resized, bbox_scale = resize_im(img,
scale=TextLineCfg.SCALE,
max_scale=TextLineCfg.MAX_SCALE)
# print(scale)
run_list, feed_dict, im_scales = run(net, img_resized)
return config, run_list, feed_dict, img_resized.shape, im_scales, bbox_scale
def load_tf_model():
cfg.TEST.HAS_RPN = True # Use RPN for proposals
# init session
config = tf.ConfigProto(allow_soft_placement=True)
sess = tf.Session(config=config)
# load network
net = get_network("VGGnet_test")
# load model
#print ('Loading network {:s}... '.format("VGGnet_test")),
saver = tf.train.Saver()
saver.restore(
sess,
os.path.join(os.getcwd(), "ctpn/checkpoints/model_final_tf13.ckpt"))
#print (' done.')
return sess, saver, net
def local_result(list_img_path, save_dir, ctpn_path, base_net):
save_dir = os.path.join(cfg.ROOT_DIR, 'data', save_dir)
if os.path.exists(save_dir):
shutil.rmtree(save_dir)
os.makedirs(save_dir)
yml_path = os.path.join(cfg.ROOT_DIR, 'ctpn/text.yml')
cfg_from_file(yml_path)
# init session
config = tf.ConfigProto(allow_soft_placement=True)
sess = tf.Session(config=config)
# load network
training_flag = tf.placeholder(tf.bool)
net = get_network(base_net, training_flag)
# load model
print(('Loading network {:s}... '.format(base_net)), end=' ')
saver = tf.train.Saver()
try:
ckpt = tf.train.get_checkpoint_state(ctpn_path)
print('Restoring from {}...'.format(ckpt.model_checkpoint_path), end=' ')
saver.restore(sess, ckpt.model_checkpoint_path)
print('done')
except:
raise 'Check your pretrained {:s}'.format(ckpt.model_checkpoint_path)
im = 128 * np.ones((300, 300, 3), dtype=np.uint8)
for i in range(2):
_, _ = test_ctpn(sess, training_flag, net, im)
im_names = list_img_path
# im_names = glob.glob(os.path.join(cfg.DATA_DIR, 'demo', '*.png')) + \
# glob.glob(os.path.join(cfg.DATA_DIR, 'demo', '*.jpg'))
xml_dir = '/data/kuaidi01/dataset_detect/VOC2007/Annotations'
for im_name in im_names:
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(('Demo for {:s}'.format(im_name)))
if im_name[-7:]=='164.jpg':
continue
xml_path = os.path.join(xml_dir, os.path.basename(im_name)[:-4]+'.xml')
ctpn(sess,training_flag, net, im_name, save_dir, xml_path)
def load_tf_model():
cfg.TEST.HAS_RPN = True # Use RPN for proposals
# init session
config = tf.ConfigProto(allow_soft_placement=True)
net = get_network("VGGnet_test")
# load model
saver = tf.train.Saver()
# sess = tf.Session(config=config)
sess = tf.Session()
ckpt_path = '/home/zihao/Downloads/project/ctpn/ctpn/checkpoints'
ckpt = tf.train.get_checkpoint_state(ckpt_path)
reader = tf.train.NewCheckpointReader(ckpt.model_checkpoint_path)
var_to_shape_map = reader.get_variable_to_shape_map()
for key in var_to_shape_map:
print("Tensor_name is : ", key)
# print(reader.get_tensor(key))
saver.restore(sess, ckpt.model_checkpoint_path)
print("load vggnet done")
return sess, saver, net
def load_tf_model():
cfg.TEST.HAS_RPN = True # Use RPN for proposals
# init session
config = tf.ConfigProto(allow_soft_placement=True)
net = get_network("VGGnet_test")
# load model
saver = tf.train.Saver()
# sess = tf.Session(config=config)
sess = tf.Session()
ckpt = tf.train.get_checkpoint_state(
'/Users/xiaofeng/Code/Github/dataset/CHINESE_OCR/ctpn/checkpoints/')
reader = tf.train.NewCheckpointReader(ckpt.model_checkpoint_path)
var_to_shape_map = reader.get_variable_to_shape_map()
for key in var_to_shape_map:
print("Tensor_name is : ", key)
# print(reader.get_tensor(key))
saver.restore(sess, ckpt.model_checkpoint_path)
print("load vggnet done")
return sess, saver, net
def load_tf_model():
cfg.TEST.HAS_RPN = True # Use RPN for proposals
# init session
config = tf.ConfigProto(allow_soft_placement=True)
net = get_network("VGGnet_test")
# load model
saver = tf.train.Saver()
# sess = tf.Session(config=config)
sess = tf.Session()
data_dir = os.path.expanduser('~/data')
ckpt = tf.train.get_checkpoint_state(data_dir + os.sep + 'checkpoints')
reader = tf.train.NewCheckpointReader(ckpt.model_checkpoint_path)
var_to_shape_map = reader.get_variable_to_shape_map()
for key in var_to_shape_map:
print("Tensor_name is : ", key)
# print(reader.get_tensor(key))
saver.restore(sess, ckpt.model_checkpoint_path)
print("load vggnet done")
return sess, saver, net
def local_result(img_path, save_dir, ctpn_path, base_net):
save_dir = os.path.join(cfg.ROOT_DIR, 'data', save_dir)
if os.path.exists(save_dir):
shutil.rmtree(save_dir)
os.makedirs(save_dir)
yml_path = os.path.join(cfg.ROOT_DIR, 'ctpn/text.yml')
cfg_from_file(yml_path)
# init session
config = tf.ConfigProto(allow_soft_placement=True)
sess = tf.Session(config=config)
# load network
training_flag = tf.placeholder(tf.bool)
net = get_network(base_net, training_flag)
# load model
print(('Loading network {:s}... '.format(base_net)), end=' ')
saver = tf.train.Saver()
ckpt = tf.train.get_checkpoint_state(ctpn_path)
print('Restoring from {}...'.format(ckpt.model_checkpoint_path), end=' ')
saver.restore(sess, ckpt.model_checkpoint_path)
print('done')
try:
ckpt = tf.train.get_checkpoint_state(ctpn_path)
print('Restoring from {}...'.format(ckpt.model_checkpoint_path),
end=' ')
saver.restore(sess, ckpt.model_checkpoint_path)
print('done')
except Exception:
print('Check your pretrained {:s}'.format(ckpt.model_checkpoint_path))
im = 128 * np.ones((300, 300, 3), dtype=np.uint8)
for i in range(2):
_, _ = test_ctpn(sess, training_flag, net, im)
for im_name in img_path:
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(('Demo for {:s}'.format(im_name)))
# if os.path.basename(im_name)!='30_5357176972_20180629182525_20180629202528_163.jpg':
# continue
ctpn(sess, training_flag, net, im_name, save_dir)
def local_result(save_dir):
save_dir = os.path.join(cfg.ROOT_DIR, 'data', save_dir)
if os.path.exists(save_dir):
shutil.rmtree(save_dir)
os.makedirs(save_dir)
yml_path = os.path.join(cfg.ROOT_DIR, 'ctpn/text.yml')
cfg_from_file(yml_path)
# init session
config = tf.ConfigProto(allow_soft_placement=True)
sess = tf.Session(config=config)
# load network
net = get_network("VGGnet_test")
# load model
print(('Loading network {:s}... '.format("VGGnet_test")), end=' ')
saver = tf.train.Saver()
_ = os.path.join(cfg.ROOT_DIR,
'output_old_auto/ctpn_end2end/voc_2007_trainval')
try:
ckpt = tf.train.get_checkpoint_state(_)
print('Restoring from {}...'.format(ckpt.model_checkpoint_path),
end=' ')
saver.restore(sess, ckpt.model_checkpoint_path)
print('done')
except:
raise 'Check your pretrained {:s}'.format(ckpt.model_checkpoint_path)
im = 128 * np.ones((300, 300, 3), dtype=np.uint8)
for i in range(2):
_, _ = test_ctpn(sess, net, im)
im_names = glob.glob(os.path.join(cfg.DATA_DIR, 'gt_set', '*.jpg'))
# im_names = glob.glob(os.path.join(cfg.DATA_DIR, 'demo', '*.png')) + \
# glob.glob(os.path.join(cfg.DATA_DIR, 'demo', '*.jpg'))
for im_name in im_names:
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(('Demo for {:s}'.format(im_name)))
ctpn(sess, net, im_name, save_dir)
def detect_area(pics_file,
dst_file,
draw_img=False,
show_area=False,
area_min=-0.0,
area_max=1.1):
cfg_from_file('ctpn/text.yml')
# init session
config = tf.ConfigProto(allow_soft_placement=True)
sess = tf.Session(config=config)
# load network
net = get_network("VGGnet_test")
# load model
print(('Loading network {:s}... '.format("VGGnet_test")), end=' ')
saver = tf.train.Saver()
try:
ckpt = tf.train.get_checkpoint_state(cfg.TEST.checkpoints_path)
print('Restoring from {}...'.format(ckpt.model_checkpoint_path),
end=' ')
saver.restore(sess, ckpt.model_checkpoint_path)
print('done')
except:
raise 'Check your pretrained {:s}'.format(ckpt.model_checkpoint_path)
im = 128 * np.ones((300, 300, 3), dtype=np.uint8)
for i in range(2):
_, _ = test_ctpn(sess, net, im)
dst_dir = os.path.split(dst_file)[0]
pic_area = pd.DataFrame(columns=['pic', 'area'])
pics = joblib.load(pics_file)
print("images:{}".format(len(pics)))
for im_name in pics:
area = ctpn_area(sess, net, im_name, dst_dir)
pic = os.path.split(im_name)[1]
pic_area.loc[len(pic_area)] = {'pic': pic, 'area': area}
pic_area.to_csv(dst_file, index=None, encoding='UTF8')
from lib.fast_rcnn.train import get_training_roidb, train_net
from lib.fast_rcnn.config import cfg_from_file, get_output_dir, get_log_dir
from lib.datasets.factory import get_imdb
from lib.networks.factory import get_network
from lib.fast_rcnn.config import cfg
if __name__ == '__main__':
cfg_from_file('ctpn/text.yml')
print('Using config:')
pprint.pprint(cfg)
imdb = get_imdb('voc_2007_trainval')
print('Loaded dataset `{:s}` for training'.format(imdb.name))
roidb = get_training_roidb(imdb)
output_dir = get_output_dir(imdb, None)
log_dir = get_log_dir(imdb)
print('Output will be saved to `{:s}`'.format(output_dir))
print('Logs will be saved to `{:s}`'.format(log_dir))
device_name = '/gpu:0'
print(device_name)
network = get_network('VGGnet_train')
train_net(network, imdb, roidb,
output_dir=output_dir,
log_dir=log_dir,
pretrained_model='data/pretrain_model/VGG_imagenet.npy',
max_iters=int(cfg.TRAIN.max_steps),
restore=bool(int(cfg.TRAIN.restore)))
'{:d} object proposals').format(timer.total_time, boxes.shape[0]))
if __name__ == '__main__':
if os.path.exists("data/results/"):
shutil.rmtree("data/results/")
os.makedirs("data/results/")
cfg_from_file('ctpn/text.yml')
# init session
config = tf.ConfigProto(allow_soft_placement=True)
sess = tf.Session(config=config)
# load network
net = get_network("VGGnet_test")
# load model
print(('Loading network {:s}... '.format("VGGnet_test")), end=' ')
saver = tf.train.Saver()
try:
ckpt = tf.train.get_checkpoint_state(cfg.TEST.checkpoints_path)
print('Restoring from {}...'.format(ckpt.model_checkpoint_path), end=' ')
saver.restore(sess, ckpt.model_checkpoint_path)
print('done')
except:
raise 'Check your pretrained {:s}'.format(ckpt.model_checkpoint_path)
im = 128 * np.ones((300, 300, 3), dtype=np.uint8)
for i in range(2):
_, _ = test_ctpn(sess, net, im)
return args
if __name__ == '__main__':
cfg.TEST.HAS_RPN = True # Use RPN for proposals
args = parse_args()
if args.model == ' ' or not os.path.exists(args.model):
print ('current path is ' + os.path.abspath(__file__))
raise IOError(('Error: Model not found.\n'))
# init session
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
# load network
net = get_network(args.demo_net)
# load model
print ('Loading network {:s}... '.format(args.demo_net)),
saver = tf.train.Saver()
saver.restore(sess, args.model)
print (' done.')
# Warmup on a dummy image
im = 128 * np.ones((300, 300, 3), dtype=np.uint8)
for i in xrange(2):
_, _ = im_detect(sess, net, im)
im_names = glob.glob(os.path.join(cfg.DATA_DIR, 'demo', '*.png')) + \
glob.glob(os.path.join(cfg.DATA_DIR, 'demo', '*.jpg'))
for im_name in im_names:
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
print('Using config:')
pprint.pprint(cfg)
if not args.randomize:
# fix the random seeds (numpy and caffe) for reproducibility
np.random.seed(cfg.RNG_SEED)
imdb = get_imdb(args.imdb_name)
print 'Loaded dataset `{:s}` for training'.format(imdb.name)
roidb = get_training_roidb(imdb)
output_dir = get_output_dir(imdb, None)
log_dir = get_log_dir(imdb)
print 'Output will be saved to `{:s}`'.format(output_dir)
print 'Logs will be saved to `{:s}`'.format(log_dir)
device_name = '/gpu:{:d}'.format(args.gpu_id)
print device_name
network = get_network(args.network_name)
print 'Use network `{:s}` in training'.format(args.network_name)
train_net(network, imdb, roidb,
output_dir=output_dir,
log_dir=log_dir,
pretrained_model=args.pretrained_model,
max_iters=args.max_iters,
restore=bool(int(args.restore)))
def export_saved_model(version, path, sess=None):
net = get_network("VGGnet_test")
im = 128 * np.ones((300, 300, 3), dtype=np.uint8)
tf.app.flags.DEFINE_integer('version', version, 'version number of the model.')
tf.app.flags.DEFINE_string('work_dir', path, 'your older model directory.')
tf.app.flags.DEFINE_string('model_dir', '/tmp/model_name', 'saved model directory')
FLAGS = tf.app.flags.FLAGS
# you can give the session and export your model immediately after training
if not sess:
saver = tf.train.import_meta_graph(os.path.join(path, 'xxx.ckpt.meta'))
saver.restore(sess, tf.train.latest_checkpoint(path))
export_path = os.path.join(
tf.compat.as_bytes(FLAGS.model_dir),
tf.compat.as_bytes(str(FLAGS.version)))
builder = tf.saved_model.builder.SavedModelBuilder(export_path)
# define the signature def map here
# ...
os.environ["CUDA_VISIBLE_DEVICES"] = "7"
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.2
graph = tf.Graph().as_default()
# with tf.Graph().as_default() as graph:
# sess = tf.Session()
# MODEL_FILE = 'model/tensorflow_inception_graph.pb'
# BOTTLENECK_TENSOR_NAME = 'pool_3/_reshape:0' # tensor name for the bottleneck of inception-v3 model
# # JPEG_DATA_TENSOR_NAME = 'DecodeJpeg/contents:0'
# JPEG_DATA_TENSOR_NAME = 'DecodeJpeg:0' # image tensor
# f = tf.gfile.FastGFile(MODEL_FILE, 'rb')
# # import initial pb model
# graph_def = tf.GraphDef()
# graph_def.ParseFromString(f.read())
# bottleneck_tensor, jpeg_data_tensor = tf.import_graph_def(graph_def,return_elements=[BOTTLENECK_TENSOR_NAME, JPEG_DATA_TENSOR_NAME])
prediction_signature = predict_signature_def(inputs={'image': graph.get_operation_by_name('BottleneckInputPlaceholder').outputs[0]},
outputs={'scores': graph.get_operation_by_name('evaluation/ArgMax').outputs[0]})
# legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
builder.add_meta_graph_and_variables(
sess=sess,
tags=[tf.saved_model.tag_constants.SERVING],
signature_def_map={
'predict':
prediction_signature
},
)
builder.save()
print('Export SavedModel!')
