def __init__(
self,
model_path='weights/MaskRCNN-R50C41x-COCO_finetune-docrop_and_rotate_24500.pb',
canvas_size=512,
debug=False):
if not tf.test.is_gpu_available():
from tensorflow.python.framework import test_util
assert get_tf_version_tuple() >= (1, 7) and test_util.IsMklEnabled(), \
"Inference requires either GPU support or MKL support!"
self.canvas_size = canvas_size
self.debug = debug
self.id_to_class_name = {
1: 'page',
2: 'profile_image',
3: 'van_tay',
4: 'passport_code'
}
self.resizer = CustomResize(self.canvas_size, self.canvas_size)
print('Loading model at', model_path)
self.graph = load_graph(model_path)
self.input_tensor = self.graph.get_tensor_by_name('import/image:0')
self.output_node_name = [
'output/boxes', 'output/scores', 'output/labels', 'output/masks'
]
self.outputs_tensor = [
self.graph.get_tensor_by_name('import/{}:0'.format(each_node))
for each_node in self.output_node_name
]
self.config = tf.compat.v1.ConfigProto()
# self.config.gpu_options.allow_growth = True
self.config.gpu_options.per_process_gpu_memory_fraction = 0.1
self.sess = tf.compat.v1.Session(config=self.config, graph=self.graph)
self.predict_crop(np.zeros((200, 200, 3), dtype=np.uint8))
print('Loaded model!')
def _setup_graph(self):
num_gpu = cfg.TRAIN.NUM_GPUS
if cfg.TRAINER == 'replicated':
# TF bug in version 1.11, 1.12: https://github.com/tensorflow/tensorflow/issues/22750
buggy_tf = get_tf_version_tuple() in [(1, 11), (1, 12)]
# Use two predictor threads per GPU to get better throughput
self.num_predictor = num_gpu if buggy_tf else num_gpu * 2
self.predictors = [
self._build_predictor(k % num_gpu)
for k in range(self.num_predictor)
]
self.dataflows = [
get_eval_dataflow(self._eval_dataset,
shard=k,
num_shards=self.num_predictor)
for k in range(self.num_predictor)
]
else:
# Only eval on the first machine.
# Alternatively, can eval on all ranks and use allgather, but allgather sometimes hangs
self._horovod_run_eval = hvd.rank() == hvd.local_rank()
if self._horovod_run_eval:
self.predictor = self._build_predictor(0)
self.dataflow = get_eval_dataflow(self._eval_dataset,
shard=hvd.local_rank(),
num_shards=hvd.local_size())
self.barrier = hvd.allreduce(tf.random_normal(shape=[1]))
parser.add_argument(
'--config',
help="A list of KEY=VALUE to overwrite those defined in config.py",
nargs='+')
parser.add_argument('--compact', help='Save a model to .pb')
parser.add_argument('--serving', help='Save a model to serving file')
args = parser.parse_args()
if args.config:
cfg.update_args(args.config)
register_coco(cfg.DATA.BASEDIR) # add COCO datasets to the registry
MODEL = ResNetFPNModel() if cfg.MODE_FPN else ResNetC4Model()
if not tf.test.is_gpu_available():
from tensorflow.python.framework import test_util
assert get_tf_version_tuple() >= (1, 7) and test_util.IsMklEnabled(), \
"Inference requires either GPU support or MKL support!"
assert args.load
finalize_configs(is_training=False)
if args.predict or args.visualize:
cfg.TEST.RESULT_SCORE_THRESH = cfg.TEST.RESULT_SCORE_THRESH_VIS
if args.visualize:
do_visualize(MODEL, args.load)
else:
predcfg = PredictConfig(
model=MODEL,
session_init=get_model_loader(args.load),
input_names=MODEL.get_inference_tensor_names()[0],
output_names=MODEL.get_inference_tensor_names()[1])
gradients = tf.sqrt(tf.reduce_sum(tf.square(gradients), [1, 2, 3]))
gradients_rms = symbolic_functions.rms(gradients, 'gradient_rms')
gradient_penalty = tf.reduce_mean(tf.square(gradients - 1), name='gradient_penalty')
add_moving_summary(self.d_loss, self.g_loss, gradient_penalty, gradients_rms)
self.d_loss = tf.add(self.d_loss, 10 * gradient_penalty)
self.collect_variables()
def optimizer(self):
opt = tf.train.AdamOptimizer(1e-4, beta1=0.5, beta2=0.9)
return opt
if __name__ == '__main__':
assert get_tf_version_tuple() >= (1, 4)
args = DCGAN.get_args(default_batch=64, default_z_dim=128)
M = Model(shape=args.final_size, batch=args.batch, z_dim=args.z_dim)
if args.sample:
DCGAN.sample(M, args.load)
else:
logger.auto_set_dir()
SeparateGANTrainer(
QueueInput(DCGAN.get_data()),
M, g_period=6).train_with_defaults(
callbacks=[ModelSaver()],
steps_per_epoch=300,
max_epoch=200,
session_init=SaverRestore(args.load) if args.load else None
)
im = im.astype(np.float32)[:, :, ::-1]
saliency_images = predictor(im)[0]
abs_saliency = np.abs(saliency_images).max(axis=-1)
pos_saliency = np.maximum(0, saliency_images)
neg_saliency = np.maximum(0, -saliency_images)
pos_saliency -= pos_saliency.min()
pos_saliency /= pos_saliency.max()
cv2.imwrite('pos.jpg', pos_saliency * 255)
neg_saliency -= neg_saliency.min()
neg_saliency /= neg_saliency.max()
cv2.imwrite('neg.jpg', neg_saliency * 255)
abs_saliency = viz.intensity_to_rgb(abs_saliency,
normalize=True)[:, :, ::-1] # bgr
cv2.imwrite("abs-saliency.jpg", abs_saliency)
rsl = im * 0.2 + abs_saliency * 0.8
cv2.imwrite("blended.jpg", rsl)
if __name__ == '__main__':
if len(sys.argv) != 2:
tp.logger.error("Usage: {} image.jpg".format(sys.argv[0]))
sys.exit(1)
assert get_tf_version_tuple() >= (1, 7), "requires TF >= 1.7"
run("resnet_v1_50.ckpt", sys.argv[1])
gradients_rms = tf.sqrt(tf.reduce_mean(tf.square(gradients)),
name='gradient_rms')
gradient_penalty = tf.reduce_mean(tf.square(gradients - 1),
name='gradient_penalty')
add_moving_summary(self.d_loss, self.g_loss, gradient_penalty,
gradients_rms)
self.d_loss = tf.add(self.d_loss, 10 * gradient_penalty)
self.collect_variables()
def optimizer(self):
opt = tf.train.AdamOptimizer(1e-4, beta1=0.5, beta2=0.9)
return opt
if __name__ == '__main__':
assert get_tf_version_tuple() >= (1, 4)
args = DCGAN.get_args(default_batch=64, default_z_dim=128)
M = Model(shape=args.final_size, batch=args.batch, z_dim=args.z_dim)
if args.sample:
DCGAN.sample(M, args.load)
else:
logger.auto_set_dir()
SeparateGANTrainer(QueueInput(
DCGAN.get_data()), M, g_period=5).train_with_defaults(
callbacks=[ModelSaver()],
steps_per_epoch=300,
max_epoch=200,
session_init=SaverRestore(args.load) if args.load else None)
