You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
644 lines
27 KiB
Python
644 lines
27 KiB
Python
import sys
|
|
import os
|
|
import numpy as np
|
|
import warnings
|
|
import cv2
|
|
import seaborn as sns
|
|
from tensorflow.keras.models import load_model
|
|
import tensorflow as tf
|
|
from tensorflow.keras import backend as K
|
|
from tensorflow.keras import layers
|
|
import tensorflow.keras.losses
|
|
from tensorflow.keras.layers import *
|
|
from models import *
|
|
from gt_gen_utils import *
|
|
import click
|
|
import json
|
|
from tensorflow.python.keras import backend as tensorflow_backend
|
|
import xml.etree.ElementTree as ET
|
|
import matplotlib.pyplot as plt
|
|
|
|
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore")
|
|
|
|
__doc__=\
|
|
"""
|
|
Tool to load model and predict for given image.
|
|
"""
|
|
|
|
class sbb_predict:
|
|
def __init__(self,image, model, task, config_params_model, patches, save, save_layout, ground_truth, xml_file, out, min_area):
|
|
self.image=image
|
|
self.patches=patches
|
|
self.save=save
|
|
self.save_layout=save_layout
|
|
self.model_dir=model
|
|
self.ground_truth=ground_truth
|
|
self.task=task
|
|
self.config_params_model=config_params_model
|
|
self.xml_file = xml_file
|
|
self.out = out
|
|
if min_area:
|
|
self.min_area = float(min_area)
|
|
else:
|
|
self.min_area = 0
|
|
|
|
def resize_image(self,img_in,input_height,input_width):
|
|
return cv2.resize( img_in, ( input_width,input_height) ,interpolation=cv2.INTER_NEAREST)
|
|
|
|
|
|
def color_images(self,seg):
|
|
ann_u=range(self.n_classes)
|
|
if len(np.shape(seg))==3:
|
|
seg=seg[:,:,0]
|
|
|
|
seg_img=np.zeros((np.shape(seg)[0],np.shape(seg)[1],3)).astype(np.uint8)
|
|
colors=sns.color_palette("hls", self.n_classes)
|
|
|
|
for c in ann_u:
|
|
c=int(c)
|
|
segl=(seg==c)
|
|
seg_img[:,:,0][seg==c]=c
|
|
seg_img[:,:,1][seg==c]=c
|
|
seg_img[:,:,2][seg==c]=c
|
|
return seg_img
|
|
|
|
def otsu_copy_binary(self,img):
|
|
img_r=np.zeros((img.shape[0],img.shape[1],3))
|
|
img1=img[:,:,0]
|
|
|
|
#print(img.min())
|
|
#print(img[:,:,0].min())
|
|
#blur = cv2.GaussianBlur(img,(5,5))
|
|
#ret3,th3 = cv2.threshold(blur,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
|
|
retval1, threshold1 = cv2.threshold(img1, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
|
|
|
|
|
|
|
|
img_r[:,:,0]=threshold1
|
|
img_r[:,:,1]=threshold1
|
|
img_r[:,:,2]=threshold1
|
|
#img_r=img_r/float(np.max(img_r))*255
|
|
return img_r
|
|
|
|
def otsu_copy(self,img):
|
|
img_r=np.zeros((img.shape[0],img.shape[1],3))
|
|
#img1=img[:,:,0]
|
|
|
|
#print(img.min())
|
|
#print(img[:,:,0].min())
|
|
#blur = cv2.GaussianBlur(img,(5,5))
|
|
#ret3,th3 = cv2.threshold(blur,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
|
|
_, threshold1 = cv2.threshold(img[:,:,0], 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
|
|
_, threshold2 = cv2.threshold(img[:,:,1], 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
|
|
_, threshold3 = cv2.threshold(img[:,:,2], 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
|
|
|
|
|
|
|
|
img_r[:,:,0]=threshold1
|
|
img_r[:,:,1]=threshold2
|
|
img_r[:,:,2]=threshold3
|
|
###img_r=img_r/float(np.max(img_r))*255
|
|
return img_r
|
|
|
|
def soft_dice_loss(self,y_true, y_pred, epsilon=1e-6):
|
|
|
|
axes = tuple(range(1, len(y_pred.shape)-1))
|
|
|
|
numerator = 2. * K.sum(y_pred * y_true, axes)
|
|
|
|
denominator = K.sum(K.square(y_pred) + K.square(y_true), axes)
|
|
return 1.00 - K.mean(numerator / (denominator + epsilon)) # average over classes and batch
|
|
|
|
def weighted_categorical_crossentropy(self,weights=None):
|
|
|
|
def loss(y_true, y_pred):
|
|
labels_floats = tf.cast(y_true, tf.float32)
|
|
per_pixel_loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=labels_floats,logits=y_pred)
|
|
|
|
if weights is not None:
|
|
weight_mask = tf.maximum(tf.reduce_max(tf.constant(
|
|
np.array(weights, dtype=np.float32)[None, None, None])
|
|
* labels_floats, axis=-1), 1.0)
|
|
per_pixel_loss = per_pixel_loss * weight_mask[:, :, :, None]
|
|
return tf.reduce_mean(per_pixel_loss)
|
|
return self.loss
|
|
|
|
|
|
def IoU(self,Yi,y_predi):
|
|
## mean Intersection over Union
|
|
## Mean IoU = TP/(FN + TP + FP)
|
|
|
|
IoUs = []
|
|
Nclass = np.unique(Yi)
|
|
for c in Nclass:
|
|
TP = np.sum( (Yi == c)&(y_predi==c) )
|
|
FP = np.sum( (Yi != c)&(y_predi==c) )
|
|
FN = np.sum( (Yi == c)&(y_predi != c))
|
|
IoU = TP/float(TP + FP + FN)
|
|
if self.n_classes>2:
|
|
print("class {:02.0f}: #TP={:6.0f}, #FP={:6.0f}, #FN={:5.0f}, IoU={:4.3f}".format(c,TP,FP,FN,IoU))
|
|
IoUs.append(IoU)
|
|
if self.n_classes>2:
|
|
mIoU = np.mean(IoUs)
|
|
print("_________________")
|
|
print("Mean IoU: {:4.3f}".format(mIoU))
|
|
return mIoU
|
|
elif self.n_classes==2:
|
|
mIoU = IoUs[1]
|
|
print("_________________")
|
|
print("IoU: {:4.3f}".format(mIoU))
|
|
return mIoU
|
|
|
|
def start_new_session_and_model(self):
|
|
|
|
config = tf.compat.v1.ConfigProto()
|
|
config.gpu_options.allow_growth = True
|
|
|
|
session = tf.compat.v1.Session(config=config) # tf.InteractiveSession()
|
|
tensorflow_backend.set_session(session)
|
|
#tensorflow.keras.layers.custom_layer = PatchEncoder
|
|
#tensorflow.keras.layers.custom_layer = Patches
|
|
self.model = load_model(self.model_dir , compile=False,custom_objects = {"PatchEncoder": PatchEncoder, "Patches": Patches})
|
|
#config = tf.ConfigProto()
|
|
#config.gpu_options.allow_growth=True
|
|
|
|
#self.session = tf.InteractiveSession()
|
|
#keras.losses.custom_loss = self.weighted_categorical_crossentropy
|
|
#self.model = load_model(self.model_dir , compile=False)
|
|
|
|
|
|
##if self.weights_dir!=None:
|
|
##self.model.load_weights(self.weights_dir)
|
|
|
|
if (self.task != 'classification' and self.task != 'reading_order'):
|
|
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 visualize_model_output(self, prediction, img, task):
|
|
if task == "binarization":
|
|
prediction = prediction * -1
|
|
prediction = prediction + 1
|
|
added_image = prediction * 255
|
|
layout_only = None
|
|
else:
|
|
unique_classes = np.unique(prediction[:,:,0])
|
|
rgb_colors = {'0' : [255, 255, 255],
|
|
'1' : [255, 0, 0],
|
|
'2' : [255, 125, 0],
|
|
'3' : [255, 0, 125],
|
|
'4' : [125, 125, 125],
|
|
'5' : [125, 125, 0],
|
|
'6' : [0, 125, 255],
|
|
'7' : [0, 125, 0],
|
|
'8' : [125, 125, 125],
|
|
'9' : [0, 125, 255],
|
|
'10' : [125, 0, 125],
|
|
'11' : [0, 255, 0],
|
|
'12' : [0, 0, 255],
|
|
'13' : [0, 255, 255],
|
|
'14' : [255, 125, 125],
|
|
'15' : [255, 0, 255]}
|
|
|
|
layout_only = np.zeros(prediction.shape)
|
|
|
|
for unq_class in unique_classes:
|
|
rgb_class_unique = rgb_colors[str(int(unq_class))]
|
|
layout_only[:,:,0][prediction[:,:,0]==unq_class] = rgb_class_unique[0]
|
|
layout_only[:,:,1][prediction[:,:,0]==unq_class] = rgb_class_unique[1]
|
|
layout_only[:,:,2][prediction[:,:,0]==unq_class] = rgb_class_unique[2]
|
|
|
|
|
|
|
|
img = self.resize_image(img, layout_only.shape[0], layout_only.shape[1])
|
|
|
|
layout_only = layout_only.astype(np.int32)
|
|
img = img.astype(np.int32)
|
|
|
|
|
|
|
|
added_image = cv2.addWeighted(img,0.5,layout_only,0.1,0)
|
|
|
|
return added_image, layout_only
|
|
|
|
def predict(self):
|
|
self.start_new_session_and_model()
|
|
if self.task == 'classification':
|
|
classes_names = self.config_params_model['classification_classes_name']
|
|
img_1ch = img=cv2.imread(self.image, 0)
|
|
|
|
img_1ch = img_1ch / 255.0
|
|
img_1ch = cv2.resize(img_1ch, (self.config_params_model['input_height'], self.config_params_model['input_width']), interpolation=cv2.INTER_NEAREST)
|
|
img_in = np.zeros((1, img_1ch.shape[0], img_1ch.shape[1], 3))
|
|
img_in[0, :, :, 0] = img_1ch[:, :]
|
|
img_in[0, :, :, 1] = img_1ch[:, :]
|
|
img_in[0, :, :, 2] = img_1ch[:, :]
|
|
|
|
label_p_pred = self.model.predict(img_in, verbose=0)
|
|
index_class = np.argmax(label_p_pred[0])
|
|
|
|
print("Predicted Class: {}".format(classes_names[str(int(index_class))]))
|
|
elif self.task == 'reading_order':
|
|
img_height = self.config_params_model['input_height']
|
|
img_width = self.config_params_model['input_width']
|
|
|
|
tree_xml, root_xml, bb_coord_printspace, file_name, id_paragraph, id_header, co_text_paragraph, co_text_header, tot_region_ref, x_len, y_len, index_tot_regions, img_poly = read_xml(self.xml_file)
|
|
_, cy_main, x_min_main, x_max_main, y_min_main, y_max_main, _ = find_new_features_of_contours(co_text_header)
|
|
|
|
img_header_and_sep = np.zeros((y_len,x_len), dtype='uint8')
|
|
|
|
|
|
for j in range(len(cy_main)):
|
|
img_header_and_sep[int(y_max_main[j]):int(y_max_main[j])+12,int(x_min_main[j]):int(x_max_main[j]) ] = 1
|
|
|
|
co_text_all = co_text_paragraph + co_text_header
|
|
id_all_text = id_paragraph + id_header
|
|
|
|
|
|
##texts_corr_order_index = [index_tot_regions[tot_region_ref.index(i)] for i in id_all_text ]
|
|
##texts_corr_order_index_int = [int(x) for x in texts_corr_order_index]
|
|
texts_corr_order_index_int = list(np.array(range(len(co_text_all))))
|
|
|
|
#print(texts_corr_order_index_int)
|
|
|
|
max_area = 1
|
|
#print(np.shape(co_text_all[0]), len( np.shape(co_text_all[0]) ),'co_text_all')
|
|
#co_text_all = filter_contours_area_of_image_tables(img_poly, co_text_all, _, max_area, min_area)
|
|
#print(co_text_all,'co_text_all')
|
|
co_text_all, texts_corr_order_index_int = filter_contours_area_of_image(img_poly, co_text_all, texts_corr_order_index_int, max_area, self.min_area)
|
|
|
|
#print(texts_corr_order_index_int)
|
|
|
|
#co_text_all = [co_text_all[index] for index in texts_corr_order_index_int]
|
|
id_all_text = [id_all_text[index] for index in texts_corr_order_index_int]
|
|
|
|
labels_con = np.zeros((y_len,x_len,len(co_text_all)),dtype='uint8')
|
|
for i in range(len(co_text_all)):
|
|
img_label = np.zeros((y_len,x_len,3),dtype='uint8')
|
|
img_label=cv2.fillPoly(img_label, pts =[co_text_all[i]], color=(1,1,1))
|
|
labels_con[:,:,i] = img_label[:,:,0]
|
|
|
|
if bb_coord_printspace:
|
|
#bb_coord_printspace[x,y,w,h,_,_]
|
|
x = bb_coord_printspace[0]
|
|
y = bb_coord_printspace[1]
|
|
w = bb_coord_printspace[2]
|
|
h = bb_coord_printspace[3]
|
|
labels_con = labels_con[y:y+h, x:x+w, :]
|
|
img_poly = img_poly[y:y+h, x:x+w, :]
|
|
img_header_and_sep = img_header_and_sep[y:y+h, x:x+w]
|
|
|
|
|
|
|
|
img3= np.copy(img_poly)
|
|
labels_con = resize_image(labels_con, img_height, img_width)
|
|
|
|
img_header_and_sep = resize_image(img_header_and_sep, img_height, img_width)
|
|
|
|
img3= resize_image (img3, img_height, img_width)
|
|
img3 = img3.astype(np.uint16)
|
|
|
|
inference_bs = 1#4
|
|
|
|
input_1= np.zeros( (inference_bs, img_height, img_width,3))
|
|
|
|
|
|
starting_list_of_regions = []
|
|
starting_list_of_regions.append( list(range(labels_con.shape[2])) )
|
|
|
|
index_update = 0
|
|
index_selected = starting_list_of_regions[0]
|
|
|
|
scalibility_num = 0
|
|
while index_update>=0:
|
|
ij_list = starting_list_of_regions[index_update]
|
|
i = ij_list[0]
|
|
ij_list.pop(0)
|
|
|
|
|
|
pr_list = []
|
|
post_list = []
|
|
|
|
batch_counter = 0
|
|
tot_counter = 1
|
|
|
|
tot_iteration = len(ij_list)
|
|
full_bs_ite= tot_iteration//inference_bs
|
|
last_bs = tot_iteration % inference_bs
|
|
|
|
jbatch_indexer =[]
|
|
for j in ij_list:
|
|
img1= np.repeat(labels_con[:,:,i][:, :, np.newaxis], 3, axis=2)
|
|
img2 = np.repeat(labels_con[:,:,j][:, :, np.newaxis], 3, axis=2)
|
|
|
|
|
|
img2[:,:,0][img3[:,:,0]==5] = 2
|
|
img2[:,:,0][img_header_and_sep[:,:]==1] = 3
|
|
|
|
|
|
|
|
img1[:,:,0][img3[:,:,0]==5] = 2
|
|
img1[:,:,0][img_header_and_sep[:,:]==1] = 3
|
|
|
|
#input_1= np.zeros( (height1, width1,3))
|
|
|
|
|
|
jbatch_indexer.append(j)
|
|
|
|
input_1[batch_counter,:,:,0] = img1[:,:,0]/3.
|
|
input_1[batch_counter,:,:,2] = img2[:,:,0]/3.
|
|
input_1[batch_counter,:,:,1] = img3[:,:,0]/5.
|
|
#input_1[batch_counter,:,:,:]= np.zeros( (batch_counter, height1, width1,3))
|
|
batch_counter = batch_counter+1
|
|
|
|
#input_1[:,:,0] = img1[:,:,0]/3.
|
|
#input_1[:,:,2] = img2[:,:,0]/3.
|
|
#input_1[:,:,1] = img3[:,:,0]/5.
|
|
|
|
if batch_counter==inference_bs or ( (tot_counter//inference_bs)==full_bs_ite and tot_counter%inference_bs==last_bs):
|
|
y_pr = self.model.predict(input_1 , verbose=0)
|
|
scalibility_num = scalibility_num+1
|
|
|
|
if batch_counter==inference_bs:
|
|
iteration_batches = inference_bs
|
|
else:
|
|
iteration_batches = last_bs
|
|
for jb in range(iteration_batches):
|
|
if y_pr[jb][0]>=0.5:
|
|
post_list.append(jbatch_indexer[jb])
|
|
else:
|
|
pr_list.append(jbatch_indexer[jb])
|
|
|
|
batch_counter = 0
|
|
jbatch_indexer = []
|
|
|
|
tot_counter = tot_counter+1
|
|
|
|
starting_list_of_regions, index_update = update_list_and_return_first_with_length_bigger_than_one(index_update, i, pr_list, post_list,starting_list_of_regions)
|
|
|
|
|
|
index_sort = [i[0] for i in starting_list_of_regions ]
|
|
|
|
id_all_text = np.array(id_all_text)[index_sort]
|
|
|
|
alltags=[elem.tag for elem in root_xml.iter()]
|
|
|
|
|
|
|
|
link=alltags[0].split('}')[0]+'}'
|
|
name_space = alltags[0].split('}')[0]
|
|
name_space = name_space.split('{')[1]
|
|
|
|
page_element = root_xml.find(link+'Page')
|
|
|
|
"""
|
|
ro_subelement = ET.SubElement(page_element, 'ReadingOrder')
|
|
#print(page_element, 'page_element')
|
|
|
|
#new_element = ET.Element('ReadingOrder')
|
|
|
|
new_element_element = ET.Element('OrderedGroup')
|
|
new_element_element.set('id', "ro357564684568544579089")
|
|
|
|
for index, id_text in enumerate(id_all_text):
|
|
new_element_2 = ET.Element('RegionRefIndexed')
|
|
new_element_2.set('regionRef', id_all_text[index])
|
|
new_element_2.set('index', str(index_sort[index]))
|
|
|
|
new_element_element.append(new_element_2)
|
|
|
|
ro_subelement.append(new_element_element)
|
|
"""
|
|
##ro_subelement = ET.SubElement(page_element, 'ReadingOrder')
|
|
|
|
ro_subelement = ET.Element('ReadingOrder')
|
|
|
|
ro_subelement2 = ET.SubElement(ro_subelement, 'OrderedGroup')
|
|
ro_subelement2.set('id', "ro357564684568544579089")
|
|
|
|
for index, id_text in enumerate(id_all_text):
|
|
new_element_2 = ET.SubElement(ro_subelement2, 'RegionRefIndexed')
|
|
new_element_2.set('regionRef', id_all_text[index])
|
|
new_element_2.set('index', str(index))
|
|
|
|
if (link+'PrintSpace' in alltags) or (link+'Border' in alltags):
|
|
page_element.insert(1, ro_subelement)
|
|
else:
|
|
page_element.insert(0, ro_subelement)
|
|
|
|
alltags=[elem.tag for elem in root_xml.iter()]
|
|
|
|
ET.register_namespace("",name_space)
|
|
tree_xml.write(os.path.join(self.out, file_name+'.xml'),xml_declaration=True,method='xml',encoding="utf8",default_namespace=None)
|
|
#tree_xml.write('library2.xml')
|
|
|
|
else:
|
|
if self.patches:
|
|
#def textline_contours(img,input_width,input_height,n_classes,model):
|
|
|
|
img=cv2.imread(self.image)
|
|
self.img_org = np.copy(img)
|
|
|
|
if img.shape[0] < self.img_height:
|
|
img = cv2.resize(img, (img.shape[1], self.img_width), interpolation=cv2.INTER_NEAREST)
|
|
|
|
if img.shape[1] < self.img_width:
|
|
img = cv2.resize(img, (self.img_height, img.shape[0]), interpolation=cv2.INTER_NEAREST)
|
|
margin = int(0.1 * self.img_width)
|
|
width_mid = self.img_width - 2 * margin
|
|
height_mid = self.img_height - 2 * margin
|
|
img = img / float(255.0)
|
|
|
|
img_h = img.shape[0]
|
|
img_w = img.shape[1]
|
|
|
|
prediction_true = np.zeros((img_h, img_w, 3))
|
|
nxf = img_w / float(width_mid)
|
|
nyf = img_h / float(height_mid)
|
|
|
|
nxf = int(nxf) + 1 if nxf > int(nxf) else int(nxf)
|
|
nyf = int(nyf) + 1 if nyf > int(nyf) else int(nyf)
|
|
|
|
for i in range(nxf):
|
|
for j in range(nyf):
|
|
if i == 0:
|
|
index_x_d = i * width_mid
|
|
index_x_u = index_x_d + self.img_width
|
|
else:
|
|
index_x_d = i * width_mid
|
|
index_x_u = index_x_d + self.img_width
|
|
if j == 0:
|
|
index_y_d = j * height_mid
|
|
index_y_u = index_y_d + self.img_height
|
|
else:
|
|
index_y_d = j * height_mid
|
|
index_y_u = index_y_d + self.img_height
|
|
|
|
if index_x_u > img_w:
|
|
index_x_u = img_w
|
|
index_x_d = img_w - self.img_width
|
|
if index_y_u > img_h:
|
|
index_y_u = img_h
|
|
index_y_d = img_h - self.img_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]),
|
|
verbose=0)
|
|
|
|
if self.task == 'enhancement':
|
|
seg = label_p_pred[0, :, :, :]
|
|
seg = seg * 255
|
|
elif self.task == 'segmentation' or self.task == 'binarization':
|
|
seg = np.argmax(label_p_pred, axis=3)[0]
|
|
seg = np.repeat(seg[:, :, np.newaxis], 3, axis=2)
|
|
|
|
|
|
if i == 0 and j == 0:
|
|
seg = seg[0 : seg.shape[0] - margin, 0 : seg.shape[1] - margin]
|
|
prediction_true[index_y_d + 0 : index_y_u - margin, index_x_d + 0 : index_x_u - margin, :] = seg
|
|
elif i == nxf - 1 and j == nyf - 1:
|
|
seg = seg[margin : seg.shape[0] - 0, margin : seg.shape[1] - 0]
|
|
prediction_true[index_y_d + margin : index_y_u - 0, index_x_d + margin : index_x_u - 0, :] = seg
|
|
elif i == 0 and j == nyf - 1:
|
|
seg = seg[margin : seg.shape[0] - 0, 0 : seg.shape[1] - margin]
|
|
prediction_true[index_y_d + margin : index_y_u - 0, index_x_d + 0 : index_x_u - margin, :] = seg
|
|
elif i == nxf - 1 and j == 0:
|
|
seg = seg[0 : seg.shape[0] - margin, margin : seg.shape[1] - 0]
|
|
prediction_true[index_y_d + 0 : index_y_u - margin, index_x_d + margin : index_x_u - 0, :] = seg
|
|
elif i == 0 and j != 0 and j != nyf - 1:
|
|
seg = seg[margin : seg.shape[0] - margin, 0 : seg.shape[1] - margin]
|
|
prediction_true[index_y_d + margin : index_y_u - margin, index_x_d + 0 : index_x_u - margin, :] = seg
|
|
elif i == nxf - 1 and j != 0 and j != nyf - 1:
|
|
seg = seg[margin : seg.shape[0] - margin, margin : seg.shape[1] - 0]
|
|
prediction_true[index_y_d + margin : index_y_u - margin, index_x_d + margin : index_x_u - 0, :] = seg
|
|
elif i != 0 and i != nxf - 1 and j == 0:
|
|
seg = seg[0 : seg.shape[0] - margin, margin : seg.shape[1] - margin]
|
|
prediction_true[index_y_d + 0 : index_y_u - margin, index_x_d + margin : index_x_u - margin, :] = seg
|
|
elif i != 0 and i != nxf - 1 and j == nyf - 1:
|
|
seg = seg[margin : seg.shape[0] - 0, margin : seg.shape[1] - margin]
|
|
prediction_true[index_y_d + margin : index_y_u - 0, index_x_d + margin : index_x_u - margin, :] = seg
|
|
else:
|
|
seg = seg[margin : seg.shape[0] - margin, margin : seg.shape[1] - margin]
|
|
prediction_true[index_y_d + margin : index_y_u - margin, index_x_d + margin : index_x_u - margin, :] = seg
|
|
prediction_true = prediction_true.astype(int)
|
|
prediction_true = cv2.resize(prediction_true, (self.img_org.shape[1], self.img_org.shape[0]), interpolation=cv2.INTER_NEAREST)
|
|
return prediction_true
|
|
|
|
else:
|
|
|
|
img=cv2.imread(self.image)
|
|
self.img_org = np.copy(img)
|
|
|
|
width=self.img_width
|
|
height=self.img_height
|
|
|
|
img=img/255.0
|
|
img=self.resize_image(img,self.img_height,self.img_width)
|
|
|
|
|
|
label_p_pred=self.model.predict(
|
|
img.reshape(1,img.shape[0],img.shape[1],img.shape[2]))
|
|
|
|
if self.task == 'enhancement':
|
|
seg = label_p_pred[0, :, :, :]
|
|
seg = seg * 255
|
|
elif self.task == 'segmentation' or self.task == 'binarization':
|
|
seg = np.argmax(label_p_pred, axis=3)[0]
|
|
seg = np.repeat(seg[:, :, np.newaxis], 3, axis=2)
|
|
|
|
prediction_true = seg.astype(int)
|
|
|
|
prediction_true = cv2.resize(prediction_true, (self.img_org.shape[1], self.img_org.shape[0]), interpolation=cv2.INTER_NEAREST)
|
|
return prediction_true
|
|
|
|
|
|
|
|
def run(self):
|
|
res=self.predict()
|
|
if (self.task == 'classification' or self.task == 'reading_order'):
|
|
pass
|
|
elif self.task == 'enhancement':
|
|
if self.save:
|
|
cv2.imwrite(self.save,res)
|
|
else:
|
|
img_seg_overlayed, only_layout = self.visualize_model_output(res, self.img_org, self.task)
|
|
if self.save:
|
|
cv2.imwrite(self.save,img_seg_overlayed)
|
|
cv2.imwrite(self.save_layout, only_layout)
|
|
|
|
if self.ground_truth:
|
|
gt_img=cv2.imread(self.ground_truth)
|
|
self.IoU(gt_img[:,:,0],res[:,:,0])
|
|
|
|
@click.command()
|
|
@click.option(
|
|
"--image",
|
|
"-i",
|
|
help="image filename",
|
|
type=click.Path(exists=True, dir_okay=False),
|
|
)
|
|
@click.option(
|
|
"--out",
|
|
"-o",
|
|
help="output directory where xml with detected reading order will be written.",
|
|
type=click.Path(exists=True, file_okay=False),
|
|
)
|
|
@click.option(
|
|
"--patches/--no-patches",
|
|
"-p/-nop",
|
|
is_flag=True,
|
|
help="if this parameter set to true, this tool will try to do inference in patches.",
|
|
)
|
|
@click.option(
|
|
"--save",
|
|
"-s",
|
|
help="save prediction as a png file in current folder.",
|
|
)
|
|
@click.option(
|
|
"--save_layout",
|
|
"-sl",
|
|
help="save layout prediction only as a png file in current folder.",
|
|
)
|
|
@click.option(
|
|
"--model",
|
|
"-m",
|
|
help="directory of models",
|
|
type=click.Path(exists=True, file_okay=False),
|
|
required=True,
|
|
)
|
|
@click.option(
|
|
"--ground_truth",
|
|
"-gt",
|
|
help="ground truth directory if you want to see the iou of prediction.",
|
|
)
|
|
@click.option(
|
|
"--xml_file",
|
|
"-xml",
|
|
help="xml file with layout coordinates that reading order detection will be implemented on. The result will be written in the same xml file.",
|
|
)
|
|
|
|
@click.option(
|
|
"--min_area",
|
|
"-min",
|
|
help="min area size of regions considered for reading order detection. The default value is zero and means that all text regions are considered for reading order.",
|
|
)
|
|
def main(image, model, patches, save, save_layout, ground_truth, xml_file, out, min_area):
|
|
with open(os.path.join(model,'config.json')) as f:
|
|
config_params_model = json.load(f)
|
|
task = config_params_model['task']
|
|
if (task != 'classification' and task != 'reading_order'):
|
|
if not save:
|
|
print("Error: You used one of segmentation or binarization task but not set -s, you need a filename to save visualized output with -s")
|
|
sys.exit(1)
|
|
x=sbb_predict(image, model, task, config_params_model, patches, save, save_layout, ground_truth, xml_file, out, min_area)
|
|
x.run()
|
|
|
|
if __name__=="__main__":
|
|
main()
|
|
|
|
|
|
|
|
|