def inference(_):
# Creating the graphs
evalGraph = tf.Graph()
polyGraph = tf.Graph()
# Evaluator Network
tf.logging.info("Building EvalNet...")
with evalGraph.as_default():
with tf.variable_scope("discriminator_network"):
evaluator = EvalNet(_BATCH_SIZE)
evaluator.build_graph()
saver = tf.train.Saver()
# Start session
evalSess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True),
graph=evalGraph)
saver.restore(evalSess, FLAGS.EvalNet_checkpoint)
# PolygonRNN++
tf.logging.info("Building PolygonRNN++ ...")
model = PolygonModel(FLAGS.PolyRNN_metagraph, polyGraph)
model.register_eval_fn(lambda input_: evaluator.do_test(evalSess, input_))
polySess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True),
graph=polyGraph)
model.saver.restore(polySess, FLAGS.PolyRNN_checkpoint)
if FLAGS.Use_ggnn:
ggnnGraph = tf.Graph()
tf.logging.info("Building GGNN ...")
ggnnModel = GGNNPolygonModel(FLAGS.GGNN_metagraph, ggnnGraph)
ggnnSess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True),
graph=ggnnGraph)
ggnnModel.saver.restore(ggnnSess, FLAGS.GGNN_checkpoint)
tf.logging.info("Testing...")
if not os.path.isdir(FLAGS.OutputFolder):
tf.gfile.MakeDirs(FLAGS.OutputFolder)
crops_path = glob.glob(os.path.join(FLAGS.InputFolder, '*.png'))
for crop_path in tqdm.tqdm(crops_path):
image_np = io.imread(crop_path)
image_np = np.expand_dims(image_np, axis=0)
preds = [
model.do_test(polySess, image_np, top_k)
for top_k in range(_FIRST_TOP_K)
]
# sort predictions based on the eval score and pick the best
preds = sorted(preds, key=lambda x: x['scores'][0], reverse=True)[0]
if FLAGS.Use_ggnn:
polys = np.copy(preds['polys'][0])
feature_indexs, poly, mask = utils.preprocess_ggnn_input(polys)
preds_gnn = ggnnModel.do_test(ggnnSess, image_np, feature_indexs,
poly, mask)
output = {
'polys': preds['polys'],
'polys_ggnn': preds_gnn['polys_ggnn']
}
else:
output = {'polys': preds['polys']}
# dumping to json files
save_to_json(crop_path, output)
def inference(_):
# Creating the graphs
evalGraph = tf.Graph()
polyGraph = tf.Graph()
# Evaluator Network
tf.logging.info("Building EvalNet...")
with evalGraph.as_default():
with tf.variable_scope("discriminator_network"):
evaluator = EvalNet(_BATCH_SIZE)
evaluator.build_graph()
saver = tf.train.Saver()
# Start session
evalSess = tf.Session(config=tf.ConfigProto(
allow_soft_placement=True
), graph=evalGraph)
saver.restore(evalSess, FLAGS.EvalNet_checkpoint)
# PolygonRNN++
tf.logging.info("Building PolygonRNN++ ...")
model = PolygonModel(FLAGS.PolyRNN_metagraph, polyGraph)
model.register_eval_fn(lambda input_: evaluator.do_test(evalSess, input_))
polySess = tf.Session(config=tf.ConfigProto(
allow_soft_placement=True
), graph=polyGraph)
model.saver.restore(polySess, FLAGS.PolyRNN_checkpoint)
if FLAGS.Use_ggnn:
ggnnGraph = tf.Graph()
tf.logging.info("Building GGNN ...")
ggnnModel = GGNNPolygonModel(FLAGS.GGNN_metagraph, ggnnGraph)
ggnnSess = tf.Session(config=tf.ConfigProto(
allow_soft_placement=True
), graph=ggnnGraph)
ggnnModel.saver.restore(ggnnSess, FLAGS.GGNN_checkpoint)
tf.logging.info("Testing...")
if not os.path.isdir(FLAGS.OutputFolder):
tf.gfile.MakeDirs(FLAGS.OutputFolder)
crops_path = glob.glob(os.path.join(FLAGS.InputFolder, '*.png'))
for crop_path in tqdm.tqdm(crops_path):
image_np = io.imread(crop_path)
image_np = np.expand_dims(image_np, axis=0)
preds = [model.do_test(polySess, image_np, top_k) for top_k in range(_FIRST_TOP_K)]
# sort predictions based on the eval score and pick the best
preds = sorted(preds, key=lambda x: x['scores'][0], reverse=True)[0]
if FLAGS.Use_ggnn:
polys = np.copy(preds['polys'][0])
feature_indexs, poly, mask = utils.preprocess_ggnn_input(polys)
preds_gnn = ggnnModel.do_test(ggnnSess, image_np, feature_indexs, poly, mask)
output = {'polys': preds['polys'], 'polys_ggnn': preds_gnn['polys_ggnn']}
else:
output = {'polys': preds['polys']}
# dumping to json files
save_to_json(crop_path, output)
class PolygonRefiner:
def __init__(self, infer=False, show_ggnn=False):
self.run_inference = infer
self.show_ggnn = show_ggnn
self.fig, axes = plt.subplots(1,
2 if show_ggnn else 1,
num=0,
figsize=(12, 6))
self.axes = np.array(axes).flatten()
self.last_image = None
self.last_image_name = None
self.index = 0
if not self.run_inference:
return
# Model setup
self.evalGraph = tf.Graph()
self.polyGraph = tf.Graph()
# Evaluator Network
tf.logging.info("Building EvalNet...")
with self.evalGraph.as_default():
with tf.variable_scope("discriminator_network"):
self.evaluator = EvalNet(_BATCH_SIZE)
self.evaluator.build_graph()
self.saver = tf.train.Saver()
# Start session
self.evalSess = tf.Session(
config=tf.ConfigProto(allow_soft_placement=True),
graph=self.evalGraph)
self.saver.restore(self.evalSess, FLAGS.EvalNet_checkpoint)
# PolygonRNN++
tf.logging.info("Building PolygonRNN++ ...")
self.model = PolygonModel(FLAGS.PolyRNN_metagraph, self.polyGraph)
self.model.register_eval_fn(
lambda input_: self.evaluator.do_test(self.evalSess, input_))
self.polySess = tf.Session(
config=tf.ConfigProto(allow_soft_placement=True),
graph=self.polyGraph)
self.model.saver.restore(self.polySess, FLAGS.PolyRNN_checkpoint)
if FLAGS.Use_ggnn:
self.ggnnGraph = tf.Graph()
tf.logging.info("Building GGNN ...")
self.ggnnModel = GGNNPolygonModel(FLAGS.GGNN_metagraph,
self.ggnnGraph)
self.ggnnSess = tf.Session(
config=tf.ConfigProto(allow_soft_placement=True),
graph=self.ggnnGraph)
self.ggnnModel.saver.restore(self.ggnnSess, FLAGS.GGNN_checkpoint)
def vis(self, pred_path):
vis_single(pred_path, self.fig, self.axes, self.show_ggnn)
def infer(self, image_np, crop_path, output_folder):
# TODO see if we can get some batch parellism
image_np = np.expand_dims(image_np, axis=0)
preds = [
self.model.do_test(self.polySess, image_np, top_k)
for top_k in range(_FIRST_TOP_K)
]
# sort predictions based on the eval score and pick the best
preds = sorted(preds, key=lambda x: x["scores"][0], reverse=True)[0]
if FLAGS.Use_ggnn:
polys = np.copy(preds["polys"][0])
feature_indexs, poly, mask = utils.preprocess_ggnn_input(polys)
preds_gnn = self.ggnnModel.do_test(self.ggnnSess, image_np,
feature_indexs, poly, mask)
output = {
"polys": preds["polys"],
"polys_ggnn": preds_gnn["polys_ggnn"]
}
else:
output = {"polys": preds["polys"]}
# dumping to json files
json_name = save_to_json(output_folder, crop_path, output)
self.vis(json_name)
@ex.capture
def refine(self, image_file, corners, chip_dir, output_dir, _log):
if self.last_image_name == image_file:
image_np = self.last_image
else:
image_np = io.imread(image_file)
self.last_image_name = image_file
self.last_image = image_np.copy()
# Creating the graphs
lx, ty, rx, by = rect_to_box(corners)
image_np = image_np[ty:by, lx:rx]
if image_np.size == 0:
return
image_np = transform.resize(image_np, (224, 224))
output_file = os.path.join(
chip_dir,
"{}_{:06d}.png".format(
os.path.basename(image_file.replace(".", "_")), self.index),
)
# TODO Consider saving to a consistent temp file
_log.info(f"output_file {output_file}")
io.imsave(output_file, image_np)
if self.run_inference:
self.infer(image_np, output_file, output_dir)
self.index += 1
def process_file(self, annotation_file, image_dir, chip_dir, output_dir):
ub.ensuredir(chip_dir, mode=0o0777, recreate=True)
ub.ensuredir(output_dir, mode=0o0777, recreate=True)
df = pd.read_csv(annotation_file, names=range(16), skiprows=2)
corners = df.iloc[:, 3:7]
filenames = df.iloc[:, 1]
for (corner, filename) in tqdm.tqdm(zip(corners.iterrows(),
filenames)):
self.refine(os.path.join(image_dir, filename), corner[1], chip_dir,
output_dir)
def segment_image(path):
with evalGraph.as_default():
with tf.variable_scope("discriminator_network"):
evaluator = EvalNet(_BATCH_SIZE)
evaluator.build_graph()
saver = tf.train.Saver()
# Start session
evalSess = tf.Session(config=tf.ConfigProto(
allow_soft_placement=True
), graph=evalGraph)
saver.restore(evalSess, EvalNet_checkpoint)
#Initializing and restoring PolyRNN++
model = PolygonModel(PolyRNN_metagraph, polyGraph)
model.register_eval_fn(lambda input_: evaluator.do_test(evalSess, input_))
polySess = tf.Session(config=tf.ConfigProto(
allow_soft_placement=True
), graph=polyGraph)
model.saver.restore(polySess, PolyRNN_checkpoint)
#Initializing and restoring GGNN
ggnnGraph = tf.Graph()
ggnnModel = GGNNPolygonModel(GGNN_metagraph, ggnnGraph)
ggnnSess = tf.Session(config=tf.ConfigProto(
allow_soft_placement=True
), graph=ggnnGraph)
ggnnModel.saver.restore(ggnnSess,GGNN_checkpoint)
#INPUT IMG CROP (224x224x3) -> object should be centered
crop_path= path
sys.path.append("/usr/src/lego_classification/assembly_segmentation/python_files")
#Testing
image_np = io.imread(crop_path)
image_np = np.expand_dims(image_np, axis=0)
preds = [model.do_test(polySess, image_np, top_k) for top_k in range(_FIRST_TOP_K)]
# sort predictions based on the eval score to pick the best.
preds = sorted(preds, key=lambda x: x['scores'][0], reverse=True)
print(preds[0]['polys'][0])
# #Visualizing TOP_K and scores
# %matplotlib inline
# import matplotlib.pyplot as plt
# fig, axes = plt.subplots(2,3)
# axes=np.array(axes).flatten()
# [vis_polys(axes[i], image_np[0], np.array(pred['polys'][0]), title='score=%.2f' % pred['scores'][0]) for i,pred in enumerate(preds)]
#Let's run GGNN now on the bestPoly
bestPoly = preds[0]['polys'][0]
feature_indexs, poly, mask = utils.preprocess_ggnn_input(bestPoly)
preds_gnn = ggnnModel.do_test(ggnnSess, image_np, feature_indexs, poly, mask)
refinedPoly=preds_gnn['polys_ggnn']
print ("---------")
print(refinedPoly[0])
return refinedPoly[0]
#Visualize the final prediction
# fig, ax = plt.subplots(1,1)
# vis_polys(ax,image_np[0],refinedPoly[0], title='PolygonRNN++')