@ -1,67 +1,56 @@
#! /usr/bin/env python3
__version__ = ' 1.0 '
"""
Tool to load model and binarize a given image .
"""
import argparse
import sys
import os
from os import listdir , environ , devnull
from os . path import join
from warnings import catch_warnings , simplefilter
import numpy as np
import warnings
from PIL import Image
import cv2
environ [ ' TF_CPP_MIN_LOG_LEVEL ' ] = ' 3 '
stderr = sys . stderr
sys . stderr = open ( devnull , ' w ' )
from keras . models import load_model
sys . stderr = stderr
import tensorflow as tf
def resize_image ( img_in , input_height , input_width ) :
return cv2 . resize ( img_in , ( input_width , input_height ) , interpolation = cv2 . INTER_NEAREST )
class SbbBinarizer :
with warnings . catch_warnings ( ) :
warnings . simplefilter ( " ignore " )
__doc__ = \
"""
Tool to load model and binarize a given image .
"""
class sbb_binarize :
def __init__ ( self , image , model , patches = ' false ' , save = None ) :
self . image = image
self . patches = patches
self . save = save
self . model_dir = model
def resize_image ( self , img_in , input_height , input_width ) :
return cv2 . resize ( img_in , ( input_width , input_height ) , interpolation = cv2 . INTER_NEAREST )
def __init__ ( self , model_dir ) :
self . model_dir = model_dir
def start_new_session _and_model ( self ) :
def start_new_session ( self ) :
config = tf . ConfigProto ( )
config . gpu_options . allow_growth = True
self . session = tf . Session ( config = config ) # tf.InteractiveSession()
def load_model ( self , model_name ) :
self . model = load_model ( self . model_dir + ' / ' + model_name , compile = False )
self . img_height = self . model . layers [ len ( self . model . layers ) - 1 ] . output_shape [ 1 ]
self . img_width = self . model . layers [ len ( self . model . layers ) - 1 ] . output_shape [ 2 ]
self . n_classes = self . model . layers [ len ( self . model . layers ) - 1 ] . output_shape [ 3 ]
def end_session ( self ) :
self . session . close ( )
del self . session
def load_model ( self , model_name ) :
model = load_model ( join ( self . model_dir , model_name ) , compile = False )
model_height = model . layers [ len ( model . layers ) - 1 ] . output_shape [ 1 ]
model_width = model . layers [ len ( model . layers ) - 1 ] . output_shape [ 2 ]
n_classes = model . layers [ len ( model . layers ) - 1 ] . output_shape [ 3 ]
return model , model_height , model_width , n_classes
del self . model
del self . session
def predict ( self , model_name ) :
self . load_model ( model_name )
img = cv2 . imread ( self . image )
img_width_model = self . img_width
img_height_model = self . img_height
def predict ( self , model_name , img , use_patches ) :
model , model_height , model_width , n_classes = self . load_model ( model_name )
if self . patches == ' true ' or self . patches == ' True ' :
if use_patches :
margin = int ( 0.1 * img_width_model )
margin = int ( 0.1 * model_width )
width_mid = img_width_model - 2 * margin
height_mid = img_height_model - 2 * margin
width_mid = model_width - 2 * margin
height_mid = model_height - 2 * margin
img = img / float ( 255.0 )
@ -89,31 +78,28 @@ class sbb_binarize:
if i == 0 :
index_x_d = i * width_mid
index_x_u = index_x_d + img_width_model
index_x_u = index_x_d + model_width
elif i > 0 :
index_x_d = i * width_mid
index_x_u = index_x_d + img_width_model
index_x_u = index_x_d + model_width
if j == 0 :
index_y_d = j * height_mid
index_y_u = index_y_d + img_height_model
index_y_u = index_y_d + model_height
elif j > 0 :
index_y_d = j * height_mid
index_y_u = index_y_d + img_height_model
index_y_u = index_y_d + model_height
if index_x_u > img_w :
index_x_u = img_w
index_x_d = img_w - img_width_model
index_x_d = img_w - model_width
if index_y_u > img_h :
index_y_u = img_h
index_y_d = img_h - img_height_model
index_y_d = img_h - model_height
img_patch = img [ index_y_d : index_y_u , index_x_d : index_x_u , : ]
label_p_pred = self . model . predict (
img_patch . reshape ( 1 , img_patch . shape [ 0 ] , img_patch . shape [ 1 ] , img_patch . shape [ 2 ] ) )
label_p_pred = model . predict ( img_patch . reshape ( 1 , img_patch . shape [ 0 ] , img_patch . shape [ 1 ] , img_patch . shape [ 2 ] ) )
seg = np . argmax ( label_p_pred , axis = 3 ) [ 0 ]
@ -124,72 +110,63 @@ class sbb_binarize:
seg = seg [ 0 : seg . shape [ 0 ] - margin , 0 : seg . shape [ 1 ] - margin ]
mask_true [ index_y_d + 0 : index_y_u - margin , index_x_d + 0 : index_x_u - margin ] = seg
prediction_true [ index_y_d + 0 : index_y_u - margin , index_x_d + 0 : index_x_u - margin ,
: ] = seg_color
prediction_true [ index_y_d + 0 : index_y_u - margin , index_x_d + 0 : index_x_u - margin , : ] = seg_color
elif i == nxf - 1 and j == nyf - 1 :
seg_color = seg_color [ margin : seg_color . shape [ 0 ] - 0 , margin : seg_color . shape [ 1 ] - 0 , : ]
seg = seg [ margin : seg . shape [ 0 ] - 0 , margin : seg . shape [ 1 ] - 0 ]
mask_true [ index_y_d + margin : index_y_u - 0 , index_x_d + margin : index_x_u - 0 ] = seg
prediction_true [ index_y_d + margin : index_y_u - 0 , index_x_d + margin : index_x_u - 0 ,
: ] = seg_color
prediction_true [ index_y_d + margin : index_y_u - 0 , index_x_d + margin : index_x_u - 0 , : ] = seg_color
elif i == 0 and j == nyf - 1 :
seg_color = seg_color [ margin : seg_color . shape [ 0 ] - 0 , 0 : seg_color . shape [ 1 ] - margin , : ]
seg = seg [ margin : seg . shape [ 0 ] - 0 , 0 : seg . shape [ 1 ] - margin ]
mask_true [ index_y_d + margin : index_y_u - 0 , index_x_d + 0 : index_x_u - margin ] = seg
prediction_true [ index_y_d + margin : index_y_u - 0 , index_x_d + 0 : index_x_u - margin ,
: ] = seg_color
prediction_true [ index_y_d + margin : index_y_u - 0 , index_x_d + 0 : index_x_u - margin , : ] = seg_color
elif i == nxf - 1 and j == 0 :
seg_color = seg_color [ 0 : seg_color . shape [ 0 ] - margin , margin : seg_color . shape [ 1 ] - 0 , : ]
seg = seg [ 0 : seg . shape [ 0 ] - margin , margin : seg . shape [ 1 ] - 0 ]
mask_true [ index_y_d + 0 : index_y_u - margin , index_x_d + margin : index_x_u - 0 ] = seg
prediction_true [ index_y_d + 0 : index_y_u - margin , index_x_d + margin : index_x_u - 0 ,
: ] = seg_color
prediction_true [ index_y_d + 0 : index_y_u - margin , index_x_d + margin : index_x_u - 0 , : ] = seg_color
elif i == 0 and j != 0 and j != nyf - 1 :
seg_color = seg_color [ margin : seg_color . shape [ 0 ] - margin , 0 : seg_color . shape [ 1 ] - margin , : ]
seg = seg [ margin : seg . shape [ 0 ] - margin , 0 : seg . shape [ 1 ] - margin ]
mask_true [ index_y_d + margin : index_y_u - margin , index_x_d + 0 : index_x_u - margin ] = seg
prediction_true [ index_y_d + margin : index_y_u - margin , index_x_d + 0 : index_x_u - margin ,
: ] = seg_color
prediction_true [ index_y_d + margin : index_y_u - margin , index_x_d + 0 : index_x_u - margin , : ] = seg_color
elif i == nxf - 1 and j != 0 and j != nyf - 1 :
seg_color = seg_color [ margin : seg_color . shape [ 0 ] - margin , margin : seg_color . shape [ 1 ] - 0 , : ]
seg = seg [ margin : seg . shape [ 0 ] - margin , margin : seg . shape [ 1 ] - 0 ]
mask_true [ index_y_d + margin : index_y_u - margin , index_x_d + margin : index_x_u - 0 ] = seg
prediction_true [ index_y_d + margin : index_y_u - margin , index_x_d + margin : index_x_u - 0 ,
: ] = seg_color
prediction_true [ index_y_d + margin : index_y_u - margin , index_x_d + margin : index_x_u - 0 , : ] = seg_color
elif i != 0 and i != nxf - 1 and j == 0 :
seg_color = seg_color [ 0 : seg_color . shape [ 0 ] - margin , margin : seg_color . shape [ 1 ] - margin , : ]
seg = seg [ 0 : seg . shape [ 0 ] - margin , margin : seg . shape [ 1 ] - margin ]
mask_true [ index_y_d + 0 : index_y_u - margin , index_x_d + margin : index_x_u - margin ] = seg
prediction_true [ index_y_d + 0 : index_y_u - margin , index_x_d + margin : index_x_u - margin ,
: ] = seg_color
prediction_true [ index_y_d + 0 : index_y_u - margin , index_x_d + margin : index_x_u - margin , : ] = seg_color
elif i != 0 and i != nxf - 1 and j == nyf - 1 :
seg_color = seg_color [ margin : seg_color . shape [ 0 ] - 0 , margin : seg_color . shape [ 1 ] - margin , : ]
seg = seg [ margin : seg . shape [ 0 ] - 0 , margin : seg . shape [ 1 ] - margin ]
mask_true [ index_y_d + margin : index_y_u - 0 , index_x_d + margin : index_x_u - margin ] = seg
prediction_true [ index_y_d + margin : index_y_u - 0 , index_x_d + margin : index_x_u - margin ,
: ] = seg_color
prediction_true [ index_y_d + margin : index_y_u - 0 , index_x_d + margin : index_x_u - margin , : ] = seg_color
else :
seg_color = seg_color [ margin : seg_color . shape [ 0 ] - margin , margin : seg_color . shape [ 1 ] - margin , : ]
seg = seg [ margin : seg . shape [ 0 ] - margin , margin : seg . shape [ 1 ] - margin ]
mask_true [ index_y_d + margin : index_y_u - margin , index_x_d + margin : index_x_u - margin ] = seg
prediction_true [ index_y_d + margin : index_y_u - margin , index_x_d + margin : index_x_u - margin ,
: ] = seg_color
prediction_true [ index_y_d + margin : index_y_u - margin , index_x_d + margin : index_x_u - margin , : ] = seg_color
prediction_true = prediction_true . astype ( np . uint8 )
@ -197,24 +174,28 @@ class sbb_binarize:
img_h_page = img . shape [ 0 ]
img_w_page = img . shape [ 1 ]
img = img / float ( 255.0 )
img = self . resize_image ( img , img_height_model, img_width_model )
img = resize_image ( img , model_height, model_width )
label_p_pred = self . model . predict (
img . reshape ( 1 , img . shape [ 0 ] , img . shape [ 1 ] , img . shape [ 2 ] ) )
label_p_pred = model . predict ( img . reshape ( 1 , img . shape [ 0 ] , img . shape [ 1 ] , img . shape [ 2 ] ) )
seg = np . argmax ( label_p_pred , axis = 3 ) [ 0 ]
seg_color = np . repeat ( seg [ : , : , np . newaxis ] , 3 , axis = 2 )
prediction_true = self . resize_image ( seg_color , img_h_page , img_w_page )
prediction_true = resize_image ( seg_color , img_h_page , img_w_page )
prediction_true = prediction_true . astype ( np . uint8 )
return prediction_true [ : , : , 0 ]
def run ( self ) :
self . start_new_session_and_model ( )
models_n = os . listdir ( self . model_dir )
def run ( self , image = None , image_path = None , save = None , use_patches = False ) :
if ( image is not None and image_path is not None ) or \
( image is None and image_path is None ) :
raise ValueError ( " Must pass either a opencv2 image or an image_path " )
if image_path is not None :
image = cv2 . imread ( image )
self . start_new_session ( )
list_of_model_files = listdir ( self . model_dir )
img_last = 0
for model_in in models_n :
for model_in in list_of_ model_file s:
res = self . predict ( model_in )
res = self . predict ( model_in , image , use_patches )
img_fin = np . zeros ( ( res . shape [ 0 ] , res . shape [ 1 ] , 3 ) )
res [ : , : ] [ res [ : , : ] == 0 ] = 2
@ -227,29 +208,10 @@ class sbb_binarize:
img_fin = img_fin . astype ( np . uint8 )
img_fin = ( res [ : , : ] == 0 ) * 255
img_last = img_last + img_fin
kernel = np . ones ( ( 5 , 5 ) , np . uint8 )
img_last [ : , : ] [ img_last [ : , : ] > 0 ] = 255
img_last = ( img_last [ : , : ] == 0 ) * 255
if self . save is not None :
cv2 . imwrite ( self . save , img_last )
def main ( ) :
parser = argparse . ArgumentParser ( )
parser . add_argument ( ' -i ' , ' --image ' , dest = ' inp1 ' , default = None , help = ' image. ' )
parser . add_argument ( ' -p ' , ' --patches ' , dest = ' inp3 ' , default = False , help = ' by setting this parameter to true you let the model to see the image in patches. ' )
parser . add_argument ( ' -s ' , ' --save ' , dest = ' inp4 ' , default = False , help = ' save prediction with a given name here. The name and format should be given (outputname.tif). ' )
parser . add_argument ( ' -m ' , ' --model ' , dest = ' inp2 ' , default = None , help = ' models directory. ' )
options = parser . parse_args ( )
possibles = globals ( )
possibles . update ( locals ( ) )
x = sbb_binarize ( options . inp1 , options . inp2 , options . inp3 , options . inp4 )
x . run ( )
if __name__ == " __main__ " :
main ( )
if save :
cv2 . imwrite ( save , img_last )
return img_last