From 1a95bca22dcdda46b689816ffc0d8da052f43530 Mon Sep 17 00:00:00 2001 From: cneud <952378+cneud@users.noreply.github.com> Date: Fri, 23 Feb 2024 22:44:27 +0100 Subject: [PATCH] more code formatting --- eynollah/eynollah/cli.py | 5 +- eynollah/eynollah/ocrd_cli.py | 2 + eynollah/eynollah/plot.py | 91 ++- eynollah/eynollah/processor.py | 1 + ..._model_load_pretrained_weights_and_save.py | 16 +- eynollah/eynollah/train/metrics.py | 189 +++-- eynollah/eynollah/train/models.py | 237 +++--- eynollah/eynollah/train/pagexml2img.py | 361 +++++---- eynollah/eynollah/train/train.py | 277 +++---- eynollah/eynollah/train/utils.py | 761 +++++++++--------- eynollah/eynollah/utils/contour.py | 74 +- eynollah/eynollah/utils/counter.py | 1 + eynollah/eynollah/utils/drop_capitals.py | 203 ++--- eynollah/eynollah/utils/marginals.py | 304 ++++--- eynollah/eynollah/utils/pil_cv2.py | 5 +- eynollah/eynollah/utils/resize.py | 1 + eynollah/eynollah/utils/rotate.py | 12 +- eynollah/eynollah/utils/xml.py | 3 + eynollah/eynollah/writer.py | 189 +++-- tests/base.py | 5 +- tests/test_counter.py | 4 + tests/test_dpi.py | 2 + tests/test_run.py | 2 + tests/test_xml.py | 2 + 24 files changed, 1418 insertions(+), 1329 deletions(-) diff --git a/eynollah/eynollah/cli.py b/eynollah/eynollah/cli.py index b720d83..94bf211 100644 --- a/eynollah/eynollah/cli.py +++ b/eynollah/eynollah/cli.py @@ -199,8 +199,9 @@ def main( ignore_page_extraction=ignore_page_extraction, ) eynollah.run() - #pcgts = eynollah.run() - ##eynollah.writer.write_pagexml(pcgts) + # pcgts = eynollah.run() + # eynollah.writer.write_pagexml(pcgts) + if __name__ == "__main__": main() diff --git a/eynollah/eynollah/ocrd_cli.py b/eynollah/eynollah/ocrd_cli.py index 8929927..499661b 100644 --- a/eynollah/eynollah/ocrd_cli.py +++ b/eynollah/eynollah/ocrd_cli.py @@ -2,10 +2,12 @@ from click import command from ocrd.decorators import ocrd_cli_options, ocrd_cli_wrap_processor + @command() @ocrd_cli_options def main(*args, **kwargs): return ocrd_cli_wrap_processor(EynollahProcessor, *args, **kwargs) + if __name__ == '__main__': main() diff --git a/eynollah/eynollah/plot.py b/eynollah/eynollah/plot.py index b01fc04..ae85f2a 100644 --- a/eynollah/eynollah/plot.py +++ b/eynollah/eynollah/plot.py @@ -9,24 +9,25 @@ from .utils.rotate import rotate_image_different from .utils.resize import resize_image + class EynollahPlotter(): """ Class collecting all the plotting and image writing methods """ def __init__( - self, - *, - dir_out, - dir_of_all, - dir_save_page, - dir_of_deskewed, - dir_of_layout, - dir_of_cropped_images, - image_filename_stem, - image_org=None, - scale_x=1, - scale_y=1, + self, + *, + dir_out, + dir_of_all, + dir_save_page, + dir_of_deskewed, + dir_of_layout, + dir_of_cropped_images, + image_filename_stem, + image_org=None, + scale_x=1, + scale_y=1, ): self.dir_out = dir_out self.dir_of_all = dir_of_all @@ -44,22 +45,23 @@ def save_plot_of_layout_main(self, text_regions_p, image_page): if self.dir_of_layout is not None: values = np.unique(text_regions_p[:, :]) # pixels=['Background' , 'Main text' , 'Heading' , 'Marginalia' ,'Drop capitals' , 'Images' , 'Seperators' , 'Tables', 'Graphics'] - pixels=['Background' , 'Main text' , 'Image' , 'Separator','Marginalia'] + pixels = ['Background', 'Main text', 'Image', 'Separator', 'Marginalia'] values_indexes = [0, 1, 2, 3, 4] plt.figure(figsize=(40, 40)) plt.rcParams["font.size"] = "40" im = plt.imshow(text_regions_p[:, :]) colors = [im.cmap(im.norm(value)) for value in values] - patches = [mpatches.Patch(color=colors[np.where(values == i)[0][0]], label="{l}".format(l=pixels[int(np.where(values_indexes == i)[0][0])])) for i in values] + patches = [mpatches.Patch(color=colors[np.where(values == i)[0][0]], + label="{l}".format(l=pixels[int(np.where(values_indexes == i)[0][0])])) for i in + values] plt.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0, fontsize=40) plt.savefig(os.path.join(self.dir_of_layout, self.image_filename_stem + "_layout_main.png")) - def save_plot_of_layout_main_all(self, text_regions_p, image_page): if self.dir_of_all is not None: values = np.unique(text_regions_p[:, :]) # pixels=['Background' , 'Main text' , 'Heading' , 'Marginalia' ,'Drop capitals' , 'Images' , 'Seperators' , 'Tables', 'Graphics'] - pixels=['Background' , 'Main text' , 'Image' , 'Separator','Marginalia'] + pixels = ['Background', 'Main text', 'Image', 'Separator', 'Marginalia'] values_indexes = [0, 1, 2, 3, 4] plt.figure(figsize=(80, 40)) plt.rcParams["font.size"] = "40" @@ -68,7 +70,9 @@ def save_plot_of_layout_main_all(self, text_regions_p, image_page): plt.subplot(1, 2, 2) im = plt.imshow(text_regions_p[:, :]) colors = [im.cmap(im.norm(value)) for value in values] - patches = [mpatches.Patch(color=colors[np.where(values == i)[0][0]], label="{l}".format(l=pixels[int(np.where(values_indexes == i)[0][0])])) for i in values] + patches = [mpatches.Patch(color=colors[np.where(values == i)[0][0]], + label="{l}".format(l=pixels[int(np.where(values_indexes == i)[0][0])])) for i in + values] plt.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0, fontsize=60) plt.savefig(os.path.join(self.dir_of_all, self.image_filename_stem + "_layout_main_and_page.png")) @@ -82,7 +86,9 @@ def save_plot_of_layout(self, text_regions_p, image_page): plt.rcParams["font.size"] = "40" im = plt.imshow(text_regions_p[:, :]) colors = [im.cmap(im.norm(value)) for value in values] - patches = [mpatches.Patch(color=colors[np.where(values == i)[0][0]], label="{l}".format(l=pixels[int(np.where(values_indexes == i)[0][0])])) for i in values] + patches = [mpatches.Patch(color=colors[np.where(values == i)[0][0]], + label="{l}".format(l=pixels[int(np.where(values_indexes == i)[0][0])])) for i in + values] plt.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0, fontsize=40) plt.savefig(os.path.join(self.dir_of_layout, self.image_filename_stem + "_layout.png")) @@ -99,7 +105,9 @@ def save_plot_of_layout_all(self, text_regions_p, image_page): plt.subplot(1, 2, 2) im = plt.imshow(text_regions_p[:, :]) colors = [im.cmap(im.norm(value)) for value in values] - patches = [mpatches.Patch(color=colors[np.where(values == i)[0][0]], label="{l}".format(l=pixels[int(np.where(values_indexes == i)[0][0])])) for i in values] + patches = [mpatches.Patch(color=colors[np.where(values == i)[0][0]], + label="{l}".format(l=pixels[int(np.where(values_indexes == i)[0][0])])) for i in + values] plt.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0, fontsize=60) plt.savefig(os.path.join(self.dir_of_all, self.image_filename_stem + "_layout_and_page.png")) @@ -115,7 +123,9 @@ def save_plot_of_textlines(self, textline_mask_tot_ea, image_page): plt.subplot(1, 2, 2) im = plt.imshow(textline_mask_tot_ea[:, :]) colors = [im.cmap(im.norm(value)) for value in values] - patches = [mpatches.Patch(color=colors[np.where(values == i)[0][0]], label="{l}".format(l=pixels[int(np.where(values_indexes == i)[0][0])])) for i in values] + patches = [mpatches.Patch(color=colors[np.where(values == i)[0][0]], + label="{l}".format(l=pixels[int(np.where(values_indexes == i)[0][0])])) for i in + values] plt.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0, fontsize=60) plt.savefig(os.path.join(self.dir_of_all, self.image_filename_stem + "_textline_and_page.png")) @@ -131,33 +141,36 @@ def save_page_image(self, image_page): cv2.imwrite(os.path.join(self.dir_of_all, self.image_filename_stem + "_page.png"), image_page) if self.dir_save_page is not None: cv2.imwrite(os.path.join(self.dir_save_page, self.image_filename_stem + "_page.png"), image_page) + def save_enhanced_image(self, img_res): cv2.imwrite(os.path.join(self.dir_out, self.image_filename_stem + "_enhanced.png"), img_res) - + def save_plot_of_textline_density(self, img_patch_org): if self.dir_of_all is not None: - plt.figure(figsize=(80,40)) - plt.rcParams['font.size']='50' - plt.subplot(1,2,1) + plt.figure(figsize=(80, 40)) + plt.rcParams['font.size'] = '50' + plt.subplot(1, 2, 1) plt.imshow(img_patch_org) - plt.subplot(1,2,2) - plt.plot(gaussian_filter1d(img_patch_org.sum(axis=1), 3),np.array(range(len(gaussian_filter1d(img_patch_org.sum(axis=1), 3)))),linewidth=8) - plt.xlabel('Density of textline prediction in direction of X axis',fontsize=60) - plt.ylabel('Height',fontsize=60) - plt.yticks([0,len(gaussian_filter1d(img_patch_org.sum(axis=1), 3))]) + plt.subplot(1, 2, 2) + plt.plot(gaussian_filter1d(img_patch_org.sum(axis=1), 3), + np.array(range(len(gaussian_filter1d(img_patch_org.sum(axis=1), 3)))), linewidth=8) + plt.xlabel('Density of textline prediction in direction of X axis', fontsize=60) + plt.ylabel('Height', fontsize=60) + plt.yticks([0, len(gaussian_filter1d(img_patch_org.sum(axis=1), 3))]) plt.gca().invert_yaxis() - plt.savefig(os.path.join(self.dir_of_all, self.image_filename_stem+'_density_of_textline.png')) + plt.savefig(os.path.join(self.dir_of_all, self.image_filename_stem + '_density_of_textline.png')) def save_plot_of_rotation_angle(self, angels, var_res): if self.dir_of_all is not None: - plt.figure(figsize=(60,30)) - plt.rcParams['font.size']='50' - plt.plot(angels,np.array(var_res),'-o',markersize=25,linewidth=4) - plt.xlabel('angle',fontsize=50) - plt.ylabel('variance of sum of rotated textline in direction of x axis',fontsize=50) - plt.plot(angels[np.argmax(var_res)],var_res[np.argmax(np.array(var_res))] ,'*',markersize=50,label='Angle of deskewing=' +str("{:.2f}".format(angels[np.argmax(var_res)]))+r'$\degree$') + plt.figure(figsize=(60, 30)) + plt.rcParams['font.size'] = '50' + plt.plot(angels, np.array(var_res), '-o', markersize=25, linewidth=4) + plt.xlabel('angle', fontsize=50) + plt.ylabel('variance of sum of rotated textline in direction of x axis', fontsize=50) + plt.plot(angels[np.argmax(var_res)], var_res[np.argmax(np.array(var_res))], '*', markersize=50, + label='Angle of deskewing=' + str("{:.2f}".format(angels[np.argmax(var_res)])) + r'$\degree$') plt.legend(loc='best') - plt.savefig(os.path.join(self.dir_of_all, self.image_filename_stem+'_rotation_angle.png')) + plt.savefig(os.path.join(self.dir_of_all, self.image_filename_stem + '_rotation_angle.png')) def write_images_into_directory(self, img_contours, image_page): if self.dir_of_cropped_images is not None: @@ -167,9 +180,9 @@ def write_images_into_directory(self, img_contours, image_page): box = [x, y, w, h] croped_page, page_coord = crop_image_inside_box(box, image_page) - croped_page = resize_image(croped_page, int(croped_page.shape[0] / self.scale_y), int(croped_page.shape[1] / self.scale_x)) + croped_page = resize_image(croped_page, int(croped_page.shape[0] / self.scale_y), + int(croped_page.shape[1] / self.scale_x)) path = os.path.join(self.dir_of_cropped_images, self.image_filename_stem + "_" + str(index) + ".jpg") cv2.imwrite(path, croped_page) index += 1 - diff --git a/eynollah/eynollah/processor.py b/eynollah/eynollah/processor.py index ccec456..042d081 100644 --- a/eynollah/eynollah/processor.py +++ b/eynollah/eynollah/processor.py @@ -22,6 +22,7 @@ OCRD_TOOL = loads(resource_string(__name__, 'ocrd-tool.json').decode('utf8')) + class EynollahProcessor(Processor): def __init__(self, *args, **kwargs): diff --git a/eynollah/eynollah/train/build_model_load_pretrained_weights_and_save.py b/eynollah/eynollah/train/build_model_load_pretrained_weights_and_save.py index 251e698..40cc1b6 100644 --- a/eynollah/eynollah/train/build_model_load_pretrained_weights_and_save.py +++ b/eynollah/eynollah/train/build_model_load_pretrained_weights_and_save.py @@ -1,7 +1,7 @@ import os import sys import tensorflow as tf -import keras , warnings +import keras, warnings from keras.optimizers import * from sacred import Experiment from models import * @@ -9,25 +9,21 @@ from metrics import * - - def configuration(): gpu_options = tf.compat.v1.GPUOptions(allow_growth=True) session = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(gpu_options=gpu_options)) -if __name__=='__main__': +if __name__ == '__main__': n_classes = 2 input_height = 224 input_width = 448 weight_decay = 1e-6 pretraining = False dir_of_weights = 'model_bin_sbb_ens.h5' - - #configuration() - - model = resnet50_unet(n_classes, input_height, input_width,weight_decay,pretraining) + + # configuration() + + model = resnet50_unet(n_classes, input_height, input_width, weight_decay, pretraining) model.load_weights(dir_of_weights) model.save('./name_in_another_python_version.h5') - - diff --git a/eynollah/eynollah/train/metrics.py b/eynollah/eynollah/train/metrics.py index c63cc22..9d41d9e 100644 --- a/eynollah/eynollah/train/metrics.py +++ b/eynollah/eynollah/train/metrics.py @@ -2,8 +2,8 @@ import tensorflow as tf import numpy as np -def focal_loss(gamma=2., alpha=4.): +def focal_loss(gamma=2., alpha=4.): gamma = float(gamma) alpha = float(alpha) @@ -37,8 +37,10 @@ def focal_loss_fixed(y_true, y_pred): fl = tf.multiply(alpha, tf.multiply(weight, ce)) reduced_fl = tf.reduce_max(fl, axis=1) return tf.reduce_mean(reduced_fl) + return focal_loss_fixed + def weighted_categorical_crossentropy(weights=None): """ weighted_categorical_crossentropy @@ -50,90 +52,102 @@ def weighted_categorical_crossentropy(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) - + 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) + * labels_floats, axis=-1), 1.0) per_pixel_loss = per_pixel_loss * weight_mask[:, :, :, None] return tf.reduce_mean(per_pixel_loss) + return loss + + def image_categorical_cross_entropy(y_true, y_pred, weights=None): """ :param y_true: tensor of shape (batch_size, height, width) representing the ground truth. :param y_pred: tensor of shape (batch_size, height, width) representing the prediction. :return: The mean cross-entropy on softmaxed tensors. """ - + labels_floats = tf.cast(y_true, tf.float32) - per_pixel_loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=labels_floats,logits=y_pred) - + 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) + 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) + + def class_tversky(y_true, y_pred): - smooth = 1.0#1.00 + smooth = 1.0 # 1.00 - y_true = K.permute_dimensions(y_true, (3,1,2,0)) - y_pred = K.permute_dimensions(y_pred, (3,1,2,0)) + y_true = K.permute_dimensions(y_true, (3, 1, 2, 0)) + y_pred = K.permute_dimensions(y_pred, (3, 1, 2, 0)) y_true_pos = K.batch_flatten(y_true) y_pred_pos = K.batch_flatten(y_pred) true_pos = K.sum(y_true_pos * y_pred_pos, 1) - false_neg = K.sum(y_true_pos * (1-y_pred_pos), 1) - false_pos = K.sum((1-y_true_pos)*y_pred_pos, 1) - alpha = 0.2#0.5 - beta=0.8 - return (true_pos + smooth)/(true_pos + alpha*false_neg + (beta)*false_pos + smooth) + false_neg = K.sum(y_true_pos * (1 - y_pred_pos), 1) + false_pos = K.sum((1 - y_true_pos) * y_pred_pos, 1) + alpha = 0.2 # 0.5 + beta = 0.8 + return (true_pos + smooth) / (true_pos + alpha * false_neg + (beta) * false_pos + smooth) -def focal_tversky_loss(y_true,y_pred): + +def focal_tversky_loss(y_true, y_pred): pt_1 = class_tversky(y_true, y_pred) - gamma =1.3#4./3.0#1.3#4.0/3.00# 0.75 - return K.sum(K.pow((1-pt_1), gamma)) + gamma = 1.3 # 4./3.0#1.3#4.0/3.00# 0.75 + return K.sum(K.pow((1 - pt_1), gamma)) + def generalized_dice_coeff2(y_true, y_pred): n_el = 1 - for dim in y_true.shape: + for dim in y_true.shape: n_el *= int(dim) n_cl = y_true.shape[-1] w = K.zeros(shape=(n_cl,)) - w = (K.sum(y_true, axis=(0,1,2)))/(n_el) - w = 1/(w**2+0.000001) - numerator = y_true*y_pred - numerator = w*K.sum(numerator,(0,1,2)) + w = (K.sum(y_true, axis=(0, 1, 2))) / (n_el) + w = 1 / (w ** 2 + 0.000001) + numerator = y_true * y_pred + numerator = w * K.sum(numerator, (0, 1, 2)) numerator = K.sum(numerator) - denominator = y_true+y_pred - denominator = w*K.sum(denominator,(0,1,2)) + denominator = y_true + y_pred + denominator = w * K.sum(denominator, (0, 1, 2)) denominator = K.sum(denominator) - return 2*numerator/denominator + return 2 * numerator / denominator + + def generalized_dice_coeff(y_true, y_pred): - axes = tuple(range(1, len(y_pred.shape)-1)) + axes = tuple(range(1, len(y_pred.shape) - 1)) Ncl = y_pred.shape[-1] w = K.zeros(shape=(Ncl,)) w = K.sum(y_true, axis=axes) - w = 1/(w**2+0.000001) + w = 1 / (w ** 2 + 0.000001) # Compute gen dice coef: - numerator = y_true*y_pred - numerator = w*K.sum(numerator,axes) + numerator = y_true * y_pred + numerator = w * K.sum(numerator, axes) numerator = K.sum(numerator) - denominator = y_true+y_pred - denominator = w*K.sum(denominator,axes) + denominator = y_true + y_pred + denominator = w * K.sum(denominator, axes) denominator = K.sum(denominator) - gen_dice_coef = 2*numerator/denominator + gen_dice_coef = 2 * numerator / denominator return gen_dice_coef + def generalized_dice_loss(y_true, y_pred): return 1 - generalized_dice_coeff2(y_true, y_pred) -def soft_dice_loss(y_true, y_pred, epsilon=1e-6): + + +def soft_dice_loss(y_true, y_pred, epsilon=1e-6): ''' Soft dice loss calculation for arbitrary batch size, number of classes, and number of spatial dimensions. Assumes the `channels_last` format. @@ -151,16 +165,18 @@ def soft_dice_loss(y_true, y_pred, epsilon=1e-6): Adapted from https://github.com/Lasagne/Recipes/issues/99#issuecomment-347775022 ''' - + # skip the batch and class axis for calculating Dice score - axes = tuple(range(1, len(y_pred.shape)-1)) - + 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 + return 1.00 - K.mean(numerator / (denominator + epsilon)) # average over classes and batch + -def seg_metrics(y_true, y_pred, metric_name, metric_type='standard', drop_last = True, mean_per_class=False, verbose=False): +def seg_metrics(y_true, y_pred, metric_name, metric_type='standard', drop_last=True, mean_per_class=False, + verbose=False): """ Compute mean metrics of two segmentation masks, via Keras. @@ -193,13 +209,13 @@ def seg_metrics(y_true, y_pred, metric_name, metric_type='standard', drop_last = H = height, N = number of classes """ - + flag_soft = (metric_type == 'soft') flag_naive_mean = (metric_type == 'naive') - + # always assume one or more classes num_classes = K.shape(y_true)[-1] - + if not flag_soft: # get one-hot encoded masks from y_pred (true masks should already be one-hot) y_pred = K.one_hot(K.argmax(y_pred), num_classes) @@ -211,29 +227,29 @@ def seg_metrics(y_true, y_pred, metric_name, metric_type='standard', drop_last = y_pred = K.cast(y_pred, 'float32') # intersection and union shapes are batch_size * n_classes (values = area in pixels) - axes = (1,2) # W,H axes of each image + axes = (1, 2) # W,H axes of each image intersection = K.sum(K.abs(y_true * y_pred), axis=axes) mask_sum = K.sum(K.abs(y_true), axis=axes) + K.sum(K.abs(y_pred), axis=axes) - union = mask_sum - intersection # or, np.logical_or(y_pred, y_true) for one-hot + union = mask_sum - intersection # or, np.logical_or(y_pred, y_true) for one-hot smooth = .001 iou = (intersection + smooth) / (union + smooth) - dice = 2 * (intersection + smooth)/(mask_sum + smooth) + dice = 2 * (intersection + smooth) / (mask_sum + smooth) metric = {'iou': iou, 'dice': dice}[metric_name] # define mask to be 0 when no pixels are present in either y_true or y_pred, 1 otherwise - mask = K.cast(K.not_equal(union, 0), 'float32') - + mask = K.cast(K.not_equal(union, 0), 'float32') + if drop_last: - metric = metric[:,:-1] - mask = mask[:,:-1] - + metric = metric[:, :-1] + mask = mask[:, :-1] + if verbose: print('intersection, union') print(K.eval(intersection), K.eval(union)) - print(K.eval(intersection/union)) - + print(K.eval(intersection / union)) + # return mean metrics: remaining axes are (batch, classes) if flag_naive_mean: return K.mean(metric) @@ -243,13 +259,14 @@ def seg_metrics(y_true, y_pred, metric_name, metric_type='standard', drop_last = non_zero = tf.greater(class_count, 0) non_zero_sum = tf.boolean_mask(K.sum(metric * mask, axis=0), non_zero) non_zero_count = tf.boolean_mask(class_count, non_zero) - + if verbose: print('Counts of inputs with class present, metrics for non-absent classes') print(K.eval(class_count), K.eval(non_zero_sum / non_zero_count)) - + return K.mean(non_zero_sum / non_zero_count) + def mean_iou(y_true, y_pred, **kwargs): """ Compute mean Intersection over Union of two segmentation masks, via Keras. @@ -257,65 +274,69 @@ def mean_iou(y_true, y_pred, **kwargs): Calls metrics_k(y_true, y_pred, metric_name='iou'), see there for allowed kwargs. """ return seg_metrics(y_true, y_pred, metric_name='iou', **kwargs) + + def Mean_IOU(y_true, y_pred): nb_classes = K.int_shape(y_pred)[-1] iou = [] true_pixels = K.argmax(y_true, axis=-1) pred_pixels = K.argmax(y_pred, axis=-1) void_labels = K.equal(K.sum(y_true, axis=-1), 0) - for i in range(0, nb_classes): # exclude first label (background) and last label (void) - true_labels = K.equal(true_pixels, i)# & ~void_labels - pred_labels = K.equal(pred_pixels, i)# & ~void_labels + for i in range(0, nb_classes): # exclude first label (background) and last label (void) + true_labels = K.equal(true_pixels, i) # & ~void_labels + pred_labels = K.equal(pred_pixels, i) # & ~void_labels inter = tf.to_int32(true_labels & pred_labels) union = tf.to_int32(true_labels | pred_labels) - legal_batches = K.sum(tf.to_int32(true_labels), axis=1)>0 - ious = K.sum(inter, axis=1)/K.sum(union, axis=1) - iou.append(K.mean(tf.gather(ious, indices=tf.where(legal_batches)))) # returns average IoU of the same objects + legal_batches = K.sum(tf.to_int32(true_labels), axis=1) > 0 + ious = K.sum(inter, axis=1) / K.sum(union, axis=1) + iou.append(K.mean(tf.gather(ious, indices=tf.where(legal_batches)))) # returns average IoU of the same objects iou = tf.stack(iou) legal_labels = ~tf.debugging.is_nan(iou) iou = tf.gather(iou, indices=tf.where(legal_labels)) return K.mean(iou) + def iou_vahid(y_true, y_pred): - nb_classes = tf.shape(y_true)[-1]+tf.to_int32(1) + nb_classes = tf.shape(y_true)[-1] + tf.to_int32(1) true_pixels = K.argmax(y_true, axis=-1) pred_pixels = K.argmax(y_pred, axis=-1) iou = [] - + for i in tf.range(nb_classes): - tp=K.sum( tf.to_int32( K.equal(true_pixels, i) & K.equal(pred_pixels, i) ) ) - fp=K.sum( tf.to_int32( K.not_equal(true_pixels, i) & K.equal(pred_pixels, i) ) ) - fn=K.sum( tf.to_int32( K.equal(true_pixels, i) & K.not_equal(pred_pixels, i) ) ) - iouh=tp/(tp+fp+fn) + tp = K.sum(tf.to_int32(K.equal(true_pixels, i) & K.equal(pred_pixels, i))) + fp = K.sum(tf.to_int32(K.not_equal(true_pixels, i) & K.equal(pred_pixels, i))) + fn = K.sum(tf.to_int32(K.equal(true_pixels, i) & K.not_equal(pred_pixels, i))) + iouh = tp / (tp + fp + fn) iou.append(iouh) return K.mean(iou) - - -def IoU_metric(Yi,y_predi): - ## mean Intersection over Union - ## Mean IoU = TP/(FN + TP + FP) + + +def IoU_metric(Yi, y_predi): + # mean Intersection over Union + # Mean IoU = TP/(FN + TP + FP) y_predi = np.argmax(y_predi, axis=3) y_testi = np.argmax(Yi, axis=3) IoUs = [] Nclass = int(np.max(Yi)) + 1 for c in range(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) + 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) IoUs.append(IoU) - return K.cast( np.mean(IoUs) ,dtype='float32' ) + return K.cast(np.mean(IoUs), dtype='float32') def IoU_metric_keras(y_true, y_pred): - ## mean Intersection over Union - ## Mean IoU = TP/(FN + TP + FP) + # mean Intersection over Union + # Mean IoU = TP/(FN + TP + FP) init = tf.global_variables_initializer() sess = tf.Session() sess.run(init) - + return IoU_metric(y_true.eval(session=sess), y_pred.eval(session=sess)) + def jaccard_distance_loss(y_true, y_pred, smooth=100): """ Jaccard = (|X & Y|)/ (|X|+ |Y| - |X & Y|) @@ -334,5 +355,3 @@ def jaccard_distance_loss(y_true, y_pred, smooth=100): sum_ = K.sum(K.abs(y_true) + K.abs(y_pred), axis=-1) jac = (intersection + smooth) / (sum_ - intersection + smooth) return (1 - jac) * smooth - - diff --git a/eynollah/eynollah/train/models.py b/eynollah/eynollah/train/models.py index 7c806b4..7a1e246 100644 --- a/eynollah/eynollah/train/models.py +++ b/eynollah/eynollah/train/models.py @@ -3,19 +3,20 @@ from keras import layers from keras.regularizers import l2 -resnet50_Weights_path='./pretrained_model/resnet50_weights_tf_dim_ordering_tf_kernels_notop.h5' -IMAGE_ORDERING ='channels_last' -MERGE_AXIS=-1 +resnet50_Weights_path = './pretrained_model/resnet50_weights_tf_dim_ordering_tf_kernels_notop.h5' +IMAGE_ORDERING = 'channels_last' +MERGE_AXIS = -1 -def one_side_pad( x ): +def one_side_pad(x): x = ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING)(x) if IMAGE_ORDERING == 'channels_first': - x = Lambda(lambda x : x[: , : , :-1 , :-1 ] )(x) + x = Lambda(lambda x: x[:, :, :-1, :-1])(x) elif IMAGE_ORDERING == 'channels_last': - x = Lambda(lambda x : x[: , :-1 , :-1 , : ] )(x) + x = Lambda(lambda x: x[:, :-1, :-1, :])(x) return x + def identity_block(input_tensor, kernel_size, filters, stage, block): """The identity block is the block that has no conv layer at shortcut. # Arguments @@ -28,7 +29,7 @@ def identity_block(input_tensor, kernel_size, filters, stage, block): Output tensor for the block. """ filters1, filters2, filters3 = filters - + if IMAGE_ORDERING == 'channels_last': bn_axis = 3 else: @@ -37,16 +38,16 @@ def identity_block(input_tensor, kernel_size, filters, stage, block): conv_name_base = 'res' + str(stage) + block + '_branch' bn_name_base = 'bn' + str(stage) + block + '_branch' - x = Conv2D(filters1, (1, 1) , data_format=IMAGE_ORDERING , name=conv_name_base + '2a')(input_tensor) + x = Conv2D(filters1, (1, 1), data_format=IMAGE_ORDERING, name=conv_name_base + '2a')(input_tensor) x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2a')(x) x = Activation('relu')(x) - x = Conv2D(filters2, kernel_size , data_format=IMAGE_ORDERING , + x = Conv2D(filters2, kernel_size, data_format=IMAGE_ORDERING, padding='same', name=conv_name_base + '2b')(x) x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2b')(x) x = Activation('relu')(x) - x = Conv2D(filters3 , (1, 1), data_format=IMAGE_ORDERING , name=conv_name_base + '2c')(x) + x = Conv2D(filters3, (1, 1), data_format=IMAGE_ORDERING, name=conv_name_base + '2c')(x) x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2c')(x) x = layers.add([x, input_tensor]) @@ -68,7 +69,7 @@ def conv_block(input_tensor, kernel_size, filters, stage, block, strides=(2, 2)) And the shortcut should have strides=(2,2) as well """ filters1, filters2, filters3 = filters - + if IMAGE_ORDERING == 'channels_last': bn_axis = 3 else: @@ -77,20 +78,20 @@ def conv_block(input_tensor, kernel_size, filters, stage, block, strides=(2, 2)) conv_name_base = 'res' + str(stage) + block + '_branch' bn_name_base = 'bn' + str(stage) + block + '_branch' - x = Conv2D(filters1, (1, 1) , data_format=IMAGE_ORDERING , strides=strides, + x = Conv2D(filters1, (1, 1), data_format=IMAGE_ORDERING, strides=strides, name=conv_name_base + '2a')(input_tensor) x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2a')(x) x = Activation('relu')(x) - x = Conv2D(filters2, kernel_size , data_format=IMAGE_ORDERING , padding='same', + x = Conv2D(filters2, kernel_size, data_format=IMAGE_ORDERING, padding='same', name=conv_name_base + '2b')(x) x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2b')(x) x = Activation('relu')(x) - x = Conv2D(filters3, (1, 1) , data_format=IMAGE_ORDERING , name=conv_name_base + '2c')(x) + x = Conv2D(filters3, (1, 1), data_format=IMAGE_ORDERING, name=conv_name_base + '2c')(x) x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2c')(x) - shortcut = Conv2D(filters3, (1, 1) , data_format=IMAGE_ORDERING , strides=strides, + shortcut = Conv2D(filters3, (1, 1), data_format=IMAGE_ORDERING, strides=strides, name=conv_name_base + '1')(input_tensor) shortcut = BatchNormalization(axis=bn_axis, name=bn_name_base + '1')(shortcut) @@ -99,12 +100,11 @@ def conv_block(input_tensor, kernel_size, filters, stage, block, strides=(2, 2)) return x -def resnet50_unet_light(n_classes,input_height=224,input_width=224,weight_decay=1e-6,pretraining=False): - assert input_height%32 == 0 - assert input_width%32 == 0 +def resnet50_unet_light(n_classes, input_height=224, input_width=224, weight_decay=1e-6, pretraining=False): + assert input_height % 32 == 0 + assert input_width % 32 == 0 - - img_input = Input(shape=(input_height,input_width , 3 )) + img_input = Input(shape=(input_height, input_width, 3)) if IMAGE_ORDERING == 'channels_last': bn_axis = 3 @@ -112,25 +112,24 @@ def resnet50_unet_light(n_classes,input_height=224,input_width=224,weight_decay= bn_axis = 1 x = ZeroPadding2D((3, 3), data_format=IMAGE_ORDERING)(img_input) - x = Conv2D(64, (7, 7), data_format=IMAGE_ORDERING, strides=(2, 2),kernel_regularizer=l2(weight_decay), name='conv1')(x) + x = Conv2D(64, (7, 7), data_format=IMAGE_ORDERING, strides=(2, 2), kernel_regularizer=l2(weight_decay), + name='conv1')(x) f1 = x x = BatchNormalization(axis=bn_axis, name='bn_conv1')(x) x = Activation('relu')(x) - x = MaxPooling2D((3, 3) , data_format=IMAGE_ORDERING , strides=(2, 2))(x) - + x = MaxPooling2D((3, 3), data_format=IMAGE_ORDERING, strides=(2, 2))(x) x = conv_block(x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1)) x = identity_block(x, 3, [64, 64, 256], stage=2, block='b') x = identity_block(x, 3, [64, 64, 256], stage=2, block='c') - f2 = one_side_pad(x ) - + f2 = one_side_pad(x) x = conv_block(x, 3, [128, 128, 512], stage=3, block='a') x = identity_block(x, 3, [128, 128, 512], stage=3, block='b') x = identity_block(x, 3, [128, 128, 512], stage=3, block='c') x = identity_block(x, 3, [128, 128, 512], stage=3, block='d') - f3 = x + f3 = x x = conv_block(x, 3, [256, 256, 1024], stage=4, block='a') x = identity_block(x, 3, [256, 256, 1024], stage=4, block='b') @@ -138,85 +137,72 @@ def resnet50_unet_light(n_classes,input_height=224,input_width=224,weight_decay= x = identity_block(x, 3, [256, 256, 1024], stage=4, block='d') x = identity_block(x, 3, [256, 256, 1024], stage=4, block='e') x = identity_block(x, 3, [256, 256, 1024], stage=4, block='f') - f4 = x + f4 = x x = conv_block(x, 3, [512, 512, 2048], stage=5, block='a') x = identity_block(x, 3, [512, 512, 2048], stage=5, block='b') x = identity_block(x, 3, [512, 512, 2048], stage=5, block='c') - f5 = x - + f5 = x if pretraining: - model=Model( img_input , x ).load_weights(resnet50_Weights_path) + model = Model(img_input, x).load_weights(resnet50_Weights_path) - - v512_2048 = Conv2D( 512 , (1, 1) , padding='same', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay) )( f5 ) - v512_2048 = ( BatchNormalization(axis=bn_axis))(v512_2048) + v512_2048 = Conv2D(512, (1, 1), padding='same', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(f5) + v512_2048 = (BatchNormalization(axis=bn_axis))(v512_2048) v512_2048 = Activation('relu')(v512_2048) - - - v512_1024=Conv2D( 512 , (1, 1) , padding='same', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay) )( f4 ) - v512_1024 = ( BatchNormalization(axis=bn_axis))(v512_1024) + v512_1024 = Conv2D(512, (1, 1), padding='same', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(f4) + v512_1024 = (BatchNormalization(axis=bn_axis))(v512_1024) v512_1024 = Activation('relu')(v512_1024) - - - o = ( UpSampling2D( (2,2), data_format=IMAGE_ORDERING))(v512_2048) - o = ( concatenate([ o ,v512_1024],axis=MERGE_AXIS ) ) - o = ( ZeroPadding2D( (1,1), data_format=IMAGE_ORDERING))(o) - o = ( Conv2D(512, (3, 3), padding='valid', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay)))(o) - o = ( BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = ( UpSampling2D( (2,2), data_format=IMAGE_ORDERING))(o) - o = ( concatenate([ o ,f3],axis=MERGE_AXIS ) ) - o = ( ZeroPadding2D( (1,1), data_format=IMAGE_ORDERING))(o) - o = ( Conv2D( 256, (3, 3), padding='valid', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay)))(o) - o = ( BatchNormalization(axis=bn_axis))(o) + o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(v512_2048) + o = (concatenate([o, v512_1024], axis=MERGE_AXIS)) + o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) + o = (Conv2D(512, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) + o = (BatchNormalization(axis=bn_axis))(o) o = Activation('relu')(o) - - o = ( UpSampling2D( (2,2), data_format=IMAGE_ORDERING))(o) - o = ( concatenate([o,f2],axis=MERGE_AXIS ) ) - o = ( ZeroPadding2D((1,1) , data_format=IMAGE_ORDERING))(o) - o = ( Conv2D( 128 , (3, 3), padding='valid' , data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay) ) )(o) - o = ( BatchNormalization(axis=bn_axis))(o) + o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) + o = (concatenate([o, f3], axis=MERGE_AXIS)) + o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) + o = (Conv2D(256, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) + o = (BatchNormalization(axis=bn_axis))(o) o = Activation('relu')(o) - - - o = ( UpSampling2D( (2,2), data_format=IMAGE_ORDERING))(o) - o = ( concatenate([o,f1],axis=MERGE_AXIS ) ) - o = ( ZeroPadding2D((1,1) , data_format=IMAGE_ORDERING ))(o) - o = ( Conv2D( 64 , (3, 3), padding='valid' , data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay) ))(o) - o = ( BatchNormalization(axis=bn_axis))(o) + o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) + o = (concatenate([o, f2], axis=MERGE_AXIS)) + o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) + o = (Conv2D(128, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) + o = (BatchNormalization(axis=bn_axis))(o) o = Activation('relu')(o) + o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) + o = (concatenate([o, f1], axis=MERGE_AXIS)) + o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) + o = (Conv2D(64, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) + o = (BatchNormalization(axis=bn_axis))(o) + o = Activation('relu')(o) - o = ( UpSampling2D( (2,2), data_format=IMAGE_ORDERING))(o) - o = ( concatenate([o,img_input],axis=MERGE_AXIS ) ) - o = ( ZeroPadding2D((1,1) , data_format=IMAGE_ORDERING ))(o) - o = ( Conv2D( 32 , (3, 3), padding='valid' , data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay) ))(o) - o = ( BatchNormalization(axis=bn_axis))(o) + o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) + o = (concatenate([o, img_input], axis=MERGE_AXIS)) + o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) + o = (Conv2D(32, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) + o = (BatchNormalization(axis=bn_axis))(o) o = Activation('relu')(o) - - - o = Conv2D( n_classes , (1, 1) , padding='same', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay) )( o ) - o = ( BatchNormalization(axis=bn_axis))(o) + o = Conv2D(n_classes, (1, 1), padding='same', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(o) + o = (BatchNormalization(axis=bn_axis))(o) o = (Activation('softmax'))(o) - - model = Model( img_input , o ) + model = Model(img_input, o) return model -def resnet50_unet(n_classes,input_height=224,input_width=224,weight_decay=1e-6,pretraining=False): - assert input_height%32 == 0 - assert input_width%32 == 0 - - img_input = Input(shape=(input_height,input_width , 3 )) +def resnet50_unet(n_classes, input_height=224, input_width=224, weight_decay=1e-6, pretraining=False): + assert input_height % 32 == 0 + assert input_width % 32 == 0 + + img_input = Input(shape=(input_height, input_width, 3)) if IMAGE_ORDERING == 'channels_last': bn_axis = 3 @@ -224,25 +210,24 @@ def resnet50_unet(n_classes,input_height=224,input_width=224,weight_decay=1e-6,p bn_axis = 1 x = ZeroPadding2D((3, 3), data_format=IMAGE_ORDERING)(img_input) - x = Conv2D(64, (7, 7), data_format=IMAGE_ORDERING, strides=(2, 2),kernel_regularizer=l2(weight_decay), name='conv1')(x) + x = Conv2D(64, (7, 7), data_format=IMAGE_ORDERING, strides=(2, 2), kernel_regularizer=l2(weight_decay), + name='conv1')(x) f1 = x x = BatchNormalization(axis=bn_axis, name='bn_conv1')(x) x = Activation('relu')(x) - x = MaxPooling2D((3, 3) , data_format=IMAGE_ORDERING , strides=(2, 2))(x) - + x = MaxPooling2D((3, 3), data_format=IMAGE_ORDERING, strides=(2, 2))(x) x = conv_block(x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1)) x = identity_block(x, 3, [64, 64, 256], stage=2, block='b') x = identity_block(x, 3, [64, 64, 256], stage=2, block='c') - f2 = one_side_pad(x ) - + f2 = one_side_pad(x) x = conv_block(x, 3, [128, 128, 512], stage=3, block='a') x = identity_block(x, 3, [128, 128, 512], stage=3, block='b') x = identity_block(x, 3, [128, 128, 512], stage=3, block='c') x = identity_block(x, 3, [128, 128, 512], stage=3, block='d') - f3 = x + f3 = x x = conv_block(x, 3, [256, 256, 1024], stage=4, block='a') x = identity_block(x, 3, [256, 256, 1024], stage=4, block='b') @@ -250,68 +235,60 @@ def resnet50_unet(n_classes,input_height=224,input_width=224,weight_decay=1e-6,p x = identity_block(x, 3, [256, 256, 1024], stage=4, block='d') x = identity_block(x, 3, [256, 256, 1024], stage=4, block='e') x = identity_block(x, 3, [256, 256, 1024], stage=4, block='f') - f4 = x + f4 = x x = conv_block(x, 3, [512, 512, 2048], stage=5, block='a') x = identity_block(x, 3, [512, 512, 2048], stage=5, block='b') x = identity_block(x, 3, [512, 512, 2048], stage=5, block='c') - f5 = x + f5 = x if pretraining: - Model( img_input , x ).load_weights(resnet50_Weights_path) + Model(img_input, x).load_weights(resnet50_Weights_path) - v1024_2048 = Conv2D( 1024 , (1, 1) , padding='same', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay) )( f5 ) - v1024_2048 = ( BatchNormalization(axis=bn_axis))(v1024_2048) + v1024_2048 = Conv2D(1024, (1, 1), padding='same', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))( + f5) + v1024_2048 = (BatchNormalization(axis=bn_axis))(v1024_2048) v1024_2048 = Activation('relu')(v1024_2048) - - - o = ( UpSampling2D( (2,2), data_format=IMAGE_ORDERING))(v1024_2048) - o = ( concatenate([ o ,f4],axis=MERGE_AXIS ) ) - o = ( ZeroPadding2D( (1,1), data_format=IMAGE_ORDERING))(o) - o = ( Conv2D(512, (3, 3), padding='valid', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay)))(o) - o = ( BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = ( UpSampling2D( (2,2), data_format=IMAGE_ORDERING))(o) - o = ( concatenate([ o ,f3],axis=MERGE_AXIS ) ) - o = ( ZeroPadding2D( (1,1), data_format=IMAGE_ORDERING))(o) - o = ( Conv2D( 256, (3, 3), padding='valid', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay)))(o) - o = ( BatchNormalization(axis=bn_axis))(o) + o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(v1024_2048) + o = (concatenate([o, f4], axis=MERGE_AXIS)) + o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) + o = (Conv2D(512, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) + o = (BatchNormalization(axis=bn_axis))(o) o = Activation('relu')(o) - - o = ( UpSampling2D( (2,2), data_format=IMAGE_ORDERING))(o) - o = ( concatenate([o,f2],axis=MERGE_AXIS ) ) - o = ( ZeroPadding2D((1,1) , data_format=IMAGE_ORDERING))(o) - o = ( Conv2D( 128 , (3, 3), padding='valid' , data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay) ) )(o) - o = ( BatchNormalization(axis=bn_axis))(o) + o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) + o = (concatenate([o, f3], axis=MERGE_AXIS)) + o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) + o = (Conv2D(256, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) + o = (BatchNormalization(axis=bn_axis))(o) o = Activation('relu')(o) - - o = ( UpSampling2D( (2,2), data_format=IMAGE_ORDERING))(o) - o = ( concatenate([o,f1],axis=MERGE_AXIS ) ) - o = ( ZeroPadding2D((1,1) , data_format=IMAGE_ORDERING ))(o) - o = ( Conv2D( 64 , (3, 3), padding='valid' , data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay) ))(o) - o = ( BatchNormalization(axis=bn_axis))(o) + o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) + o = (concatenate([o, f2], axis=MERGE_AXIS)) + o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) + o = (Conv2D(128, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) + o = (BatchNormalization(axis=bn_axis))(o) o = Activation('relu')(o) + o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) + o = (concatenate([o, f1], axis=MERGE_AXIS)) + o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) + o = (Conv2D(64, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) + o = (BatchNormalization(axis=bn_axis))(o) + o = Activation('relu')(o) - o = ( UpSampling2D( (2,2), data_format=IMAGE_ORDERING))(o) - o = ( concatenate([o,img_input],axis=MERGE_AXIS ) ) - o = ( ZeroPadding2D((1,1) , data_format=IMAGE_ORDERING ))(o) - o = ( Conv2D( 32 , (3, 3), padding='valid' , data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay) ))(o) - o = ( BatchNormalization(axis=bn_axis))(o) + o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) + o = (concatenate([o, img_input], axis=MERGE_AXIS)) + o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) + o = (Conv2D(32, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) + o = (BatchNormalization(axis=bn_axis))(o) o = Activation('relu')(o) - - - o = Conv2D( n_classes , (1, 1) , padding='same', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay) )( o ) - o = ( BatchNormalization(axis=bn_axis))(o) + + o = Conv2D(n_classes, (1, 1), padding='same', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(o) + o = (BatchNormalization(axis=bn_axis))(o) o = (Activation('softmax'))(o) - - model = Model( img_input , o ) - - + model = Model(img_input, o) return model diff --git a/eynollah/eynollah/train/pagexml2img.py b/eynollah/eynollah/train/pagexml2img.py index c489315..8570f4f 100644 --- a/eynollah/eynollah/train/pagexml2img.py +++ b/eynollah/eynollah/train/pagexml2img.py @@ -1,6 +1,6 @@ #! /usr/bin/env python3 -__version__= '1.0' +__version__ = '1.0' import argparse import sys @@ -14,235 +14,260 @@ with warnings.catch_warnings(): warnings.simplefilter("ignore") -__doc__=\ -""" +__doc__ = \ + """ tool to extract 2d or 3d RGB images from page xml data. In former case output will be 1 2D image array which each class has filled with a pixel value. In the case of 3D RGB image each class will be defined with a RGB value and beside images a text file of classes also will be produced. This classes.txt file is required for dhsegment tool. """ + class pagexml2img: - def __init__(self,dir_in, out_dir,output_type): - self.dir=dir_in - self.output_dir=out_dir - self.output_type=output_type + def __init__(self, dir_in, out_dir, output_type): + self.dir = dir_in + self.output_dir = out_dir + self.output_type = output_type def get_content_of_dir(self): """ Listing all ground truth page xml files. All files are needed to have xml format. """ - - gt_all=os.listdir(self.dir) - self.gt_list=[file for file in gt_all if file.split('.')[ len(file.split('.'))-1 ]=='xml' ] + gt_all = os.listdir(self.dir) + self.gt_list = [file for file in gt_all if file.split('.')[len(file.split('.')) - 1] == 'xml'] def get_images_of_ground_truth(self): """ Reading the page xml files and write the ground truth images into given output directory. """ - if self.output_type=='3d' or self.output_type=='3D': - classes=np.array([ [0,0,0, 1, 0, 0, 0, 0], - [255,0,0, 0, 1, 0, 0, 0], - [0,255,0, 0, 0, 1, 0, 0], - [0,0,255, 0, 0, 0, 1, 0], - [0,255,255, 0, 0, 0, 0, 1] ]) - - - + if self.output_type == '3d' or self.output_type == '3D': + classes = np.array([[0, 0, 0, 1, 0, 0, 0, 0], + [255, 0, 0, 0, 1, 0, 0, 0], + [0, 255, 0, 0, 0, 1, 0, 0], + [0, 0, 255, 0, 0, 0, 1, 0], + [0, 255, 255, 0, 0, 0, 0, 1]]) for index in tqdm(range(len(self.gt_list))): try: - tree1 = ET.parse(self.dir+'/'+self.gt_list[index]) - root1=tree1.getroot() - alltags=[elem.tag for elem in root1.iter()] - link=alltags[0].split('}')[0]+'}' - - region_tags=np.unique([x for x in alltags if x.endswith('Region')]) - - for jj in root1.iter(link+'Page'): - y_len=int(jj.attrib['imageHeight']) - x_len=int(jj.attrib['imageWidth']) - - co_text=[] - co_sep=[] - co_img=[] - co_table=[] + tree1 = ET.parse(self.dir + '/' + self.gt_list[index]) + root1 = tree1.getroot() + alltags = [elem.tag for elem in root1.iter()] + link = alltags[0].split('}')[0] + '}' + + region_tags = np.unique([x for x in alltags if x.endswith('Region')]) + + for jj in root1.iter(link + 'Page'): + y_len = int(jj.attrib['imageHeight']) + x_len = int(jj.attrib['imageWidth']) + + co_text = [] + co_sep = [] + co_img = [] + co_table = [] for tag in region_tags: - if tag.endswith('}TextRegion') or tag.endswith('}Textregion') or tag.endswith('}textRegion') or tag.endswith('}textregion'): - + if tag.endswith('}TextRegion') or tag.endswith('}Textregion') or tag.endswith( + '}textRegion') or tag.endswith('}textregion'): + for nn in root1.iter(tag): - for co_it in nn.iter(link+'Coords'): - if bool(co_it.attrib)==False: - c_t_in=[] - for ll in nn.iter(link+'Point'): - c_t_in.append([ int(np.float(ll.attrib['x'])) , int(np.float(ll.attrib['y'])) ]) + for co_it in nn.iter(link + 'Coords'): + if bool(co_it.attrib) == False: + c_t_in = [] + for ll in nn.iter(link + 'Point'): + c_t_in.append( + [int(np.float(ll.attrib['x'])), int(np.float(ll.attrib['y']))]) co_text.append(np.array(c_t_in)) print(co_text) - elif bool(co_it.attrib)==True and 'points' in co_it.attrib.keys(): - p_h=co_it.attrib['points'].split(' ') - co_text.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) + elif bool(co_it.attrib) == True and 'points' in co_it.attrib.keys(): + p_h = co_it.attrib['points'].split(' ') + co_text.append( + np.array([[int(x.split(',')[0]), int(x.split(',')[1])] for x in p_h])) - - elif tag.endswith('}ImageRegion') or tag.endswith('}Imageregion') or tag.endswith('}imageRegion') or tag.endswith('}imageregion'): + elif tag.endswith('}ImageRegion') or tag.endswith('}Imageregion') or tag.endswith( + '}imageRegion') or tag.endswith('}imageregion'): for nn in root1.iter(tag): - for co_it in nn.iter(link+'Coords'): - if bool(co_it.attrib)==False: - c_i_in=[] - for ll in nn.iter(link+'Point'): - c_i_in.append([ int(np.float(ll.attrib['x'])) , int(np.float(ll.attrib['y'])) ]) + for co_it in nn.iter(link + 'Coords'): + if bool(co_it.attrib) == False: + c_i_in = [] + for ll in nn.iter(link + 'Point'): + c_i_in.append( + [int(np.float(ll.attrib['x'])), int(np.float(ll.attrib['y']))]) co_img.append(np.array(c_i_in)) - elif bool(co_it.attrib)==True and 'points' in co_it.attrib.keys(): - p_h=co_it.attrib['points'].split(' ') - co_img.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - - elif tag.endswith('}SeparatorRegion') or tag.endswith('}Separatorregion') or tag.endswith('}separatorRegion') or tag.endswith('}separatorregion'): + elif bool(co_it.attrib) == True and 'points' in co_it.attrib.keys(): + p_h = co_it.attrib['points'].split(' ') + co_img.append( + np.array([[int(x.split(',')[0]), int(x.split(',')[1])] for x in p_h])) + + elif tag.endswith('}SeparatorRegion') or tag.endswith('}Separatorregion') or tag.endswith( + '}separatorRegion') or tag.endswith('}separatorregion'): for nn in root1.iter(tag): - for co_it in nn.iter(link+'Coords'): - if bool(co_it.attrib)==False: - c_s_in=[] - for ll in nn.iter(link+'Point'): - c_s_in.append([ int(np.float(ll.attrib['x'])) , int(np.float(ll.attrib['y'])) ]) + for co_it in nn.iter(link + 'Coords'): + if bool(co_it.attrib) == False: + c_s_in = [] + for ll in nn.iter(link + 'Point'): + c_s_in.append( + [int(np.float(ll.attrib['x'])), int(np.float(ll.attrib['y']))]) co_sep.append(np.array(c_s_in)) - - elif bool(co_it.attrib)==True and 'points' in co_it.attrib.keys(): - p_h=co_it.attrib['points'].split(' ') - co_sep.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - - elif tag.endswith('}TableRegion') or tag.endswith('}tableRegion') or tag.endswith('}Tableregion') or tag.endswith('}tableregion'): + + elif bool(co_it.attrib) == True and 'points' in co_it.attrib.keys(): + p_h = co_it.attrib['points'].split(' ') + co_sep.append( + np.array([[int(x.split(',')[0]), int(x.split(',')[1])] for x in p_h])) + + elif tag.endswith('}TableRegion') or tag.endswith('}tableRegion') or tag.endswith( + '}Tableregion') or tag.endswith('}tableregion'): for nn in root1.iter(tag): - for co_it in nn.iter(link+'Coords'): - if bool(co_it.attrib)==False: - c_ta_in=[] - for ll in nn.iter(link+'Point'): - c_ta_in.append([ int(np.float(ll.attrib['x'])) , int(np.float(ll.attrib['y'])) ]) + for co_it in nn.iter(link + 'Coords'): + if bool(co_it.attrib) == False: + c_ta_in = [] + for ll in nn.iter(link + 'Point'): + c_ta_in.append( + [int(np.float(ll.attrib['x'])), int(np.float(ll.attrib['y']))]) co_table.append(np.array(c_ta_in)) - - elif bool(co_it.attrib)==True and 'points' in co_it.attrib.keys(): - p_h=co_it.attrib['points'].split(' ') - co_table.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) + + elif bool(co_it.attrib) == True and 'points' in co_it.attrib.keys(): + p_h = co_it.attrib['points'].split(' ') + co_table.append( + np.array([[int(x.split(',')[0]), int(x.split(',')[1])] for x in p_h])) else: pass - - img = np.zeros( (y_len,x_len,3) ) - img_poly=cv2.fillPoly(img, pts =co_text, color=(255,0,0)) - img_poly=cv2.fillPoly(img, pts =co_img, color=(0,255,0)) - img_poly=cv2.fillPoly(img, pts =co_sep, color=(0,0,255)) - img_poly=cv2.fillPoly(img, pts =co_table, color=(0,255,255)) - - try: - cv2.imwrite(self.output_dir+'/'+self.gt_list[index].split('-')[1].split('.')[0]+'.png',img_poly ) + + img = np.zeros((y_len, x_len, 3)) + img_poly = cv2.fillPoly(img, pts=co_text, color=(255, 0, 0)) + img_poly = cv2.fillPoly(img, pts=co_img, color=(0, 255, 0)) + img_poly = cv2.fillPoly(img, pts=co_sep, color=(0, 0, 255)) + img_poly = cv2.fillPoly(img, pts=co_table, color=(0, 255, 255)) + + try: + cv2.imwrite(self.output_dir + '/' + self.gt_list[index].split('-')[1].split('.')[0] + '.png', + img_poly) except: - cv2.imwrite(self.output_dir+'/'+self.gt_list[index].split('.')[0]+'.png',img_poly ) + cv2.imwrite(self.output_dir + '/' + self.gt_list[index].split('.')[0] + '.png', img_poly) except: pass - np.savetxt(self.output_dir+'/../classes.txt',classes) - - if self.output_type=='2d' or self.output_type=='2D': + np.savetxt(self.output_dir + '/../classes.txt', classes) + + if self.output_type == '2d' or self.output_type == '2D': for index in tqdm(range(len(self.gt_list))): try: - tree1 = ET.parse(self.dir+'/'+self.gt_list[index]) - root1=tree1.getroot() - alltags=[elem.tag for elem in root1.iter()] - link=alltags[0].split('}')[0]+'}' - - region_tags=np.unique([x for x in alltags if x.endswith('Region')]) - - for jj in root1.iter(link+'Page'): - y_len=int(jj.attrib['imageHeight']) - x_len=int(jj.attrib['imageWidth']) - - co_text=[] - co_sep=[] - co_img=[] - co_table=[] + tree1 = ET.parse(self.dir + '/' + self.gt_list[index]) + root1 = tree1.getroot() + alltags = [elem.tag for elem in root1.iter()] + link = alltags[0].split('}')[0] + '}' + + region_tags = np.unique([x for x in alltags if x.endswith('Region')]) + + for jj in root1.iter(link + 'Page'): + y_len = int(jj.attrib['imageHeight']) + x_len = int(jj.attrib['imageWidth']) + + co_text = [] + co_sep = [] + co_img = [] + co_table = [] for tag in region_tags: - if tag.endswith('}TextRegion') or tag.endswith('}Textregion') or tag.endswith('}textRegion') or tag.endswith('}textregion'): - + if tag.endswith('}TextRegion') or tag.endswith('}Textregion') or tag.endswith( + '}textRegion') or tag.endswith('}textregion'): + for nn in root1.iter(tag): - for co_it in nn.iter(link+'Coords'): - if bool(co_it.attrib)==False: - c_t_in=[] - for ll in nn.iter(link+'Point'): - c_t_in.append([ int(np.float(ll.attrib['x'])) , int(np.float(ll.attrib['y'])) ]) + for co_it in nn.iter(link + 'Coords'): + if bool(co_it.attrib) == False: + c_t_in = [] + for ll in nn.iter(link + 'Point'): + c_t_in.append( + [int(np.float(ll.attrib['x'])), int(np.float(ll.attrib['y']))]) co_text.append(np.array(c_t_in)) print(co_text) - elif bool(co_it.attrib)==True and 'points' in co_it.attrib.keys(): - p_h=co_it.attrib['points'].split(' ') - co_text.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) + elif bool(co_it.attrib) == True and 'points' in co_it.attrib.keys(): + p_h = co_it.attrib['points'].split(' ') + co_text.append( + np.array([[int(x.split(',')[0]), int(x.split(',')[1])] for x in p_h])) - - elif tag.endswith('}ImageRegion') or tag.endswith('}Imageregion') or tag.endswith('}imageRegion') or tag.endswith('}imageregion'): + elif tag.endswith('}ImageRegion') or tag.endswith('}Imageregion') or tag.endswith( + '}imageRegion') or tag.endswith('}imageregion'): for nn in root1.iter(tag): - for co_it in nn.iter(link+'Coords'): - if bool(co_it.attrib)==False: - c_i_in=[] - for ll in nn.iter(link+'Point'): - c_i_in.append([ int(np.float(ll.attrib['x'])) , int(np.float(ll.attrib['y'])) ]) + for co_it in nn.iter(link + 'Coords'): + if bool(co_it.attrib) == False: + c_i_in = [] + for ll in nn.iter(link + 'Point'): + c_i_in.append( + [int(np.float(ll.attrib['x'])), int(np.float(ll.attrib['y']))]) co_img.append(np.array(c_i_in)) - elif bool(co_it.attrib)==True and 'points' in co_it.attrib.keys(): - p_h=co_it.attrib['points'].split(' ') - co_img.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - - elif tag.endswith('}SeparatorRegion') or tag.endswith('}Separatorregion') or tag.endswith('}separatorRegion') or tag.endswith('}separatorregion'): + elif bool(co_it.attrib) == True and 'points' in co_it.attrib.keys(): + p_h = co_it.attrib['points'].split(' ') + co_img.append( + np.array([[int(x.split(',')[0]), int(x.split(',')[1])] for x in p_h])) + + elif tag.endswith('}SeparatorRegion') or tag.endswith('}Separatorregion') or tag.endswith( + '}separatorRegion') or tag.endswith('}separatorregion'): for nn in root1.iter(tag): - for co_it in nn.iter(link+'Coords'): - if bool(co_it.attrib)==False: - c_s_in=[] - for ll in nn.iter(link+'Point'): - c_s_in.append([ int(np.float(ll.attrib['x'])) , int(np.float(ll.attrib['y'])) ]) + for co_it in nn.iter(link + 'Coords'): + if bool(co_it.attrib) == False: + c_s_in = [] + for ll in nn.iter(link + 'Point'): + c_s_in.append( + [int(np.float(ll.attrib['x'])), int(np.float(ll.attrib['y']))]) co_sep.append(np.array(c_s_in)) - - elif bool(co_it.attrib)==True and 'points' in co_it.attrib.keys(): - p_h=co_it.attrib['points'].split(' ') - co_sep.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - - elif tag.endswith('}TableRegion') or tag.endswith('}tableRegion') or tag.endswith('}Tableregion') or tag.endswith('}tableregion'): + + elif bool(co_it.attrib) == True and 'points' in co_it.attrib.keys(): + p_h = co_it.attrib['points'].split(' ') + co_sep.append( + np.array([[int(x.split(',')[0]), int(x.split(',')[1])] for x in p_h])) + + elif tag.endswith('}TableRegion') or tag.endswith('}tableRegion') or tag.endswith( + '}Tableregion') or tag.endswith('}tableregion'): for nn in root1.iter(tag): - for co_it in nn.iter(link+'Coords'): - if bool(co_it.attrib)==False: - c_ta_in=[] - for ll in nn.iter(link+'Point'): - c_ta_in.append([ int(np.float(ll.attrib['x'])) , int(np.float(ll.attrib['y'])) ]) + for co_it in nn.iter(link + 'Coords'): + if bool(co_it.attrib) == False: + c_ta_in = [] + for ll in nn.iter(link + 'Point'): + c_ta_in.append( + [int(np.float(ll.attrib['x'])), int(np.float(ll.attrib['y']))]) co_table.append(np.array(c_ta_in)) - - elif bool(co_it.attrib)==True and 'points' in co_it.attrib.keys(): - p_h=co_it.attrib['points'].split(' ') - co_table.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) + + elif bool(co_it.attrib) == True and 'points' in co_it.attrib.keys(): + p_h = co_it.attrib['points'].split(' ') + co_table.append( + np.array([[int(x.split(',')[0]), int(x.split(',')[1])] for x in p_h])) else: pass - - img = np.zeros( (y_len,x_len) ) - img_poly=cv2.fillPoly(img, pts =co_text, color=(1,1,1)) - img_poly=cv2.fillPoly(img, pts =co_img, color=(2,2,2)) - img_poly=cv2.fillPoly(img, pts =co_sep, color=(3,3,3)) - img_poly=cv2.fillPoly(img, pts =co_table, color=(4,4,4)) - try: - cv2.imwrite(self.output_dir+'/'+self.gt_list[index].split('-')[1].split('.')[0]+'.png',img_poly ) + + img = np.zeros((y_len, x_len)) + img_poly = cv2.fillPoly(img, pts=co_text, color=(1, 1, 1)) + img_poly = cv2.fillPoly(img, pts=co_img, color=(2, 2, 2)) + img_poly = cv2.fillPoly(img, pts=co_sep, color=(3, 3, 3)) + img_poly = cv2.fillPoly(img, pts=co_table, color=(4, 4, 4)) + try: + cv2.imwrite(self.output_dir + '/' + self.gt_list[index].split('-')[1].split('.')[0] + '.png', + img_poly) except: - cv2.imwrite(self.output_dir+'/'+self.gt_list[index].split('.')[0]+'.png',img_poly ) + cv2.imwrite(self.output_dir + '/' + self.gt_list[index].split('.')[0] + '.png', img_poly) except: pass + def run(self): self.get_content_of_dir() self.get_images_of_ground_truth() + + def main(): - parser=argparse.ArgumentParser() - - parser.add_argument('-dir_in','--dir_in', dest='inp1', default=None, help='directory of page-xml files') - parser.add_argument('-dir_out','--dir_out', dest='inp2', default=None, help='directory where ground truth images would be written') - parser.add_argument('-type','--type', dest='inp3', default=None, help='this defines how output should be. A 2d image array or a 3d image array encoded with RGB color. Just pass 2d or 3d. The file will be saved one directory up. 2D image array is 3d but only information of one channel would be enough since all channels have the same values.') - options=parser.parse_args() - - possibles=globals() + parser = argparse.ArgumentParser() + + parser.add_argument('-dir_in', '--dir_in', dest='inp1', default=None, help='directory of page-xml files') + parser.add_argument('-dir_out', '--dir_out', dest='inp2', default=None, + help='directory where ground truth images would be written') + parser.add_argument('-type', '--type', dest='inp3', default=None, + help='this defines how output should be. A 2d image array or a 3d image array encoded with RGB color. Just pass 2d or 3d. The file will be saved one directory up. 2D image array is 3d but only information of one channel would be enough since all channels have the same values.') + options = parser.parse_args() + + possibles = globals() possibles.update(locals()) - x=pagexml2img(options.inp1,options.inp2,options.inp3) + x = pagexml2img(options.inp1, options.inp2, options.inp3) x.run() -if __name__=="__main__": - main() - - + +if __name__ == "__main__": + main() diff --git a/eynollah/eynollah/train/train.py b/eynollah/eynollah/train/train.py index 0cc5ef3..0881182 100644 --- a/eynollah/eynollah/train/train.py +++ b/eynollah/eynollah/train/train.py @@ -2,7 +2,7 @@ import sys import tensorflow as tf from keras.backend.tensorflow_backend import set_session -import keras , warnings +import keras, warnings from keras.optimizers import * from sacred import Experiment from models import * @@ -11,20 +11,21 @@ from keras.models import load_model from tqdm import tqdm + def configuration(): keras.backend.clear_session() tf.reset_default_graph() warnings.filterwarnings('ignore') - - os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID' + + os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID' config = tf.ConfigProto(log_device_placement=False, allow_soft_placement=True) - - + config.gpu_options.allow_growth = True - config.gpu_options.per_process_gpu_memory_fraction=0.95#0.95 - config.gpu_options.visible_device_list="0" + config.gpu_options.per_process_gpu_memory_fraction = 0.95 # 0.95 + config.gpu_options.visible_device_list = "0" set_session(tf.Session(config=config)) + def get_dirs_or_files(input_data): if os.path.isdir(input_data): image_input, labels_input = os.path.join(input_data, 'images/'), os.path.join(input_data, 'labels/') @@ -33,206 +34,188 @@ def get_dirs_or_files(input_data): assert os.path.isdir(labels_input), "{} is not a directory".format(labels_input) return image_input, labels_input + ex = Experiment() + @ex.config def config_params(): - n_classes=None # Number of classes. If your case study is binary case the set it to 2 and otherwise give your number of cases. - n_epochs=1 - input_height=224*1 - input_width=224*1 - weight_decay=1e-6 # Weight decay of l2 regularization of model layers. - n_batch=1 # Number of batches at each iteration. - learning_rate=1e-4 - patches=False # Make patches of image in order to use all information of image. In the case of page + n_classes = None # Number of classes. If your case study is binary case the set it to 2 and otherwise give your number of cases. + n_epochs = 1 + input_height = 224 * 1 + input_width = 224 * 1 + weight_decay = 1e-6 # Weight decay of l2 regularization of model layers. + n_batch = 1 # Number of batches at each iteration. + learning_rate = 1e-4 + patches = False # Make patches of image in order to use all information of image. In the case of page # extraction this should be set to false since model should see all image. - augmentation=False - flip_aug=False # Flip image (augmentation). - blur_aug=False # Blur patches of image (augmentation). - scaling=False # Scaling of patches (augmentation) will be imposed if this set to true. - binarization=False # Otsu thresholding. Used for augmentation in the case of binary case like textline prediction. For multicases should not be applied. - dir_train=None # Directory of training dataset (sub-folders should be named images and labels). - dir_eval=None # Directory of validation dataset (sub-folders should be named images and labels). - dir_output=None # Directory of output where the model should be saved. - pretraining=False # Set true to load pretrained weights of resnet50 encoder. - scaling_bluring=False - scaling_binarization=False - scaling_flip=False - thetha=[10,-10] - blur_k=['blur','guass','median'] # Used in order to blur image. Used for augmentation. - scales= [ 0.5, 2 ] # Scale patches with these scales. Used for augmentation. - flip_index=[0,1,-1] # Flip image. Used for augmentation. - continue_training = False # If + augmentation = False + flip_aug = False # Flip image (augmentation). + blur_aug = False # Blur patches of image (augmentation). + scaling = False # Scaling of patches (augmentation) will be imposed if this set to true. + binarization = False # Otsu thresholding. Used for augmentation in the case of binary case like textline prediction. For multicases should not be applied. + dir_train = None # Directory of training dataset (sub-folders should be named images and labels). + dir_eval = None # Directory of validation dataset (sub-folders should be named images and labels). + dir_output = None # Directory of output where the model should be saved. + pretraining = False # Set true to load pretrained weights of resnet50 encoder. + scaling_bluring = False + scaling_binarization = False + scaling_flip = False + thetha = [10, -10] + blur_k = ['blur', 'guass', 'median'] # Used in order to blur image. Used for augmentation. + scales = [0.5, 2] # Scale patches with these scales. Used for augmentation. + flip_index = [0, 1, -1] # Flip image. Used for augmentation. + continue_training = False # If index_start = 0 dir_of_start_model = '' is_loss_soft_dice = False weighted_loss = False data_is_provided = False + @ex.automain -def run(n_classes,n_epochs,input_height, - input_width,weight_decay,weighted_loss, - index_start,dir_of_start_model,is_loss_soft_dice, - n_batch,patches,augmentation,flip_aug - ,blur_aug,scaling, binarization, - blur_k,scales,dir_train,data_is_provided, - scaling_bluring,scaling_binarization,rotation, - rotation_not_90,thetha,scaling_flip,continue_training, - flip_index,dir_eval ,dir_output,pretraining,learning_rate): - - +def run(n_classes, n_epochs, input_height, + input_width, weight_decay, weighted_loss, + index_start, dir_of_start_model, is_loss_soft_dice, + n_batch, patches, augmentation, flip_aug, + blur_aug, scaling, binarization, + blur_k, scales, dir_train, data_is_provided, + scaling_bluring, scaling_binarization, rotation, + rotation_not_90, thetha, scaling_flip, continue_training, + flip_index, dir_eval, dir_output, pretraining, learning_rate): if data_is_provided: - dir_train_flowing=os.path.join(dir_output,'train') - dir_eval_flowing=os.path.join(dir_output,'eval') - - dir_flow_train_imgs=os.path.join(dir_train_flowing,'images') - dir_flow_train_labels=os.path.join(dir_train_flowing,'labels') - - dir_flow_eval_imgs=os.path.join(dir_eval_flowing,'images') - dir_flow_eval_labels=os.path.join(dir_eval_flowing,'labels') - + dir_train_flowing = os.path.join(dir_output, 'train') + dir_eval_flowing = os.path.join(dir_output, 'eval') + + dir_flow_train_imgs = os.path.join(dir_train_flowing, 'images') + dir_flow_train_labels = os.path.join(dir_train_flowing, 'labels') + + dir_flow_eval_imgs = os.path.join(dir_eval_flowing, 'images') + dir_flow_eval_labels = os.path.join(dir_eval_flowing, 'labels') + configuration() - + else: - dir_img,dir_seg=get_dirs_or_files(dir_train) - dir_img_val,dir_seg_val=get_dirs_or_files(dir_eval) - + dir_img, dir_seg = get_dirs_or_files(dir_train) + dir_img_val, dir_seg_val = get_dirs_or_files(dir_eval) + # make first a directory in output for both training and evaluations in order to flow data from these directories. - dir_train_flowing=os.path.join(dir_output,'train') - dir_eval_flowing=os.path.join(dir_output,'eval') - - dir_flow_train_imgs=os.path.join(dir_train_flowing,'images/') - dir_flow_train_labels=os.path.join(dir_train_flowing,'labels/') - - dir_flow_eval_imgs=os.path.join(dir_eval_flowing,'images/') - dir_flow_eval_labels=os.path.join(dir_eval_flowing,'labels/') - + dir_train_flowing = os.path.join(dir_output, 'train') + dir_eval_flowing = os.path.join(dir_output, 'eval') + + dir_flow_train_imgs = os.path.join(dir_train_flowing, 'images/') + dir_flow_train_labels = os.path.join(dir_train_flowing, 'labels/') + + dir_flow_eval_imgs = os.path.join(dir_eval_flowing, 'images/') + dir_flow_eval_labels = os.path.join(dir_eval_flowing, 'labels/') + if os.path.isdir(dir_train_flowing): - os.system('rm -rf '+dir_train_flowing) + os.system('rm -rf ' + dir_train_flowing) os.makedirs(dir_train_flowing) else: os.makedirs(dir_train_flowing) - + if os.path.isdir(dir_eval_flowing): - os.system('rm -rf '+dir_eval_flowing) + os.system('rm -rf ' + dir_eval_flowing) os.makedirs(dir_eval_flowing) else: os.makedirs(dir_eval_flowing) - os.mkdir(dir_flow_train_imgs) os.mkdir(dir_flow_train_labels) - + os.mkdir(dir_flow_eval_imgs) os.mkdir(dir_flow_eval_labels) - - - #set the gpu configuration - configuration() + # set the gpu configuration + configuration() - #writing patches into a sub-folder in order to be flowed from directory. - provide_patches(dir_img,dir_seg,dir_flow_train_imgs, + # writing patches into a sub-folder in order to be flowed from directory. + provide_patches(dir_img, dir_seg, dir_flow_train_imgs, dir_flow_train_labels, - input_height,input_width,blur_k,blur_aug, - flip_aug,binarization,scaling,scales,flip_index, - scaling_bluring,scaling_binarization,rotation, - rotation_not_90,thetha,scaling_flip, - augmentation=augmentation,patches=patches) - - provide_patches(dir_img_val,dir_seg_val,dir_flow_eval_imgs, + input_height, input_width, blur_k, blur_aug, + flip_aug, binarization, scaling, scales, flip_index, + scaling_bluring, scaling_binarization, rotation, + rotation_not_90, thetha, scaling_flip, + augmentation=augmentation, patches=patches) + + provide_patches(dir_img_val, dir_seg_val, dir_flow_eval_imgs, dir_flow_eval_labels, - input_height,input_width,blur_k,blur_aug, - flip_aug,binarization,scaling,scales,flip_index, - scaling_bluring,scaling_binarization,rotation, - rotation_not_90,thetha,scaling_flip, - augmentation=False,patches=patches) - - - + input_height, input_width, blur_k, blur_aug, + flip_aug, binarization, scaling, scales, flip_index, + scaling_bluring, scaling_binarization, rotation, + rotation_not_90, thetha, scaling_flip, + augmentation=False, patches=patches) + if weighted_loss: - weights=np.zeros(n_classes) + weights = np.zeros(n_classes) if data_is_provided: for obj in os.listdir(dir_flow_train_labels): try: - label_obj=cv2.imread(dir_flow_train_labels+'/'+obj) - label_obj_one_hot=get_one_hot( label_obj,label_obj.shape[0],label_obj.shape[1],n_classes) - weights+=(label_obj_one_hot.sum(axis=0)).sum(axis=0) + label_obj = cv2.imread(dir_flow_train_labels + '/' + obj) + label_obj_one_hot = get_one_hot(label_obj, label_obj.shape[0], label_obj.shape[1], n_classes) + weights += (label_obj_one_hot.sum(axis=0)).sum(axis=0) except: pass else: - + for obj in os.listdir(dir_seg): try: - label_obj=cv2.imread(dir_seg+'/'+obj) - label_obj_one_hot=get_one_hot( label_obj,label_obj.shape[0],label_obj.shape[1],n_classes) - weights+=(label_obj_one_hot.sum(axis=0)).sum(axis=0) + label_obj = cv2.imread(dir_seg + '/' + obj) + label_obj_one_hot = get_one_hot(label_obj, label_obj.shape[0], label_obj.shape[1], n_classes) + weights += (label_obj_one_hot.sum(axis=0)).sum(axis=0) except: pass - - - weights=1.00/weights - - weights=weights/float(np.sum(weights)) - weights=weights/float(np.min(weights)) - weights=weights/float(np.sum(weights)) - - - + + weights = 1.00 / weights + + weights = weights / float(np.sum(weights)) + weights = weights / float(np.min(weights)) + weights = weights / float(np.sum(weights)) + if continue_training: if is_loss_soft_dice: - model = load_model (dir_of_start_model, compile = True, custom_objects={'soft_dice_loss': soft_dice_loss}) + model = load_model(dir_of_start_model, compile=True, custom_objects={'soft_dice_loss': soft_dice_loss}) if weighted_loss: - model = load_model (dir_of_start_model, compile = True, custom_objects={'loss': weighted_categorical_crossentropy(weights)}) + model = load_model(dir_of_start_model, compile=True, + custom_objects={'loss': weighted_categorical_crossentropy(weights)}) if not is_loss_soft_dice and not weighted_loss: - model = load_model (dir_of_start_model, compile = True) + model = load_model(dir_of_start_model, compile=True) else: - #get our model. + # get our model. index_start = 0 - model = resnet50_unet(n_classes, input_height, input_width,weight_decay,pretraining) - - #if you want to see the model structure just uncomment model summary. - #model.summary() - + model = resnet50_unet(n_classes, input_height, input_width, weight_decay, pretraining) + + # if you want to see the model structure just uncomment model summary. + # model.summary() if not is_loss_soft_dice and not weighted_loss: model.compile(loss='categorical_crossentropy', - optimizer = Adam(lr=learning_rate),metrics=['accuracy']) - if is_loss_soft_dice: + optimizer=Adam(lr=learning_rate), metrics=['accuracy']) + if is_loss_soft_dice: model.compile(loss=soft_dice_loss, - optimizer = Adam(lr=learning_rate),metrics=['accuracy']) - + optimizer=Adam(lr=learning_rate), metrics=['accuracy']) + if weighted_loss: model.compile(loss=weighted_categorical_crossentropy(weights), - optimizer = Adam(lr=learning_rate),metrics=['accuracy']) - - #generating train and evaluation data - train_gen = data_gen(dir_flow_train_imgs,dir_flow_train_labels, batch_size = n_batch, - input_height=input_height, input_width=input_width,n_classes=n_classes ) - val_gen = data_gen(dir_flow_eval_imgs,dir_flow_eval_labels, batch_size = n_batch, - input_height=input_height, input_width=input_width,n_classes=n_classes ) - - for i in tqdm(range(index_start, n_epochs+index_start)): + optimizer=Adam(lr=learning_rate), metrics=['accuracy']) + + # generating train and evaluation data + train_gen = data_gen(dir_flow_train_imgs, dir_flow_train_labels, batch_size=n_batch, + input_height=input_height, input_width=input_width, n_classes=n_classes) + val_gen = data_gen(dir_flow_eval_imgs, dir_flow_eval_labels, batch_size=n_batch, + input_height=input_height, input_width=input_width, n_classes=n_classes) + + for i in tqdm(range(index_start, n_epochs + index_start)): model.fit_generator( train_gen, - steps_per_epoch=int(len(os.listdir(dir_flow_train_imgs))/n_batch)-1, + steps_per_epoch=int(len(os.listdir(dir_flow_train_imgs)) / n_batch) - 1, validation_data=val_gen, validation_steps=1, epochs=1) - model.save(dir_output+'/'+'model_'+str(i)+'.h5') - - - #os.system('rm -rf '+dir_train_flowing) - #os.system('rm -rf '+dir_eval_flowing) - - #model.save(dir_output+'/'+'model'+'.h5') - - - - - - - - + model.save(dir_output + '/' + 'model_' + str(i) + '.h5') + # os.system('rm -rf '+dir_train_flowing) + # os.system('rm -rf '+dir_eval_flowing) + # model.save(dir_output+'/'+'model'+'.h5') diff --git a/eynollah/eynollah/train/utils.py b/eynollah/eynollah/train/utils.py index 19ab46e..64263f4 100644 --- a/eynollah/eynollah/train/utils.py +++ b/eynollah/eynollah/train/utils.py @@ -10,18 +10,17 @@ import math - -def bluring(img_in,kind): - if kind=='guass': - img_blur = cv2.GaussianBlur(img_in,(5,5),0) - elif kind=="median": - img_blur = cv2.medianBlur(img_in,5) - elif kind=='blur': - img_blur=cv2.blur(img_in,(5,5)) +def bluring(img_in, kind): + if kind == 'guass': + img_blur = cv2.GaussianBlur(img_in, (5, 5), 0) + elif kind == "median": + img_blur = cv2.medianBlur(img_in, 5) + elif kind == 'blur': + img_blur = cv2.blur(img_in, (5, 5)) return img_blur -def elastic_transform(image, alpha, sigma,seedj, random_state=None): - + +def elastic_transform(image, alpha, sigma, seedj, random_state=None): """Elastic deformation of images as described in [Simard2003]_. .. [Simard2003] Simard, Steinkraus and Platt, "Best Practices for Convolutional Neural Networks applied to Visual Document Analysis", in @@ -37,461 +36,459 @@ def elastic_transform(image, alpha, sigma,seedj, random_state=None): dz = np.zeros_like(dx) x, y, z = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]), np.arange(shape[2])) - indices = np.reshape(y+dy, (-1, 1)), np.reshape(x+dx, (-1, 1)), np.reshape(z, (-1, 1)) + indices = np.reshape(y + dy, (-1, 1)), np.reshape(x + dx, (-1, 1)), np.reshape(z, (-1, 1)) distored_image = map_coordinates(image, indices, order=1, mode='reflect') return distored_image.reshape(image.shape) + def rotation_90(img): - img_rot=np.zeros((img.shape[1],img.shape[0],img.shape[2])) - img_rot[:,:,0]=img[:,:,0].T - img_rot[:,:,1]=img[:,:,1].T - img_rot[:,:,2]=img[:,:,2].T + img_rot = np.zeros((img.shape[1], img.shape[0], img.shape[2])) + img_rot[:, :, 0] = img[:, :, 0].T + img_rot[:, :, 1] = img[:, :, 1].T + img_rot[:, :, 2] = img[:, :, 2].T return img_rot + def rotatedRectWithMaxArea(w, h, angle): - """ + """ Given a rectangle of size wxh that has been rotated by 'angle' (in radians), computes the width and height of the largest possible axis-aligned rectangle (maximal area) within the rotated rectangle. """ - if w <= 0 or h <= 0: - return 0,0 - - width_is_longer = w >= h - side_long, side_short = (w,h) if width_is_longer else (h,w) - - # since the solutions for angle, -angle and 180-angle are all the same, - # if suffices to look at the first quadrant and the absolute values of sin,cos: - sin_a, cos_a = abs(math.sin(angle)), abs(math.cos(angle)) - if side_short <= 2.*sin_a*cos_a*side_long or abs(sin_a-cos_a) < 1e-10: - # half constrained case: two crop corners touch the longer side, - # the other two corners are on the mid-line parallel to the longer line - x = 0.5*side_short - wr,hr = (x/sin_a,x/cos_a) if width_is_longer else (x/cos_a,x/sin_a) - else: - # fully constrained case: crop touches all 4 sides - cos_2a = cos_a*cos_a - sin_a*sin_a - wr,hr = (w*cos_a - h*sin_a)/cos_2a, (h*cos_a - w*sin_a)/cos_2a - - return wr,hr - -def rotate_max_area(image,rotated, rotated_label,angle): + if w <= 0 or h <= 0: + return 0, 0 + + width_is_longer = w >= h + side_long, side_short = (w, h) if width_is_longer else (h, w) + + # since the solutions for angle, -angle and 180-angle are all the same, + # if suffices to look at the first quadrant and the absolute values of sin,cos: + sin_a, cos_a = abs(math.sin(angle)), abs(math.cos(angle)) + if side_short <= 2. * sin_a * cos_a * side_long or abs(sin_a - cos_a) < 1e-10: + # half constrained case: two crop corners touch the longer side, + # the other two corners are on the mid-line parallel to the longer line + x = 0.5 * side_short + wr, hr = (x / sin_a, x / cos_a) if width_is_longer else (x / cos_a, x / sin_a) + else: + # fully constrained case: crop touches all 4 sides + cos_2a = cos_a * cos_a - sin_a * sin_a + wr, hr = (w * cos_a - h * sin_a) / cos_2a, (h * cos_a - w * sin_a) / cos_2a + + return wr, hr + + +def rotate_max_area(image, rotated, rotated_label, angle): """ image: cv2 image matrix object angle: in degree """ wr, hr = rotatedRectWithMaxArea(image.shape[1], image.shape[0], math.radians(angle)) h, w, _ = rotated.shape - y1 = h//2 - int(hr/2) + y1 = h // 2 - int(hr / 2) y2 = y1 + int(hr) - x1 = w//2 - int(wr/2) + x1 = w // 2 - int(wr / 2) x2 = x1 + int(wr) - return rotated[y1:y2, x1:x2],rotated_label[y1:y2, x1:x2] -def rotation_not_90_func(img,label,thetha): - rotated=imutils.rotate(img,thetha) - rotated_label=imutils.rotate(label,thetha) - return rotate_max_area(img, rotated,rotated_label,thetha) + return rotated[y1:y2, x1:x2], rotated_label[y1:y2, x1:x2] + + +def rotation_not_90_func(img, label, thetha): + rotated = imutils.rotate(img, thetha) + rotated_label = imutils.rotate(label, thetha) + return rotate_max_area(img, rotated, rotated_label, thetha) + def color_images(seg, n_classes): - ann_u=range(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(float) - colors=sns.color_palette("hls", n_classes) - + ann_u = range(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(float) + colors = sns.color_palette("hls", n_classes) + for c in ann_u: - c=int(c) - segl=(seg==c) - seg_img[:,:,0]+=segl*(colors[c][0]) - seg_img[:,:,1]+=segl*(colors[c][1]) - seg_img[:,:,2]+=segl*(colors[c][2]) + c = int(c) + segl = (seg == c) + seg_img[:, :, 0] += segl * (colors[c][0]) + seg_img[:, :, 1] += segl * (colors[c][1]) + seg_img[:, :, 2] += segl * (colors[c][2]) return seg_img - -def resize_image(seg_in,input_height,input_width): - return cv2.resize(seg_in,(input_width,input_height),interpolation=cv2.INTER_NEAREST) -def get_one_hot(seg,input_height,input_width,n_classes): - seg=seg[:,:,0] - seg_f=np.zeros((input_height, input_width,n_classes)) + +def resize_image(seg_in, input_height, input_width): + return cv2.resize(seg_in, (input_width, input_height), interpolation=cv2.INTER_NEAREST) + + +def get_one_hot(seg, input_height, input_width, n_classes): + seg = seg[:, :, 0] + seg_f = np.zeros((input_height, input_width, n_classes)) for j in range(n_classes): - seg_f[:,:,j]=(seg==j).astype(int) + seg_f[:, :, j] = (seg == j).astype(int) return seg_f - -def IoU(Yi,y_predi): - ## mean Intersection over Union - ## Mean IoU = TP/(FN + TP + FP) + +def IoU(Yi, y_predi): + # mean Intersection over Union + # Mean IoU = TP/(FN + TP + FP) IoUs = [] - classes_true=np.unique(Yi) + classes_true = np.unique(Yi) for c in classes_true: - 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) - print("class {:02.0f}: #TP={:6.0f}, #FP={:6.0f}, #FN={:5.0f}, IoU={:4.3f}".format(c,TP,FP,FN,IoU)) + 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) + 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) mIoU = np.mean(IoUs) print("_________________") print("Mean IoU: {:4.3f}".format(mIoU)) return mIoU -def data_gen(img_folder, mask_folder, batch_size,input_height, input_width,n_classes): + + +def data_gen(img_folder, mask_folder, batch_size, input_height, input_width, n_classes): c = 0 - n = [f for f in os.listdir(img_folder) if not f.startswith('.')]# os.listdir(img_folder) #List of training images + n = [f for f in os.listdir(img_folder) if not f.startswith('.')] # os.listdir(img_folder) #List of training images random.shuffle(n) while True: img = np.zeros((batch_size, input_height, input_width, 3)).astype('float') mask = np.zeros((batch_size, input_height, input_width, n_classes)).astype('float') - - for i in range(c, c+batch_size): #initially from 0 to 16, c = 0. - #print(img_folder+'/'+n[i]) - + + for i in range(c, c + batch_size): # initially from 0 to 16, c = 0. + # print(img_folder+'/'+n[i]) + try: - filename=n[i].split('.')[0] - - train_img = cv2.imread(img_folder+'/'+n[i])/255. - train_img = cv2.resize(train_img, (input_width, input_height),interpolation=cv2.INTER_NEAREST)# Read an image from folder and resize - - img[i-c] = train_img #add to array - img[0], img[1], and so on. - train_mask = cv2.imread(mask_folder+'/'+filename+'.png') - #print(mask_folder+'/'+filename+'.png') - #print(train_mask.shape) - train_mask = get_one_hot( resize_image(train_mask,input_height,input_width),input_height,input_width,n_classes) - #train_mask = train_mask.reshape(224, 224, 1) # Add extra dimension for parity with train_img size [512 * 512 * 3] - - mask[i-c] = train_mask + filename = n[i].split('.')[0] + + train_img = cv2.imread(img_folder + '/' + n[i]) / 255. + train_img = cv2.resize(train_img, (input_width, input_height), + interpolation=cv2.INTER_NEAREST) # Read an image from folder and resize + + img[i - c] = train_img # add to array - img[0], img[1], and so on. + train_mask = cv2.imread(mask_folder + '/' + filename + '.png') + # print(mask_folder+'/'+filename+'.png') + # print(train_mask.shape) + train_mask = get_one_hot(resize_image(train_mask, input_height, input_width), input_height, input_width, + n_classes) + # train_mask = train_mask.reshape(224, 224, 1) # Add extra dimension for parity with train_img size [512 * 512 * 3] + + mask[i - c] = train_mask except: - img[i-c] = np.ones((input_height, input_width, 3)).astype('float') - mask[i-c] = np.zeros((input_height, input_width, n_classes)).astype('float') - - - - c+=batch_size - if(c+batch_size>=len(os.listdir(img_folder))): - c=0 + img[i - c] = np.ones((input_height, input_width, 3)).astype('float') + mask[i - c] = np.zeros((input_height, input_width, n_classes)).astype('float') + + c += batch_size + if c + batch_size >= len(os.listdir(img_folder)): + c = 0 random.shuffle(n) yield img, mask - + + def otsu_copy(img): - img_r=np.zeros(img.shape) - img1=img[:,:,0] - img2=img[:,:,1] - img3=img[:,:,2] - _, threshold1 = cv2.threshold(img1, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU) - _, threshold2 = cv2.threshold(img2, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU) - _, threshold3 = cv2.threshold(img3, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU) - img_r[:,:,0]=threshold1 - img_r[:,:,1]=threshold1 - img_r[:,:,2]=threshold1 + img_r = np.zeros(img.shape) + img1 = img[:, :, 0] + img2 = img[:, :, 1] + img3 = img[:, :, 2] + _, threshold1 = cv2.threshold(img1, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) + _, threshold2 = cv2.threshold(img2, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) + _, threshold3 = cv2.threshold(img3, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) + img_r[:, :, 0] = threshold1 + img_r[:, :, 1] = threshold1 + img_r[:, :, 2] = threshold1 return img_r -def get_patches(dir_img_f,dir_seg_f,img,label,height,width,indexer): - - if img.shape[0]int(nxf): - nxf=int(nxf)+1 - if nyf>int(nyf): - nyf=int(nyf)+1 - - nxf=int(nxf) - nyf=int(nyf) - + + +def get_patches(dir_img_f, dir_seg_f, img, label, height, width, indexer): + if img.shape[0] < height or img.shape[1] < width: + img, label = do_padding(img, label, height, width) + + img_h = img.shape[0] + img_w = img.shape[1] + + nxf = img_w / float(width) + nyf = img_h / float(height) + + if nxf > int(nxf): + nxf = int(nxf) + 1 + if nyf > int(nyf): + nyf = int(nyf) + 1 + + nxf = int(nxf) + nyf = int(nyf) + for i in range(nxf): for j in range(nyf): - index_x_d=i*width - index_x_u=(i+1)*width - - index_y_d=j*height - index_y_u=(j+1)*height - - if index_x_u>img_w: - index_x_u=img_w - index_x_d=img_w-width - if index_y_u>img_h: - index_y_u=img_h - index_y_d=img_h-height - - - img_patch=img[index_y_d:index_y_u,index_x_d:index_x_u,:] - label_patch=label[index_y_d:index_y_u,index_x_d:index_x_u,:] - - cv2.imwrite(dir_img_f+'/img_'+str(indexer)+'.png', img_patch ) - cv2.imwrite(dir_seg_f+'/img_'+str(indexer)+'.png' , label_patch ) - indexer+=1 - + index_x_d = i * width + index_x_u = (i + 1) * width + + index_y_d = j * height + index_y_u = (j + 1) * height + + if index_x_u > img_w: + index_x_u = img_w + index_x_d = img_w - width + if index_y_u > img_h: + index_y_u = img_h + index_y_d = img_h - height + + img_patch = img[index_y_d:index_y_u, index_x_d:index_x_u, :] + label_patch = label[index_y_d:index_y_u, index_x_d:index_x_u, :] + + cv2.imwrite(dir_img_f + '/img_' + str(indexer) + '.png', img_patch) + cv2.imwrite(dir_seg_f + '/img_' + str(indexer) + '.png', label_patch) + indexer += 1 + return indexer -def do_padding(img,label,height,width): - - height_new=img.shape[0] - width_new=img.shape[1] - - h_start=0 - w_start=0 - - if img.shape[0]int(nxf): - nxf=int(nxf)+1 - if nyf>int(nyf): - nyf=int(nyf)+1 - - nxf=int(nxf) - nyf=int(nyf) - + +def do_padding(img, label, height, width): + height_new = img.shape[0] + width_new = img.shape[1] + + h_start = 0 + w_start = 0 + + if img.shape[0] < height: + h_start = int(abs(height - img.shape[0]) / 2.) + height_new = height + + if img.shape[1] < width: + w_start = int(abs(width - img.shape[1]) / 2.) + width_new = width + + img_new = np.ones((height_new, width_new, img.shape[2])).astype(float) * 255 + label_new = np.zeros((height_new, width_new, label.shape[2])).astype(float) + + img_new[h_start:h_start + img.shape[0], w_start:w_start + img.shape[1], :] = np.copy(img[:, :, :]) + label_new[h_start:h_start + label.shape[0], w_start:w_start + label.shape[1], :] = np.copy(label[:, :, :]) + + return img_new, label_new + + +def get_patches_num_scale(dir_img_f, dir_seg_f, img, label, height, width, indexer, n_patches, scaler): + if img.shape[0] < height or img.shape[1] < width: + img, label = do_padding(img, label, height, width) + + img_h = img.shape[0] + img_w = img.shape[1] + + height_scale = int(height * scaler) + width_scale = int(width * scaler) + + nxf = img_w / float(width_scale) + nyf = img_h / float(height_scale) + + if nxf > int(nxf): + nxf = int(nxf) + 1 + if nyf > int(nyf): + nyf = int(nyf) + 1 + + nxf = int(nxf) + nyf = int(nyf) + for i in range(nxf): for j in range(nyf): - index_x_d=i*width_scale - index_x_u=(i+1)*width_scale - - index_y_d=j*height_scale - index_y_u=(j+1)*height_scale - - if index_x_u>img_w: - index_x_u=img_w - index_x_d=img_w-width_scale - if index_y_u>img_h: - index_y_u=img_h - index_y_d=img_h-height_scale - - - img_patch=img[index_y_d:index_y_u,index_x_d:index_x_u,:] - label_patch=label[index_y_d:index_y_u,index_x_d:index_x_u,:] - - img_patch=resize_image(img_patch,height,width) - label_patch=resize_image(label_patch,height,width) - - cv2.imwrite(dir_img_f+'/img_'+str(indexer)+'.png', img_patch ) - cv2.imwrite(dir_seg_f+'/img_'+str(indexer)+'.png' , label_patch ) - indexer+=1 + index_x_d = i * width_scale + index_x_u = (i + 1) * width_scale + + index_y_d = j * height_scale + index_y_u = (j + 1) * height_scale + + if index_x_u > img_w: + index_x_u = img_w + index_x_d = img_w - width_scale + if index_y_u > img_h: + index_y_u = img_h + index_y_d = img_h - height_scale + + img_patch = img[index_y_d:index_y_u, index_x_d:index_x_u, :] + label_patch = label[index_y_d:index_y_u, index_x_d:index_x_u, :] + + img_patch = resize_image(img_patch, height, width) + label_patch = resize_image(label_patch, height, width) + + cv2.imwrite(dir_img_f + '/img_' + str(indexer) + '.png', img_patch) + cv2.imwrite(dir_seg_f + '/img_' + str(indexer) + '.png', label_patch) + indexer += 1 return indexer -def get_patches_num_scale_new(dir_img_f,dir_seg_f,img,label,height,width,indexer,scaler): - img=resize_image(img,int(img.shape[0]*scaler),int(img.shape[1]*scaler)) - label=resize_image(label,int(label.shape[0]*scaler),int(label.shape[1]*scaler)) - - if img.shape[0]int(nxf): - nxf=int(nxf)+1 - if nyf>int(nyf): - nyf=int(nyf)+1 - - nxf=int(nxf) - nyf=int(nyf) - + +def get_patches_num_scale_new(dir_img_f, dir_seg_f, img, label, height, width, indexer, scaler): + img = resize_image(img, int(img.shape[0] * scaler), int(img.shape[1] * scaler)) + label = resize_image(label, int(label.shape[0] * scaler), int(label.shape[1] * scaler)) + + if img.shape[0] < height or img.shape[1] < width: + img, label = do_padding(img, label, height, width) + + img_h = img.shape[0] + img_w = img.shape[1] + + height_scale = int(height * 1) + width_scale = int(width * 1) + + nxf = img_w / float(width_scale) + nyf = img_h / float(height_scale) + + if nxf > int(nxf): + nxf = int(nxf) + 1 + if nyf > int(nyf): + nyf = int(nyf) + 1 + + nxf = int(nxf) + nyf = int(nyf) + for i in range(nxf): for j in range(nyf): - index_x_d=i*width_scale - index_x_u=(i+1)*width_scale - - index_y_d=j*height_scale - index_y_u=(j+1)*height_scale - - if index_x_u>img_w: - index_x_u=img_w - index_x_d=img_w-width_scale - if index_y_u>img_h: - index_y_u=img_h - index_y_d=img_h-height_scale - - - img_patch=img[index_y_d:index_y_u,index_x_d:index_x_u,:] - label_patch=label[index_y_d:index_y_u,index_x_d:index_x_u,:] - - #img_patch=resize_image(img_patch,height,width) - #label_patch=resize_image(label_patch,height,width) - - cv2.imwrite(dir_img_f+'/img_'+str(indexer)+'.png', img_patch ) - cv2.imwrite(dir_seg_f+'/img_'+str(indexer)+'.png' , label_patch ) - indexer+=1 + index_x_d = i * width_scale + index_x_u = (i + 1) * width_scale + + index_y_d = j * height_scale + index_y_u = (j + 1) * height_scale + + if index_x_u > img_w: + index_x_u = img_w + index_x_d = img_w - width_scale + if index_y_u > img_h: + index_y_u = img_h + index_y_d = img_h - height_scale + + img_patch = img[index_y_d:index_y_u, index_x_d:index_x_u, :] + label_patch = label[index_y_d:index_y_u, index_x_d:index_x_u, :] + + # img_patch=resize_image(img_patch,height,width) + # label_patch=resize_image(label_patch,height,width) + + cv2.imwrite(dir_img_f + '/img_' + str(indexer) + '.png', img_patch) + cv2.imwrite(dir_seg_f + '/img_' + str(indexer) + '.png', label_patch) + indexer += 1 return indexer -def provide_patches(dir_img,dir_seg,dir_flow_train_imgs, +def provide_patches(dir_img, dir_seg, dir_flow_train_imgs, dir_flow_train_labels, - input_height,input_width,blur_k,blur_aug, - flip_aug,binarization,scaling,scales,flip_index, - scaling_bluring,scaling_binarization,rotation, - rotation_not_90,thetha,scaling_flip, - augmentation=False,patches=False): - - imgs_cv_train=np.array(os.listdir(dir_img)) - segs_cv_train=np.array(os.listdir(dir_seg)) - - indexer=0 - for im, seg_i in tqdm(zip(imgs_cv_train,segs_cv_train)): - img_name=im.split('.')[0] + input_height, input_width, blur_k, blur_aug, + flip_aug, binarization, scaling, scales, flip_index, + scaling_bluring, scaling_binarization, rotation, + rotation_not_90, thetha, scaling_flip, + augmentation=False, patches=False): + imgs_cv_train = np.array(os.listdir(dir_img)) + segs_cv_train = np.array(os.listdir(dir_seg)) + + indexer = 0 + for im, seg_i in tqdm(zip(imgs_cv_train, segs_cv_train)): + img_name = im.split('.')[0] if not patches: - cv2.imwrite(dir_flow_train_imgs+'/img_'+str(indexer)+'.png', resize_image(cv2.imread(dir_img+'/'+im),input_height,input_width ) ) - cv2.imwrite(dir_flow_train_labels+'/img_'+str(indexer)+'.png' , resize_image(cv2.imread(dir_seg+'/'+img_name+'.png'),input_height,input_width ) ) - indexer+=1 - + cv2.imwrite(dir_flow_train_imgs + '/img_' + str(indexer) + '.png', + resize_image(cv2.imread(dir_img + '/' + im), input_height, input_width)) + cv2.imwrite(dir_flow_train_labels + '/img_' + str(indexer) + '.png', + resize_image(cv2.imread(dir_seg + '/' + img_name + '.png'), input_height, input_width)) + indexer += 1 + if augmentation: if flip_aug: for f_i in flip_index: - cv2.imwrite(dir_flow_train_imgs+'/img_'+str(indexer)+'.png', - resize_image(cv2.flip(cv2.imread(dir_img+'/'+im),f_i),input_height,input_width) ) - - cv2.imwrite(dir_flow_train_labels+'/img_'+str(indexer)+'.png' , - resize_image(cv2.flip(cv2.imread(dir_seg+'/'+img_name+'.png'),f_i),input_height,input_width) ) - indexer+=1 - - if blur_aug: + cv2.imwrite(dir_flow_train_imgs + '/img_' + str(indexer) + '.png', + resize_image(cv2.flip(cv2.imread(dir_img + '/' + im), f_i), input_height, + input_width)) + + cv2.imwrite(dir_flow_train_labels + '/img_' + str(indexer) + '.png', + resize_image(cv2.flip(cv2.imread(dir_seg + '/' + img_name + '.png'), f_i), + input_height, input_width)) + indexer += 1 + + if blur_aug: for blur_i in blur_k: - cv2.imwrite(dir_flow_train_imgs+'/img_'+str(indexer)+'.png', - (resize_image(bluring(cv2.imread(dir_img+'/'+im),blur_i),input_height,input_width) ) ) - - cv2.imwrite(dir_flow_train_labels+'/img_'+str(indexer)+'.png' , - resize_image(cv2.imread(dir_seg+'/'+img_name+'.png'),input_height,input_width) ) - indexer+=1 - - + cv2.imwrite(dir_flow_train_imgs + '/img_' + str(indexer) + '.png', + (resize_image(bluring(cv2.imread(dir_img + '/' + im), blur_i), input_height, + input_width))) + + cv2.imwrite(dir_flow_train_labels + '/img_' + str(indexer) + '.png', + resize_image(cv2.imread(dir_seg + '/' + img_name + '.png'), input_height, + input_width)) + indexer += 1 + if binarization: - cv2.imwrite(dir_flow_train_imgs+'/img_'+str(indexer)+'.png', - resize_image(otsu_copy( cv2.imread(dir_img+'/'+im)),input_height,input_width )) - - cv2.imwrite(dir_flow_train_labels+'/img_'+str(indexer)+'.png', - resize_image( cv2.imread(dir_seg+'/'+img_name+'.png'),input_height,input_width )) - indexer+=1 - - - - - - + cv2.imwrite(dir_flow_train_imgs + '/img_' + str(indexer) + '.png', + resize_image(otsu_copy(cv2.imread(dir_img + '/' + im)), input_height, input_width)) + + cv2.imwrite(dir_flow_train_labels + '/img_' + str(indexer) + '.png', + resize_image(cv2.imread(dir_seg + '/' + img_name + '.png'), input_height, input_width)) + indexer += 1 + if patches: - - indexer=get_patches(dir_flow_train_imgs,dir_flow_train_labels, - cv2.imread(dir_img+'/'+im),cv2.imread(dir_seg+'/'+img_name+'.png'), - input_height,input_width,indexer=indexer) - + + indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, + cv2.imread(dir_img + '/' + im), cv2.imread(dir_seg + '/' + img_name + '.png'), + input_height, input_width, indexer=indexer) + if augmentation: - + if rotation: - - - indexer=get_patches(dir_flow_train_imgs,dir_flow_train_labels, - rotation_90( cv2.imread(dir_img+'/'+im) ), - rotation_90( cv2.imread(dir_seg+'/'+img_name+'.png') ), - input_height,input_width,indexer=indexer) - + indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, + rotation_90(cv2.imread(dir_img + '/' + im)), + rotation_90(cv2.imread(dir_seg + '/' + img_name + '.png')), + input_height, input_width, indexer=indexer) + if rotation_not_90: - + for thetha_i in thetha: - img_max_rotated,label_max_rotated=rotation_not_90_func(cv2.imread(dir_img+'/'+im),cv2.imread(dir_seg+'/'+img_name+'.png'),thetha_i) - indexer=get_patches(dir_flow_train_imgs,dir_flow_train_labels, - img_max_rotated, - label_max_rotated, - input_height,input_width,indexer=indexer) + img_max_rotated, label_max_rotated = rotation_not_90_func(cv2.imread(dir_img + '/' + im), + cv2.imread( + dir_seg + '/' + img_name + '.png'), + thetha_i) + indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, + img_max_rotated, + label_max_rotated, + input_height, input_width, indexer=indexer) if flip_aug: for f_i in flip_index: - indexer=get_patches(dir_flow_train_imgs,dir_flow_train_labels, - cv2.flip( cv2.imread(dir_img+'/'+im) , f_i), - cv2.flip( cv2.imread(dir_seg+'/'+img_name+'.png') ,f_i), - input_height,input_width,indexer=indexer) - if blur_aug: + indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, + cv2.flip(cv2.imread(dir_img + '/' + im), f_i), + cv2.flip(cv2.imread(dir_seg + '/' + img_name + '.png'), f_i), + input_height, input_width, indexer=indexer) + if blur_aug: for blur_i in blur_k: + indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, + bluring(cv2.imread(dir_img + '/' + im), blur_i), + cv2.imread(dir_seg + '/' + img_name + '.png'), + input_height, input_width, indexer=indexer) - indexer=get_patches(dir_flow_train_imgs,dir_flow_train_labels, - bluring( cv2.imread(dir_img+'/'+im) , blur_i), - cv2.imread(dir_seg+'/'+img_name+'.png'), - input_height,input_width,indexer=indexer) - - - if scaling: + if scaling: for sc_ind in scales: - indexer=get_patches_num_scale_new(dir_flow_train_imgs,dir_flow_train_labels, - cv2.imread(dir_img+'/'+im) , - cv2.imread(dir_seg+'/'+img_name+'.png'), - input_height,input_width,indexer=indexer,scaler=sc_ind) + indexer = get_patches_num_scale_new(dir_flow_train_imgs, dir_flow_train_labels, + cv2.imread(dir_img + '/' + im), + cv2.imread(dir_seg + '/' + img_name + '.png'), + input_height, input_width, indexer=indexer, scaler=sc_ind) if binarization: - indexer=get_patches(dir_flow_train_imgs,dir_flow_train_labels, - otsu_copy( cv2.imread(dir_img+'/'+im)), - cv2.imread(dir_seg+'/'+img_name+'.png'), - input_height,input_width,indexer=indexer) - - - - if scaling_bluring: + indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, + otsu_copy(cv2.imread(dir_img + '/' + im)), + cv2.imread(dir_seg + '/' + img_name + '.png'), + input_height, input_width, indexer=indexer) + + if scaling_bluring: for sc_ind in scales: for blur_i in blur_k: - indexer=get_patches_num_scale_new(dir_flow_train_imgs,dir_flow_train_labels, - bluring( cv2.imread(dir_img+'/'+im) , blur_i) , - cv2.imread(dir_seg+'/'+img_name+'.png') , - input_height,input_width,indexer=indexer,scaler=sc_ind) + indexer = get_patches_num_scale_new(dir_flow_train_imgs, dir_flow_train_labels, + bluring(cv2.imread(dir_img + '/' + im), blur_i), + cv2.imread(dir_seg + '/' + img_name + '.png'), + input_height, input_width, indexer=indexer, + scaler=sc_ind) - if scaling_binarization: + if scaling_binarization: for sc_ind in scales: - indexer=get_patches_num_scale_new(dir_flow_train_imgs,dir_flow_train_labels, - otsu_copy( cv2.imread(dir_img+'/'+im)) , - cv2.imread(dir_seg+'/'+img_name+'.png'), - input_height,input_width,indexer=indexer,scaler=sc_ind) - - if scaling_flip: + indexer = get_patches_num_scale_new(dir_flow_train_imgs, dir_flow_train_labels, + otsu_copy(cv2.imread(dir_img + '/' + im)), + cv2.imread(dir_seg + '/' + img_name + '.png'), + input_height, input_width, indexer=indexer, scaler=sc_ind) + + if scaling_flip: for sc_ind in scales: for f_i in flip_index: - indexer=get_patches_num_scale_new(dir_flow_train_imgs,dir_flow_train_labels, - cv2.flip( cv2.imread(dir_img+'/'+im) , f_i) , - cv2.flip(cv2.imread(dir_seg+'/'+img_name+'.png') ,f_i) , - input_height,input_width,indexer=indexer,scaler=sc_ind) - - - - - - - + indexer = get_patches_num_scale_new(dir_flow_train_imgs, dir_flow_train_labels, + cv2.flip(cv2.imread(dir_img + '/' + im), f_i), + cv2.flip(cv2.imread(dir_seg + '/' + img_name + '.png'), + f_i), + input_height, input_width, indexer=indexer, + scaler=sc_ind) diff --git a/eynollah/eynollah/utils/contour.py b/eynollah/eynollah/utils/contour.py index 53b39b5..95a1199 100644 --- a/eynollah/eynollah/utils/contour.py +++ b/eynollah/eynollah/utils/contour.py @@ -5,6 +5,8 @@ from .rotate import rotate_image, rotation_image_new from multiprocessing import Process, Queue, cpu_count from multiprocessing import Pool + + def contours_in_same_horizon(cy_main_hor): X1 = np.zeros((len(cy_main_hor), len(cy_main_hor))) X2 = np.zeros((len(cy_main_hor), len(cy_main_hor))) @@ -22,6 +24,7 @@ def contours_in_same_horizon(cy_main_hor): all_args.append(list(set(list_h))) return np.unique(np.array(all_args, dtype=object)) + def find_contours_mean_y_diff(contours_main): M_main = [cv2.moments(contours_main[j]) for j in range(len(contours_main))] cy_main = [(M_main[j]["m01"] / (M_main[j]["m00"] + 1e-32)) for j in range(len(M_main))] @@ -42,10 +45,11 @@ def get_text_region_boxes_by_given_contours(contours): del contours return boxes, contours_new + def filter_contours_area_of_image(image, contours, hierarchy, max_area, min_area): found_polygons_early = list() - for jv,c in enumerate(contours): + for jv, c in enumerate(contours): if len(c) < 3: # A polygon cannot have less than 3 points continue @@ -55,17 +59,18 @@ def filter_contours_area_of_image(image, contours, hierarchy, max_area, min_area found_polygons_early.append(np.array([[point] for point in polygon.exterior.coords], dtype=np.uint)) return found_polygons_early + def filter_contours_area_of_image_tables(image, contours, hierarchy, max_area, min_area): found_polygons_early = list() - for jv,c in enumerate(contours): + for jv, c in enumerate(contours): if len(c) < 3: # A polygon cannot have less than 3 points continue polygon = geometry.Polygon([point[0] for point in c]) # area = cv2.contourArea(c) area = polygon.area - ##print(np.prod(thresh.shape[:2])) + # print(np.prod(thresh.shape[:2])) # Check that polygon has area greater than minimal area # print(hierarchy[0][jv][3],hierarchy ) if area >= min_area * np.prod(image.shape[:2]) and area <= max_area * np.prod(image.shape[:2]): # and hierarchy[0][jv][3]==-1 : @@ -73,6 +78,7 @@ def filter_contours_area_of_image_tables(image, contours, hierarchy, max_area, m found_polygons_early.append(np.array([[point] for point in polygon.exterior.coords], dtype=np.int32)) return found_polygons_early + def find_new_features_of_contours(contours_main): areas_main = np.array([cv2.contourArea(contours_main[j]) for j in range(len(contours_main))]) @@ -107,25 +113,27 @@ def find_new_features_of_contours(contours_main): # dis_x=np.abs(x_max_main-x_min_main) return cx_main, cy_main, x_min_main, x_max_main, y_min_main, y_max_main, y_corr_x_min_from_argmin -def find_features_of_contours(contours_main): - - areas_main=np.array([cv2.contourArea(contours_main[j]) for j in range(len(contours_main))]) - M_main=[cv2.moments(contours_main[j]) for j in range(len(contours_main))] - cx_main=[(M_main[j]['m10']/(M_main[j]['m00']+1e-32)) for j in range(len(M_main))] - cy_main=[(M_main[j]['m01']/(M_main[j]['m00']+1e-32)) for j in range(len(M_main))] - x_min_main=np.array([np.min(contours_main[j][:,0,0]) for j in range(len(contours_main))]) - x_max_main=np.array([np.max(contours_main[j][:,0,0]) for j in range(len(contours_main))]) - y_min_main=np.array([np.min(contours_main[j][:,0,1]) for j in range(len(contours_main))]) - y_max_main=np.array([np.max(contours_main[j][:,0,1]) for j in range(len(contours_main))]) +def find_features_of_contours(contours_main): + areas_main = np.array([cv2.contourArea(contours_main[j]) for j in range(len(contours_main))]) + M_main = [cv2.moments(contours_main[j]) for j in range(len(contours_main))] + cx_main = [(M_main[j]['m10']/(M_main[j]['m00']+1e-32)) for j in range(len(M_main))] + cy_main = [(M_main[j]['m01']/(M_main[j]['m00']+1e-32)) for j in range(len(M_main))] + x_min_main = np.array([np.min(contours_main[j][:, 0, 0]) for j in range(len(contours_main))]) + x_max_main = np.array([np.max(contours_main[j][:, 0, 0]) for j in range(len(contours_main))]) + + y_min_main = np.array([np.min(contours_main[j][:, 0, 1]) for j in range(len(contours_main))]) + y_max_main = np.array([np.max(contours_main[j][:, 0, 1]) for j in range(len(contours_main))]) - return y_min_main, y_max_main + + def return_parent_contours(contours, hierarchy): contours_parent = [contours[i] for i in range(len(contours)) if hierarchy[0][i][3] == -1] return contours_parent + def return_contours_of_interested_region(region_pre_p, pixel, min_area=0.0002): # pixels of images are identified by 5 @@ -145,6 +153,7 @@ def return_contours_of_interested_region(region_pre_p, pixel, min_area=0.0002): return contours_imgs + def do_work_of_contours_in_image(queue_of_all_params, contours_per_process, indexes_r_con_per_pro, img, slope_first): cnts_org_per_each_subprocess = [] index_by_text_region_contours = [] @@ -165,10 +174,9 @@ def do_work_of_contours_in_image(queue_of_all_params, contours_per_process, inde cont_int[0][:, 0, 0] = cont_int[0][:, 0, 0] + np.abs(img_copy.shape[1] - img.shape[1]) cont_int[0][:, 0, 1] = cont_int[0][:, 0, 1] + np.abs(img_copy.shape[0] - img.shape[0]) - cnts_org_per_each_subprocess.append(cont_int[0]) - queue_of_all_params.put([ cnts_org_per_each_subprocess, index_by_text_region_contours]) + queue_of_all_params.put([cnts_org_per_each_subprocess, index_by_text_region_contours]) def get_textregion_contours_in_org_image_multi(cnts, img, slope_first): @@ -180,10 +188,10 @@ def get_textregion_contours_in_org_image_multi(cnts, img, slope_first): nh = np.linspace(0, len(cnts), num_cores + 1) indexes_by_text_con = np.array(range(len(cnts))) for i in range(num_cores): - contours_per_process = cnts[int(nh[i]) : int(nh[i + 1])] - indexes_text_con_per_process = indexes_by_text_con[int(nh[i]) : int(nh[i + 1])] + contours_per_process = cnts[int(nh[i]): int(nh[i + 1])] + indexes_text_con_per_process = indexes_by_text_con[int(nh[i]): int(nh[i + 1])] - processes.append(Process(target=do_work_of_contours_in_image, args=(queue_of_all_params, contours_per_process, indexes_text_con_per_process, img,slope_first ))) + processes.append(Process(target=do_work_of_contours_in_image, args=(queue_of_all_params, contours_per_process, indexes_text_con_per_process, img, slope_first))) for i in range(num_cores): processes[i].start() cnts_org = [] @@ -200,7 +208,9 @@ def get_textregion_contours_in_org_image_multi(cnts, img, slope_first): print(all_index_text_con) return cnts_org -def loop_contour_image(index_l, cnts,img, slope_first): + + +def loop_contour_image(index_l, cnts, img, slope_first): img_copy = np.zeros(img.shape) img_copy = cv2.fillPoly(img_copy, pts=[cnts[index_l]], color=(1, 1, 1)) @@ -209,7 +219,7 @@ def loop_contour_image(index_l, cnts,img, slope_first): # print(img.shape,'img') img_copy = rotation_image_new(img_copy, -slope_first) - ##print(img_copy.shape,'img_copy') + # print(img_copy.shape,'img_copy') # plt.imshow(img_copy) # plt.show() @@ -224,15 +234,17 @@ def loop_contour_image(index_l, cnts,img, slope_first): # print(np.shape(cont_int[0])) return cont_int[0] + def get_textregion_contours_in_org_image_multi2(cnts, img, slope_first): cnts_org = [] # print(cnts,'cnts') with Pool(cpu_count()) as p: - cnts_org = p.starmap(loop_contour_image, [(index_l,cnts, img,slope_first) for index_l in range(len(cnts))]) + cnts_org = p.starmap(loop_contour_image, [(index_l, cnts, img, slope_first) for index_l in range(len(cnts))]) return cnts_org + def get_textregion_contours_in_org_image(cnts, img, slope_first): cnts_org = [] @@ -246,7 +258,7 @@ def get_textregion_contours_in_org_image(cnts, img, slope_first): # print(img.shape,'img') img_copy = rotation_image_new(img_copy, -slope_first) - ##print(img_copy.shape,'img_copy') + # print(img_copy.shape,'img_copy') # plt.imshow(img_copy) # plt.show() @@ -263,17 +275,18 @@ def get_textregion_contours_in_org_image(cnts, img, slope_first): return cnts_org + def get_textregion_contours_in_org_image_light(cnts, img, slope_first): h_o = img.shape[0] w_o = img.shape[1] img = cv2.resize(img, (int(img.shape[1]/3.), int(img.shape[0]/3.)), interpolation=cv2.INTER_NEAREST) - ##cnts = list( (np.array(cnts)/2).astype(np.int16) ) - #cnts = cnts/2 - cnts = [(i/ 3).astype(np.int32) for i in cnts] + # cnts = list( (np.array(cnts)/2).astype(np.int16) ) + # cnts = cnts/2 + cnts = [(i / 3).astype(np.int32) for i in cnts] cnts_org = [] - #print(cnts,'cnts') + # print(cnts,'cnts') for i in range(len(cnts)): img_copy = np.zeros(img.shape) img_copy = cv2.fillPoly(img_copy, pts=[cnts[i]], color=(1, 1, 1)) @@ -283,7 +296,7 @@ def get_textregion_contours_in_org_image_light(cnts, img, slope_first): # print(img.shape,'img') img_copy = rotation_image_new(img_copy, -slope_first) - ##print(img_copy.shape,'img_copy') + # print(img_copy.shape,'img_copy') # plt.imshow(img_copy) # plt.show() @@ -300,6 +313,7 @@ def get_textregion_contours_in_org_image_light(cnts, img, slope_first): return cnts_org + def return_contours_of_interested_textline(region_pre_p, pixel): # pixels of images are identified by 5 @@ -317,6 +331,7 @@ def return_contours_of_interested_textline(region_pre_p, pixel): contours_imgs = filter_contours_area_of_image_tables(thresh, contours_imgs, hierarchy, max_area=1, min_area=0.000000003) return contours_imgs + def return_contours_of_image(image): if len(image.shape) == 2: @@ -329,6 +344,7 @@ def return_contours_of_image(image): contours, hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) return contours, hierarchy + def return_contours_of_interested_region_by_min_size(region_pre_p, pixel, min_size=0.00003): # pixels of images are identified by 5 @@ -348,6 +364,7 @@ def return_contours_of_interested_region_by_min_size(region_pre_p, pixel, min_si return contours_imgs + def return_contours_of_interested_region_by_size(region_pre_p, pixel, min_area, max_area): # pixels of images are identified by 5 @@ -367,4 +384,3 @@ def return_contours_of_interested_region_by_size(region_pre_p, pixel, min_area, img_ret = np.zeros((region_pre_p.shape[0], region_pre_p.shape[1], 3)) img_ret = cv2.fillPoly(img_ret, pts=contours_imgs, color=(1, 1, 1)) return img_ret[:, :, 0] - diff --git a/eynollah/eynollah/utils/counter.py b/eynollah/eynollah/utils/counter.py index 9a3ed70..ac32dc9 100644 --- a/eynollah/eynollah/utils/counter.py +++ b/eynollah/eynollah/utils/counter.py @@ -3,6 +3,7 @@ REGION_ID_TEMPLATE = 'region_%04d' LINE_ID_TEMPLATE = 'region_%04d_line_%04d' + class EynollahIdCounter(): def __init__(self, region_idx=0, line_idx=0): diff --git a/eynollah/eynollah/utils/drop_capitals.py b/eynollah/eynollah/utils/drop_capitals.py index e12028f..d464c63 100644 --- a/eynollah/eynollah/utils/drop_capitals.py +++ b/eynollah/eynollah/utils/drop_capitals.py @@ -6,6 +6,7 @@ return_parent_contours, ) + def adhere_drop_capital_region_into_corresponding_textline( text_regions_p, polygons_of_drop_capitals, @@ -26,7 +27,7 @@ def adhere_drop_capital_region_into_corresponding_textline( img_con_all = np.zeros((text_regions_p.shape[0], text_regions_p.shape[1], 3)) for j_cont in range(len(contours_only_text_parent)): - img_con_all[all_box_coord[j_cont][0] : all_box_coord[j_cont][1], all_box_coord[j_cont][2] : all_box_coord[j_cont][3], 0] = (j_cont + 1) * 3 + img_con_all[all_box_coord[j_cont][0]: all_box_coord[j_cont][1], all_box_coord[j_cont][2]: all_box_coord[j_cont][3], 0] = (j_cont + 1) * 3 # img_con_all=cv2.fillPoly(img_con_all,pts=[contours_only_text_parent[j_cont]],color=((j_cont+1)*3,(j_cont+1)*3,(j_cont+1)*3)) # plt.imshow(img_con_all[:,:,0]) @@ -44,7 +45,7 @@ def adhere_drop_capital_region_into_corresponding_textline( # plt.imshow(img_con[:,:,0]) # plt.show() - ##img_con=cv2.dilate(img_con, kernel, iterations=30) + # img_con=cv2.dilate(img_con, kernel, iterations=30) # plt.imshow(img_con[:,:,0]) # plt.show() @@ -185,7 +186,7 @@ def adhere_drop_capital_region_into_corresponding_textline( # contours_biggest[:,0,1]=contours_biggest[:,0,1]#-all_box_coord[int(region_final)][0] # print(np.shape(contours_biggest),'contours_biggest') # print(np.shape(all_found_textline_polygons[int(region_final)][arg_min])) - ##contours_biggest=contours_biggest.reshape(np.shape(contours_biggest)[0],np.shape(contours_biggest)[2]) + # contours_biggest=contours_biggest.reshape(np.shape(contours_biggest)[0],np.shape(contours_biggest)[2]) all_found_textline_polygons[int(region_final)][arg_min] = contours_biggest except: pass @@ -230,7 +231,7 @@ def adhere_drop_capital_region_into_corresponding_textline( contours_biggest[:, 0, 0] = contours_biggest[:, 0, 0] # -all_box_coord[int(region_final)][2] contours_biggest[:, 0, 1] = contours_biggest[:, 0, 1] # -all_box_coord[int(region_final)][0] - ##contours_biggest=contours_biggest.reshape(np.shape(contours_biggest)[0],np.shape(contours_biggest)[2]) + # contours_biggest=contours_biggest.reshape(np.shape(contours_biggest)[0],np.shape(contours_biggest)[2]) all_found_textline_polygons[int(region_final)][arg_min] = contours_biggest # all_found_textline_polygons[int(region_final)][arg_min]=contours_biggest @@ -239,49 +240,49 @@ def adhere_drop_capital_region_into_corresponding_textline( else: pass - ##cx_t,cy_t ,_, _, _ ,_,_= find_new_features_of_contours(all_found_textline_polygons[int(region_final)]) - ###print(all_box_coord[j_cont]) - ###print(cx_t) - ###print(cy_t) - ###print(cx_d[i_drop]) - ###print(cy_d[i_drop]) - ##y_lines=all_box_coord[int(region_final)][0]+np.array(cy_t) + # cx_t,cy_t ,_, _, _ ,_,_= find_new_features_of_contours(all_found_textline_polygons[int(region_final)]) + # ##print(all_box_coord[j_cont]) + # ##print(cx_t) + # ##print(cy_t) + # ##print(cx_d[i_drop]) + # ##print(cy_d[i_drop]) + # #y_lines=all_box_coord[int(region_final)][0]+np.array(cy_t) - ##y_lines[y_lines 1: @@ -399,71 +400,72 @@ def adhere_drop_capital_region_into_corresponding_textline( else: pass - #####for i_drop in range(len(polygons_of_drop_capitals)): - #####for j_cont in range(len(contours_only_text_parent)): - #####img_con=np.zeros((text_regions_p.shape[0],text_regions_p.shape[1],3)) - #####img_con=cv2.fillPoly(img_con,pts=[polygons_of_drop_capitals[i_drop] ],color=(255,255,255)) - #####img_con=cv2.fillPoly(img_con,pts=[contours_only_text_parent[j_cont]],color=(255,255,255)) - - #####img_con=img_con.astype(np.uint8) - ######imgray = cv2.cvtColor(img_con, cv2.COLOR_BGR2GRAY) - ######ret, thresh = cv2.threshold(imgray, 0, 255, 0) - - ######contours_new,hierarchy=cv2.findContours(thresh,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE) - - #####contours_new,hir_new=return_contours_of_image(img_con) - #####contours_new_parent=return_parent_contours( contours_new,hir_new) - ######plt.imshow(img_con) - ######plt.show() - #####try: - #####if len(contours_new_parent)==1: - ######print(all_found_textline_polygons[j_cont][0]) - #####cx_t,cy_t ,_, _, _ ,_,_= find_new_features_of_contours(all_found_textline_polygons[j_cont]) - ######print(all_box_coord[j_cont]) - ######print(cx_t) - ######print(cy_t) - ######print(cx_d[i_drop]) - ######print(cy_d[i_drop]) - #####y_lines=all_box_coord[j_cont][0]+np.array(cy_t) - - ######print(y_lines) - - #####arg_min=np.argmin(np.abs(y_lines-y_min_d[i_drop]) ) - ######print(arg_min) - - #####cnt_nearest=np.copy(all_found_textline_polygons[j_cont][arg_min]) - #####cnt_nearest[:,0]=all_found_textline_polygons[j_cont][arg_min][:,0]+all_box_coord[j_cont][2] - #####cnt_nearest[:,1]=all_found_textline_polygons[j_cont][arg_min][:,1]+all_box_coord[j_cont][0] - - #####img_textlines=np.zeros((text_regions_p.shape[0],text_regions_p.shape[1],3)) - #####img_textlines=cv2.fillPoly(img_textlines,pts=[cnt_nearest],color=(255,255,255)) - #####img_textlines=cv2.fillPoly(img_textlines,pts=[polygons_of_drop_capitals[i_drop] ],color=(255,255,255)) - - #####img_textlines=img_textlines.astype(np.uint8) - #####imgray = cv2.cvtColor(img_textlines, cv2.COLOR_BGR2GRAY) - #####ret, thresh = cv2.threshold(imgray, 0, 255, 0) - - #####contours_combined,hierarchy=cv2.findContours(thresh,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE) - - #####areas_cnt_text=np.array([cv2.contourArea(contours_combined[j]) for j in range(len(contours_combined))]) - - #####contours_biggest=contours_combined[np.argmax(areas_cnt_text)] - - ######print(np.shape(contours_biggest)) - ######print(contours_biggest[:]) - #####contours_biggest[:,0,0]=contours_biggest[:,0,0]-all_box_coord[j_cont][2] - #####contours_biggest[:,0,1]=contours_biggest[:,0,1]-all_box_coord[j_cont][0] - - #####all_found_textline_polygons[j_cont][arg_min]=contours_biggest - ######print(contours_biggest) - ######plt.imshow(img_textlines[:,:,0]) - ######plt.show() - #####else: - #####pass - #####except: - #####pass + # ####for i_drop in range(len(polygons_of_drop_capitals)): + # ####for j_cont in range(len(contours_only_text_parent)): + # ####img_con=np.zeros((text_regions_p.shape[0],text_regions_p.shape[1],3)) + # ####img_con=cv2.fillPoly(img_con,pts=[polygons_of_drop_capitals[i_drop] ],color=(255,255,255)) + # ####img_con=cv2.fillPoly(img_con,pts=[contours_only_text_parent[j_cont]],color=(255,255,255)) + + # ####img_con=img_con.astype(np.uint8) + # #####imgray = cv2.cvtColor(img_con, cv2.COLOR_BGR2GRAY) + # #####ret, thresh = cv2.threshold(imgray, 0, 255, 0) + + # #####contours_new,hierarchy=cv2.findContours(thresh,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE) + + # ####contours_new,hir_new=return_contours_of_image(img_con) + # ####contours_new_parent=return_parent_contours( contours_new,hir_new) + # #####plt.imshow(img_con) + # #####plt.show() + # ####try: + # ####if len(contours_new_parent)==1: + # #####print(all_found_textline_polygons[j_cont][0]) + # ####cx_t,cy_t ,_, _, _ ,_,_= find_new_features_of_contours(all_found_textline_polygons[j_cont]) + # #####print(all_box_coord[j_cont]) + # #####print(cx_t) + # #####print(cy_t) + # #####print(cx_d[i_drop]) + # #####print(cy_d[i_drop]) + # ####y_lines=all_box_coord[j_cont][0]+np.array(cy_t) + + # #####print(y_lines) + + # ####arg_min=np.argmin(np.abs(y_lines-y_min_d[i_drop]) ) + # #####print(arg_min) + + # ####cnt_nearest=np.copy(all_found_textline_polygons[j_cont][arg_min]) + # ####cnt_nearest[:,0]=all_found_textline_polygons[j_cont][arg_min][:,0]+all_box_coord[j_cont][2] + # ####cnt_nearest[:,1]=all_found_textline_polygons[j_cont][arg_min][:,1]+all_box_coord[j_cont][0] + + # ####img_textlines=np.zeros((text_regions_p.shape[0],text_regions_p.shape[1],3)) + # ####img_textlines=cv2.fillPoly(img_textlines,pts=[cnt_nearest],color=(255,255,255)) + # ####img_textlines=cv2.fillPoly(img_textlines,pts=[polygons_of_drop_capitals[i_drop] ],color=(255,255,255)) + + # ####img_textlines=img_textlines.astype(np.uint8) + # ####imgray = cv2.cvtColor(img_textlines, cv2.COLOR_BGR2GRAY) + # ####ret, thresh = cv2.threshold(imgray, 0, 255, 0) + + # ####contours_combined,hierarchy=cv2.findContours(thresh,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE) + + # ####areas_cnt_text=np.array([cv2.contourArea(contours_combined[j]) for j in range(len(contours_combined))]) + + # ####contours_biggest=contours_combined[np.argmax(areas_cnt_text)] + + # #####print(np.shape(contours_biggest)) + # #####print(contours_biggest[:]) + # ####contours_biggest[:,0,0]=contours_biggest[:,0,0]-all_box_coord[j_cont][2] + # ####contours_biggest[:,0,1]=contours_biggest[:,0,1]-all_box_coord[j_cont][0] + + # ####all_found_textline_polygons[j_cont][arg_min]=contours_biggest + # #####print(contours_biggest) + # #####plt.imshow(img_textlines[:,:,0]) + # #####plt.show() + # ####else: + # ####pass + # ####except: + # ####pass return all_found_textline_polygons + def filter_small_drop_capitals_from_no_patch_layout(layout_no_patch, layout1): drop_only = (layout_no_patch[:, :, 0] == 4) * 1 @@ -489,7 +491,7 @@ def filter_small_drop_capitals_from_no_patch_layout(layout_no_patch, layout1): if iou_of_box_and_contoure > 60 and weigh_to_height_ratio < 1.2 and height_to_weight_ratio < 2: map_of_drop_contour_bb = np.zeros((layout1.shape[0], layout1.shape[1])) - map_of_drop_contour_bb[y : y + h, x : x + w] = layout1[y : y + h, x : x + w] + map_of_drop_contour_bb[y: y + h, x: x + w] = layout1[y: y + h, x: x + w] if (((map_of_drop_contour_bb == 1) * 1).sum() / float(((map_of_drop_contour_bb == 5) * 1).sum()) * 100) >= 15: contours_drop_parent_final.append(contours_drop_parent[jj]) @@ -499,4 +501,3 @@ def filter_small_drop_capitals_from_no_patch_layout(layout_no_patch, layout1): layout_no_patch = cv2.fillPoly(layout_no_patch, pts=contours_drop_parent_final, color=(4, 4, 4)) return layout_no_patch - diff --git a/eynollah/eynollah/utils/marginals.py b/eynollah/eynollah/utils/marginals.py index 7c43de6..d7514bb 100644 --- a/eynollah/eynollah/utils/marginals.py +++ b/eynollah/eynollah/utils/marginals.py @@ -3,250 +3,226 @@ from scipy.signal import find_peaks from scipy.ndimage import gaussian_filter1d - from .contour import find_new_features_of_contours, return_contours_of_interested_region from .resize import resize_image from .rotate import rotate_image -def get_marginals(text_with_lines, text_regions, num_col, slope_deskew, kernel=None): - mask_marginals=np.zeros((text_with_lines.shape[0],text_with_lines.shape[1])) - mask_marginals=mask_marginals.astype(np.uint8) +def get_marginals(text_with_lines, text_regions, num_col, slope_deskew, kernel=None): + mask_marginals = np.zeros((text_with_lines.shape[0], text_with_lines.shape[1])) + mask_marginals = mask_marginals.astype(np.uint8) - text_with_lines=text_with_lines.astype(np.uint8) - ##text_with_lines=cv2.erode(text_with_lines,self.kernel,iterations=3) + text_with_lines = text_with_lines.astype(np.uint8) + # text_with_lines=cv2.erode(text_with_lines,self.kernel,iterations=3) - text_with_lines_eroded=cv2.erode(text_with_lines,kernel,iterations=5) + text_with_lines_eroded = cv2.erode(text_with_lines, kernel, iterations=5) - if text_with_lines.shape[0]<=1500: + if text_with_lines.shape[0] <= 1500: pass - elif text_with_lines.shape[0]>1500 and text_with_lines.shape[0]<=1800: - text_with_lines=resize_image(text_with_lines,int(text_with_lines.shape[0]*1.5),text_with_lines.shape[1]) - text_with_lines=cv2.erode(text_with_lines,kernel,iterations=5) - text_with_lines=resize_image(text_with_lines,text_with_lines_eroded.shape[0],text_with_lines_eroded.shape[1]) + elif text_with_lines.shape[0] > 1500 and text_with_lines.shape[0] <= 1800: + text_with_lines = resize_image(text_with_lines, int(text_with_lines.shape[0] * 1.5), text_with_lines.shape[1]) + text_with_lines = cv2.erode(text_with_lines, kernel, iterations=5) + text_with_lines = resize_image(text_with_lines, text_with_lines_eroded.shape[0], + text_with_lines_eroded.shape[1]) else: - text_with_lines=resize_image(text_with_lines,int(text_with_lines.shape[0]*1.8),text_with_lines.shape[1]) - text_with_lines=cv2.erode(text_with_lines,kernel,iterations=7) - text_with_lines=resize_image(text_with_lines,text_with_lines_eroded.shape[0],text_with_lines_eroded.shape[1]) - + text_with_lines = resize_image(text_with_lines, int(text_with_lines.shape[0] * 1.8), text_with_lines.shape[1]) + text_with_lines = cv2.erode(text_with_lines, kernel, iterations=7) + text_with_lines = resize_image(text_with_lines, text_with_lines_eroded.shape[0], + text_with_lines_eroded.shape[1]) - text_with_lines_y=text_with_lines.sum(axis=0) - text_with_lines_y_eroded=text_with_lines_eroded.sum(axis=0) + text_with_lines_y = text_with_lines.sum(axis=0) + text_with_lines_y_eroded = text_with_lines_eroded.sum(axis=0) - thickness_along_y_percent=text_with_lines_y_eroded.max()/(float(text_with_lines.shape[0]))*100 + thickness_along_y_percent = text_with_lines_y_eroded.max() / (float(text_with_lines.shape[0])) * 100 - #print(thickness_along_y_percent,'thickness_along_y_percent') + # print(thickness_along_y_percent,'thickness_along_y_percent') - if thickness_along_y_percent<30: - min_textline_thickness=8 - elif thickness_along_y_percent>=30 and thickness_along_y_percent<50: - min_textline_thickness=20 + if thickness_along_y_percent < 30: + min_textline_thickness = 8 + elif thickness_along_y_percent >= 30 and thickness_along_y_percent < 50: + min_textline_thickness = 20 else: - min_textline_thickness=40 - - - - if thickness_along_y_percent>=14: - - text_with_lines_y_rev=-1*text_with_lines_y[:] - #print(text_with_lines_y) - #print(text_with_lines_y_rev) - - - - - #plt.plot(text_with_lines_y) - #plt.show() - - - text_with_lines_y_rev=text_with_lines_y_rev-np.min(text_with_lines_y_rev) - - #plt.plot(text_with_lines_y_rev) - #plt.show() - sigma_gaus=1 - region_sum_0= gaussian_filter1d(text_with_lines_y, sigma_gaus) + min_textline_thickness = 40 - region_sum_0_rev=gaussian_filter1d(text_with_lines_y_rev, sigma_gaus) + if thickness_along_y_percent >= 14: - #plt.plot(region_sum_0_rev) - #plt.show() - region_sum_0_updown=region_sum_0[len(region_sum_0)::-1] + text_with_lines_y_rev = -1 * text_with_lines_y[:] + # print(text_with_lines_y) + # print(text_with_lines_y_rev) - first_nonzero=(next((i for i, x in enumerate(region_sum_0) if x), None)) - last_nonzero=(next((i for i, x in enumerate(region_sum_0_updown) if x), None)) + # plt.plot(text_with_lines_y) + # plt.show() + text_with_lines_y_rev = text_with_lines_y_rev - np.min(text_with_lines_y_rev) - last_nonzero=len(region_sum_0)-last_nonzero + # plt.plot(text_with_lines_y_rev) + # plt.show() + sigma_gaus = 1 + region_sum_0 = gaussian_filter1d(text_with_lines_y, sigma_gaus) - ##img_sum_0_smooth_rev=-region_sum_0 + region_sum_0_rev = gaussian_filter1d(text_with_lines_y_rev, sigma_gaus) + # plt.plot(region_sum_0_rev) + # plt.show() + region_sum_0_updown = region_sum_0[len(region_sum_0)::-1] - mid_point=(last_nonzero+first_nonzero)/2. + first_nonzero = (next((i for i, x in enumerate(region_sum_0) if x), None)) + last_nonzero = (next((i for i, x in enumerate(region_sum_0_updown) if x), None)) + last_nonzero = len(region_sum_0) - last_nonzero - one_third_right=(last_nonzero-mid_point)/3.0 - one_third_left=(mid_point-first_nonzero)/3.0 - - #img_sum_0_smooth_rev=img_sum_0_smooth_rev-np.min(img_sum_0_smooth_rev) + # img_sum_0_smooth_rev=-region_sum_0 + mid_point = (last_nonzero + first_nonzero) / 2. + one_third_right = (last_nonzero - mid_point) / 3.0 + one_third_left = (mid_point - first_nonzero) / 3.0 + # img_sum_0_smooth_rev=img_sum_0_smooth_rev-np.min(img_sum_0_smooth_rev) peaks, _ = find_peaks(text_with_lines_y_rev, height=0) + peaks = np.array(peaks) - peaks=np.array(peaks) - - - #print(region_sum_0[peaks]) - ##plt.plot(region_sum_0) - ##plt.plot(peaks,region_sum_0[peaks],'*') - ##plt.show() - #print(first_nonzero,last_nonzero,peaks) - peaks=peaks[(peaks>first_nonzero) & ((peaks first_nonzero) & (peaks < last_nonzero)] + # print(first_nonzero,last_nonzero,peaks) - #print(region_sum_0[peaks]<10) - ####peaks=peaks[region_sum_0[peaks]<25 ] + # print(region_sum_0[peaks]<10) + # ###peaks=peaks[region_sum_0[peaks]<25 ] - #print(region_sum_0[peaks]) - peaks=peaks[region_sum_0[peaks]mid_point] - peaks_left=peaks[peaks(mid_point+one_third_right)] - peaks_left=peaks[peaks<(mid_point-one_third_left)] + # print(region_sum_0[peaks]) + peaks = peaks[region_sum_0[peaks] < min_textline_thickness] + # print(peaks) + # print(first_nonzero,last_nonzero,one_third_right,one_third_left) + if num_col == 1: + peaks_right = peaks[peaks > mid_point] + peaks_left = peaks[peaks < mid_point] + if num_col == 2: + peaks_right = peaks[peaks > (mid_point + one_third_right)] + peaks_left = peaks[peaks < (mid_point - one_third_left)] try: - point_right=np.min(peaks_right) + point_right = np.min(peaks_right) except: - point_right=last_nonzero - + point_right = last_nonzero try: - point_left=np.max(peaks_left) + point_left = np.max(peaks_left) except: - point_left=first_nonzero - + point_left = first_nonzero - - - #print(point_left,point_right) - #print(text_regions.shape) - if point_right>=mask_marginals.shape[1]: - point_right=mask_marginals.shape[1]-1 + # print(point_left,point_right) + # print(text_regions.shape) + if point_right >= mask_marginals.shape[1]: + point_right = mask_marginals.shape[1] - 1 try: - mask_marginals[:,point_left:point_right]=1 + mask_marginals[:, point_left:point_right] = 1 except: - mask_marginals[:,:]=1 + mask_marginals[:, :] = 1 - #print(mask_marginals.shape,point_left,point_right,'nadosh') - mask_marginals_rotated=rotate_image(mask_marginals,-slope_deskew) + # print(mask_marginals.shape,point_left,point_right,'nadosh') + mask_marginals_rotated = rotate_image(mask_marginals, -slope_deskew) - #print(mask_marginals_rotated.shape,'nadosh') - mask_marginals_rotated_sum=mask_marginals_rotated.sum(axis=0) + # print(mask_marginals_rotated.shape,'nadosh') + mask_marginals_rotated_sum = mask_marginals_rotated.sum(axis=0) - mask_marginals_rotated_sum[mask_marginals_rotated_sum!=0]=1 - index_x=np.array(range(len(mask_marginals_rotated_sum)))+1 + mask_marginals_rotated_sum[mask_marginals_rotated_sum != 0] = 1 + index_x = np.array(range(len(mask_marginals_rotated_sum))) + 1 - index_x_interest=index_x[mask_marginals_rotated_sum==1] + index_x_interest = index_x[mask_marginals_rotated_sum == 1] - min_point_of_left_marginal=np.min(index_x_interest)-16 - max_point_of_right_marginal=np.max(index_x_interest)+16 + min_point_of_left_marginal = np.min(index_x_interest) - 16 + max_point_of_right_marginal = np.max(index_x_interest) + 16 - if min_point_of_left_marginal<0: - min_point_of_left_marginal=0 - if max_point_of_right_marginal>=text_regions.shape[1]: - max_point_of_right_marginal=text_regions.shape[1]-1 + if min_point_of_left_marginal < 0: + min_point_of_left_marginal = 0 + if max_point_of_right_marginal >= text_regions.shape[1]: + max_point_of_right_marginal = text_regions.shape[1] - 1 + # print(np.min(index_x_interest) ,np.max(index_x_interest),'minmaxnew') + # print(mask_marginals_rotated.shape,text_regions.shape,'mask_marginals_rotated') + # plt.imshow(mask_marginals) + # plt.show() - #print(np.min(index_x_interest) ,np.max(index_x_interest),'minmaxnew') - #print(mask_marginals_rotated.shape,text_regions.shape,'mask_marginals_rotated') - #plt.imshow(mask_marginals) - #plt.show() + # plt.imshow(mask_marginals_rotated) + # plt.show() - #plt.imshow(mask_marginals_rotated) - #plt.show() + text_regions[(mask_marginals_rotated[:, :] != 1) & (text_regions[:, :] == 1)] = 4 - text_regions[(mask_marginals_rotated[:,:]!=1) & (text_regions[:,:]==1)]=4 + # plt.imshow(text_regions) + # plt.show() - #plt.imshow(text_regions) - #plt.show() + pixel_img = 4 + min_area_text = 0.00001 + polygons_of_marginals = return_contours_of_interested_region(text_regions, pixel_img, min_area_text) - pixel_img=4 - min_area_text=0.00001 - polygons_of_marginals=return_contours_of_interested_region(text_regions,pixel_img,min_area_text) + cx_text_only, cy_text_only, x_min_text_only, x_max_text_only, y_min_text_only, y_max_text_only, y_cor_x_min_main = find_new_features_of_contours( + polygons_of_marginals) - cx_text_only,cy_text_only ,x_min_text_only,x_max_text_only, y_min_text_only ,y_max_text_only,y_cor_x_min_main=find_new_features_of_contours(polygons_of_marginals) + text_regions[(text_regions[:, :] == 4)] = 1 - text_regions[(text_regions[:,:]==4)]=1 + marginlas_should_be_main_text = [] - marginlas_should_be_main_text=[] - - x_min_marginals_left=[] - x_min_marginals_right=[] + x_min_marginals_left = [] + x_min_marginals_right = [] for i in range(len(cx_text_only)): - x_width_mar=abs(x_min_text_only[i]-x_max_text_only[i]) - y_height_mar=abs(y_min_text_only[i]-y_max_text_only[i]) - #print(x_width_mar,y_height_mar,y_height_mar/x_width_mar,'y_height_mar') - if x_width_mar>16 and y_height_mar/x_width_mar<18: + x_width_mar = abs(x_min_text_only[i] - x_max_text_only[i]) + y_height_mar = abs(y_min_text_only[i] - y_max_text_only[i]) + # print(x_width_mar,y_height_mar,y_height_mar/x_width_mar,'y_height_mar') + if x_width_mar > 16 and y_height_mar / x_width_mar < 18: marginlas_should_be_main_text.append(polygons_of_marginals[i]) - if x_min_text_only[i]<(mid_point-one_third_left): - x_min_marginals_left_new=x_min_text_only[i] - if len(x_min_marginals_left)==0: + if x_min_text_only[i] < (mid_point - one_third_left): + x_min_marginals_left_new = x_min_text_only[i] + if len(x_min_marginals_left) == 0: x_min_marginals_left.append(x_min_marginals_left_new) else: - x_min_marginals_left[0]=min(x_min_marginals_left[0],x_min_marginals_left_new) + x_min_marginals_left[0] = min(x_min_marginals_left[0], x_min_marginals_left_new) else: - x_min_marginals_right_new=x_min_text_only[i] - if len(x_min_marginals_right)==0: + x_min_marginals_right_new = x_min_text_only[i] + if len(x_min_marginals_right) == 0: x_min_marginals_right.append(x_min_marginals_right_new) else: - x_min_marginals_right[0]=min(x_min_marginals_right[0],x_min_marginals_right_new) - - if len(x_min_marginals_left)==0: - x_min_marginals_left=[0] - if len(x_min_marginals_right)==0: - x_min_marginals_right=[text_regions.shape[1]-1] - - - + x_min_marginals_right[0] = min(x_min_marginals_right[0], x_min_marginals_right_new) - #print(x_min_marginals_left[0],x_min_marginals_right[0],'margo') + if len(x_min_marginals_left) == 0: + x_min_marginals_left = [0] + if len(x_min_marginals_right) == 0: + x_min_marginals_right = [text_regions.shape[1] - 1] - #print(marginlas_should_be_main_text,'marginlas_should_be_main_text') - text_regions=cv2.fillPoly(text_regions, pts =marginlas_should_be_main_text, color=(4,4)) + # print(x_min_marginals_left[0],x_min_marginals_right[0],'margo') - #print(np.unique(text_regions)) + # print(marginlas_should_be_main_text,'marginlas_should_be_main_text') + text_regions = cv2.fillPoly(text_regions, pts=marginlas_should_be_main_text, color=(4, 4)) - #text_regions[:,:int(x_min_marginals_left[0])][text_regions[:,:int(x_min_marginals_left[0])]==1]=0 - #text_regions[:,int(x_min_marginals_right[0]):][text_regions[:,int(x_min_marginals_right[0]):]==1]=0 + # print(np.unique(text_regions)) - text_regions[:,:int(min_point_of_left_marginal)][text_regions[:,:int(min_point_of_left_marginal)]==1]=0 - text_regions[:,int(max_point_of_right_marginal):][text_regions[:,int(max_point_of_right_marginal):]==1]=0 + # text_regions[:,:int(x_min_marginals_left[0])][text_regions[:,:int(x_min_marginals_left[0])]==1]=0 + # text_regions[:,int(x_min_marginals_right[0]):][text_regions[:,int(x_min_marginals_right[0]):]==1]=0 - ###text_regions[:,0:point_left][text_regions[:,0:point_left]==1]=4 + text_regions[:, :int(min_point_of_left_marginal)][text_regions[:, :int(min_point_of_left_marginal)] == 1] = 0 + text_regions[:, int(max_point_of_right_marginal):][text_regions[:, int(max_point_of_right_marginal):] == 1] = 0 - ###text_regions[:,point_right:][ text_regions[:,point_right:]==1]=4 - #plt.plot(region_sum_0) - #plt.plot(peaks,region_sum_0[peaks],'*') - #plt.show() + # ##text_regions[:,0:point_left][text_regions[:,0:point_left]==1]=4 + # ##text_regions[:,point_right:][ text_regions[:,point_right:]==1]=4 + # plt.plot(region_sum_0) + # plt.plot(peaks,region_sum_0[peaks],'*') + # plt.show() - #plt.imshow(text_regions) - #plt.show() + # plt.imshow(text_regions) + # plt.show() - #sys.exit() + # sys.exit() else: pass return text_regions diff --git a/eynollah/eynollah/utils/pil_cv2.py b/eynollah/eynollah/utils/pil_cv2.py index 83ae47d..34ef9e1 100644 --- a/eynollah/eynollah/utils/pil_cv2.py +++ b/eynollah/eynollah/utils/pil_cv2.py @@ -5,15 +5,18 @@ # from sbb_binarization + def cv2pil(img): return Image.fromarray(np.array(cvtColor(img, COLOR_BGR2RGB))) + def pil2cv(img): # from ocrd/workspace.py - color_conversion = COLOR_GRAY2BGR if img.mode in ('1', 'L') else COLOR_RGB2BGR + color_conversion = COLOR_GRAY2BGR if img.mode in ('1', 'L') else COLOR_RGB2BGR pil_as_np_array = np.array(img).astype('uint8') if img.mode == '1' else np.array(img) return cvtColor(pil_as_np_array, color_conversion) + def check_dpi(img): try: if isinstance(img, Image.Image): diff --git a/eynollah/eynollah/utils/resize.py b/eynollah/eynollah/utils/resize.py index fdc49ec..8c09b04 100644 --- a/eynollah/eynollah/utils/resize.py +++ b/eynollah/eynollah/utils/resize.py @@ -1,4 +1,5 @@ import cv2 + def resize_image(img_in, input_height, input_width): return cv2.resize(img_in, (input_width, input_height), interpolation=cv2.INTER_NEAREST) diff --git a/eynollah/eynollah/utils/rotate.py b/eynollah/eynollah/utils/rotate.py index 603c2d9..4b6fbb6 100644 --- a/eynollah/eynollah/utils/rotate.py +++ b/eynollah/eynollah/utils/rotate.py @@ -3,6 +3,7 @@ import imutils import cv2 + def rotatedRectWithMaxArea(w, h, angle): if w <= 0 or h <= 0: return 0, 0 @@ -25,6 +26,7 @@ def rotatedRectWithMaxArea(w, h, angle): return wr, hr + def rotate_max_area_new(image, rotated, angle): wr, hr = rotatedRectWithMaxArea(image.shape[1], image.shape[0], math.radians(angle)) h, w, _ = rotated.shape @@ -34,17 +36,20 @@ def rotate_max_area_new(image, rotated, angle): x2 = x1 + int(wr) return rotated[y1:y2, x1:x2] + def rotation_image_new(img, thetha): rotated = imutils.rotate(img, thetha) return rotate_max_area_new(img, rotated, thetha) + def rotate_image(img_patch, slope): (h, w) = img_patch.shape[:2] center = (w // 2, h // 2) M = cv2.getRotationMatrix2D(center, slope, 1.0) return cv2.warpAffine(img_patch, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE) -def rotate_image_different( img, slope): + +def rotate_image_different(img, slope): # img = cv2.imread('images/input.jpg') num_rows, num_cols = img.shape[:2] @@ -52,6 +57,7 @@ def rotate_image_different( img, slope): img_rotation = cv2.warpAffine(img, rotation_matrix, (num_cols, num_rows)) return img_rotation + def rotate_max_area(image, rotated, rotated_textline, rotated_layout, rotated_table_prediction, angle): wr, hr = rotatedRectWithMaxArea(image.shape[1], image.shape[0], math.radians(angle)) h, w, _ = rotated.shape @@ -61,6 +67,7 @@ def rotate_max_area(image, rotated, rotated_textline, rotated_layout, rotated_ta x2 = x1 + int(wr) return rotated[y1:y2, x1:x2], rotated_textline[y1:y2, x1:x2], rotated_layout[y1:y2, x1:x2], rotated_table_prediction[y1:y2, x1:x2] + def rotation_not_90_func(img, textline, text_regions_p_1, table_prediction, thetha): rotated = imutils.rotate(img, thetha) rotated_textline = imutils.rotate(textline, thetha) @@ -68,6 +75,7 @@ def rotation_not_90_func(img, textline, text_regions_p_1, table_prediction, thet rotated_table_prediction = imutils.rotate(table_prediction, thetha) return rotate_max_area(img, rotated, rotated_textline, rotated_layout, rotated_table_prediction, thetha) + def rotation_not_90_func_full_layout(img, textline, text_regions_p_1, text_regions_p_fully, thetha): rotated = imutils.rotate(img, thetha) rotated_textline = imutils.rotate(textline, thetha) @@ -75,6 +83,7 @@ def rotation_not_90_func_full_layout(img, textline, text_regions_p_1, text_regio rotated_layout_full = imutils.rotate(text_regions_p_fully, thetha) return rotate_max_area_full_layout(img, rotated, rotated_textline, rotated_layout, rotated_layout_full, thetha) + def rotate_max_area_full_layout(image, rotated, rotated_textline, rotated_layout, rotated_layout_full, angle): wr, hr = rotatedRectWithMaxArea(image.shape[1], image.shape[0], math.radians(angle)) h, w, _ = rotated.shape @@ -83,4 +92,3 @@ def rotate_max_area_full_layout(image, rotated, rotated_textline, rotated_layout x1 = w // 2 - int(wr / 2) x2 = x1 + int(wr) return rotated[y1:y2, x1:x2], rotated_textline[y1:y2, x1:x2], rotated_layout[y1:y2, x1:x2], rotated_layout_full[y1:y2, x1:x2] - diff --git a/eynollah/eynollah/utils/xml.py b/eynollah/eynollah/utils/xml.py index 0386b25..009edd2 100644 --- a/eynollah/eynollah/utils/xml.py +++ b/eynollah/eynollah/utils/xml.py @@ -29,6 +29,7 @@ to_xml) + def create_page_xml(imageFilename, height, width): now = datetime.now() pcgts = PcGtsType( @@ -46,6 +47,7 @@ def create_page_xml(imageFilename, height, width): )) return pcgts + def xml_reading_order(page, order_of_texts, id_of_marginalia): region_order = ReadingOrderType() og = OrderedGroupType(id="ro357564684568544579089") @@ -59,6 +61,7 @@ def xml_reading_order(page, order_of_texts, id_of_marginalia): og.add_RegionRefIndexed(RegionRefIndexedType(index=str(region_counter.get('region')), regionRef=id_marginal)) region_counter.inc('region') + def order_and_id_of_texts(found_polygons_text_region, found_polygons_text_region_h, matrix_of_orders, indexes_sorted, index_of_types, kind_of_texts, ref_point): indexes_sorted = np.array(indexes_sorted) index_of_types = np.array(index_of_types) diff --git a/eynollah/eynollah/writer.py b/eynollah/eynollah/writer.py index f537f65..ee81a6a 100644 --- a/eynollah/eynollah/writer.py +++ b/eynollah/eynollah/writer.py @@ -8,21 +8,22 @@ from ocrd_utils import getLogger from ocrd_models.ocrd_page import ( - BorderType, - CoordsType, - PcGtsType, - TextLineType, - TextRegionType, - ImageRegionType, - TableRegionType, - SeparatorRegionType, - to_xml - ) + BorderType, + CoordsType, + PcGtsType, + TextLineType, + TextRegionType, + ImageRegionType, + TableRegionType, + SeparatorRegionType, + to_xml +) import numpy as np + class EynollahXmlWriter(): - def __init__(self, *, dir_out, image_filename, curved_line,textline_light, pcgts=None): + def __init__(self, *, dir_out, image_filename, curved_line, textline_light, pcgts=None): self.logger = getLogger('eynollah.writer') self.counter = EynollahIdCounter() self.dir_out = dir_out @@ -30,10 +31,10 @@ def __init__(self, *, dir_out, image_filename, curved_line,textline_light, pcgts self.curved_line = curved_line self.textline_light = textline_light self.pcgts = pcgts - self.scale_x = None # XXX set outside __init__ - self.scale_y = None # XXX set outside __init__ - self.height_org = None # XXX set outside __init__ - self.width_org = None # XXX set outside __init__ + self.scale_x = None # XXX set outside __init__ + self.scale_y = None # XXX set outside __init__ + self.height_org = None # XXX set outside __init__ + self.width_org = None # XXX set outside __init__ @property def image_filename_stem(self): @@ -50,11 +51,12 @@ def calculate_page_coords(self, cont_page): else: points_page_print += str(int((contour[0][0]) / self.scale_x)) points_page_print += ',' - points_page_print += str(int((contour[0][1] ) / self.scale_y)) + points_page_print += str(int((contour[0][1]) / self.scale_y)) points_page_print = points_page_print + ' ' return points_page_print[:-1] - def serialize_lines_in_marginal(self, marginal_region, all_found_textline_polygons_marginals, marginal_idx, page_coord, all_box_coord_marginals, slopes_marginals, counter): + def serialize_lines_in_marginal(self, marginal_region, all_found_textline_polygons_marginals, marginal_idx, + page_coord, all_box_coord_marginals, slopes_marginals, counter): for j in range(len(all_found_textline_polygons_marginals[marginal_idx])): coords = CoordsType() textline = TextLineType(id=counter.next_line_id, Coords=coords) @@ -63,37 +65,54 @@ def serialize_lines_in_marginal(self, marginal_region, all_found_textline_polygo for l in range(len(all_found_textline_polygons_marginals[marginal_idx][j])): if not (self.curved_line or self.textline_light): if len(all_found_textline_polygons_marginals[marginal_idx][j][l]) == 2: - textline_x_coord = max(0, int((all_found_textline_polygons_marginals[marginal_idx][j][l][0] + all_box_coord_marginals[marginal_idx][2] + page_coord[2]) / self.scale_x) ) - textline_y_coord = max(0, int((all_found_textline_polygons_marginals[marginal_idx][j][l][1] + all_box_coord_marginals[marginal_idx][0] + page_coord[0]) / self.scale_y) ) + textline_x_coord = max(0, int((all_found_textline_polygons_marginals[marginal_idx][j][l][0] + + all_box_coord_marginals[marginal_idx][2] + page_coord[ + 2]) / self.scale_x)) + textline_y_coord = max(0, int((all_found_textline_polygons_marginals[marginal_idx][j][l][1] + + all_box_coord_marginals[marginal_idx][0] + page_coord[ + 0]) / self.scale_y)) else: - textline_x_coord = max(0, int((all_found_textline_polygons_marginals[marginal_idx][j][l][0][0] + all_box_coord_marginals[marginal_idx][2] + page_coord[2]) / self.scale_x) ) - textline_y_coord = max(0, int((all_found_textline_polygons_marginals[marginal_idx][j][l][0][1] + all_box_coord_marginals[marginal_idx][0] + page_coord[0]) / self.scale_y) ) + textline_x_coord = max(0, int((all_found_textline_polygons_marginals[marginal_idx][j][l][0][0] + + all_box_coord_marginals[marginal_idx][2] + page_coord[ + 2]) / self.scale_x)) + textline_y_coord = max(0, int((all_found_textline_polygons_marginals[marginal_idx][j][l][0][1] + + all_box_coord_marginals[marginal_idx][0] + page_coord[ + 0]) / self.scale_y)) points_co += str(textline_x_coord) points_co += ',' points_co += str(textline_y_coord) if (self.curved_line or self.textline_light) and np.abs(slopes_marginals[marginal_idx]) <= 45: if len(all_found_textline_polygons_marginals[marginal_idx][j][l]) == 2: - points_co += str(int((all_found_textline_polygons_marginals[marginal_idx][j][l][0] + page_coord[2]) / self.scale_x)) + points_co += str(int((all_found_textline_polygons_marginals[marginal_idx][j][l][0] + page_coord[ + 2]) / self.scale_x)) points_co += ',' - points_co += str(int((all_found_textline_polygons_marginals[marginal_idx][j][l][1] + page_coord[0]) / self.scale_y)) + points_co += str(int((all_found_textline_polygons_marginals[marginal_idx][j][l][1] + page_coord[ + 0]) / self.scale_y)) else: - points_co += str(int((all_found_textline_polygons_marginals[marginal_idx][j][l][0][0] + page_coord[2]) / self.scale_x)) + points_co += str(int((all_found_textline_polygons_marginals[marginal_idx][j][l][0][0] + + page_coord[2]) / self.scale_x)) points_co += ',' - points_co += str(int((all_found_textline_polygons_marginals[marginal_idx][j][l][0][1] + page_coord[0]) / self.scale_y)) + points_co += str(int((all_found_textline_polygons_marginals[marginal_idx][j][l][0][1] + + page_coord[0]) / self.scale_y)) elif (self.curved_line or self.textline_light) and np.abs(slopes_marginals[marginal_idx]) > 45: if len(all_found_textline_polygons_marginals[marginal_idx][j][l]) == 2: - points_co += str(int((all_found_textline_polygons_marginals[marginal_idx][j][l][0] + all_box_coord_marginals[marginal_idx][2] + page_coord[2]) / self.scale_x)) + points_co += str(int((all_found_textline_polygons_marginals[marginal_idx][j][l][0] + + all_box_coord_marginals[marginal_idx][2] + page_coord[2]) / self.scale_x)) points_co += ',' - points_co += str(int((all_found_textline_polygons_marginals[marginal_idx][j][l][1] + all_box_coord_marginals[marginal_idx][0] + page_coord[0]) / self.scale_y)) + points_co += str(int((all_found_textline_polygons_marginals[marginal_idx][j][l][1] + + all_box_coord_marginals[marginal_idx][0] + page_coord[0]) / self.scale_y)) else: - points_co += str(int((all_found_textline_polygons_marginals[marginal_idx][j][l][0][0] + all_box_coord_marginals[marginal_idx][2] + page_coord[2]) / self.scale_x)) + points_co += str(int((all_found_textline_polygons_marginals[marginal_idx][j][l][0][0] + + all_box_coord_marginals[marginal_idx][2] + page_coord[2]) / self.scale_x)) points_co += ',' - points_co += str(int((all_found_textline_polygons_marginals[marginal_idx][j][l][0][1] + all_box_coord_marginals[marginal_idx][0] + page_coord[0]) / self.scale_y)) + points_co += str(int((all_found_textline_polygons_marginals[marginal_idx][j][l][0][1] + + all_box_coord_marginals[marginal_idx][0] + page_coord[0]) / self.scale_y)) points_co += ' ' coords.set_points(points_co[:-1]) - def serialize_lines_in_region(self, text_region, all_found_textline_polygons, region_idx, page_coord, all_box_coord, slopes, counter): + def serialize_lines_in_region(self, text_region, all_found_textline_polygons, region_idx, page_coord, all_box_coord, + slopes, counter): self.logger.debug('enter serialize_lines_in_region') for j in range(len(all_found_textline_polygons[region_idx])): coords = CoordsType() @@ -104,11 +123,15 @@ def serialize_lines_in_region(self, text_region, all_found_textline_polygons, re for idx_contour_textline, contour_textline in enumerate(all_found_textline_polygons[region_idx][j]): if not (self.curved_line or self.textline_light): if len(contour_textline) == 2: - textline_x_coord = max(0, int((contour_textline[0] + region_bboxes[2] + page_coord[2]) / self.scale_x)) - textline_y_coord = max(0, int((contour_textline[1] + region_bboxes[0] + page_coord[0]) / self.scale_y)) + textline_x_coord = max(0, int((contour_textline[0] + region_bboxes[2] + page_coord[ + 2]) / self.scale_x)) + textline_y_coord = max(0, int((contour_textline[1] + region_bboxes[0] + page_coord[ + 0]) / self.scale_y)) else: - textline_x_coord = max(0, int((contour_textline[0][0] + region_bboxes[2] + page_coord[2]) / self.scale_x)) - textline_y_coord = max(0, int((contour_textline[0][1] + region_bboxes[0] + page_coord[0]) / self.scale_y)) + textline_x_coord = max(0, int((contour_textline[0][0] + region_bboxes[2] + page_coord[ + 2]) / self.scale_x)) + textline_y_coord = max(0, int((contour_textline[0][1] + region_bboxes[0] + page_coord[ + 0]) / self.scale_y)) points_co += str(textline_x_coord) points_co += ',' points_co += str(textline_y_coord) @@ -121,16 +144,18 @@ def serialize_lines_in_region(self, text_region, all_found_textline_polygons, re else: points_co += str(int((contour_textline[0][0] + page_coord[2]) / self.scale_x)) points_co += ',' - points_co += str(int((contour_textline[0][1] + page_coord[0])/self.scale_y)) + points_co += str(int((contour_textline[0][1] + page_coord[0]) / self.scale_y)) elif (self.curved_line or self.textline_light) and np.abs(slopes[region_idx]) > 45: - if len(contour_textline)==2: - points_co += str(int((contour_textline[0] + region_bboxes[2] + page_coord[2])/self.scale_x)) + if len(contour_textline) == 2: + points_co += str(int((contour_textline[0] + region_bboxes[2] + page_coord[2]) / self.scale_x)) points_co += ',' - points_co += str(int((contour_textline[1] + region_bboxes[0] + page_coord[0])/self.scale_y)) + points_co += str(int((contour_textline[1] + region_bboxes[0] + page_coord[0]) / self.scale_y)) else: - points_co += str(int((contour_textline[0][0] + region_bboxes[2]+page_coord[2])/self.scale_x)) + points_co += str( + int((contour_textline[0][0] + region_bboxes[2] + page_coord[2]) / self.scale_x)) points_co += ',' - points_co += str(int((contour_textline[0][1] + region_bboxes[0]+page_coord[0])/self.scale_y)) + points_co += str( + int((contour_textline[0][1] + region_bboxes[0] + page_coord[0]) / self.scale_y)) points_co += ' ' coords.set_points(points_co[:-1]) @@ -140,7 +165,11 @@ def write_pagexml(self, pcgts): with open(out_fname, 'w') as f: f.write(to_xml(pcgts)) - def build_pagexml_no_full_layout(self, found_polygons_text_region, page_coord, order_of_texts, id_of_texts, all_found_textline_polygons, all_box_coord, found_polygons_text_region_img, found_polygons_marginals, all_found_textline_polygons_marginals, all_box_coord_marginals, slopes, slopes_marginals, cont_page, polygons_lines_to_be_written_in_xml, found_polygons_tables): + def build_pagexml_no_full_layout(self, found_polygons_text_region, page_coord, order_of_texts, id_of_texts, + all_found_textline_polygons, all_box_coord, found_polygons_text_region_img, + found_polygons_marginals, all_found_textline_polygons_marginals, + all_box_coord_marginals, slopes, slopes_marginals, cont_page, + polygons_lines_to_be_written_in_xml, found_polygons_tables): self.logger.debug('enter build_pagexml_no_full_layout') # create the file structure @@ -156,36 +185,42 @@ def build_pagexml_no_full_layout(self, found_polygons_text_region, page_coord, o for mm in range(len(found_polygons_text_region)): textregion = TextRegionType(id=counter.next_region_id, type_='paragraph', - Coords=CoordsType(points=self.calculate_polygon_coords(found_polygons_text_region[mm], page_coord)), - ) + Coords=CoordsType( + points=self.calculate_polygon_coords(found_polygons_text_region[mm], + page_coord)), + ) page.add_TextRegion(textregion) - self.serialize_lines_in_region(textregion, all_found_textline_polygons, mm, page_coord, all_box_coord, slopes, counter) + self.serialize_lines_in_region(textregion, all_found_textline_polygons, mm, page_coord, all_box_coord, + slopes, counter) for mm in range(len(found_polygons_marginals)): marginal = TextRegionType(id=counter.next_region_id, type_='marginalia', - Coords=CoordsType(points=self.calculate_polygon_coords(found_polygons_marginals[mm], page_coord))) + Coords=CoordsType( + points=self.calculate_polygon_coords(found_polygons_marginals[mm], + page_coord))) page.add_TextRegion(marginal) - self.serialize_lines_in_marginal(marginal, all_found_textline_polygons_marginals, mm, page_coord, all_box_coord_marginals, slopes_marginals, counter) + self.serialize_lines_in_marginal(marginal, all_found_textline_polygons_marginals, mm, page_coord, + all_box_coord_marginals, slopes_marginals, counter) for mm in range(len(found_polygons_text_region_img)): img_region = ImageRegionType(id=counter.next_region_id, Coords=CoordsType()) page.add_ImageRegion(img_region) points_co = '' for lmm in range(len(found_polygons_text_region_img[mm])): - points_co += str(int((found_polygons_text_region_img[mm][lmm,0,0] + page_coord[2]) / self.scale_x)) + points_co += str(int((found_polygons_text_region_img[mm][lmm, 0, 0] + page_coord[2]) / self.scale_x)) points_co += ',' - points_co += str(int((found_polygons_text_region_img[mm][lmm,0,1] + page_coord[0]) / self.scale_y)) + points_co += str(int((found_polygons_text_region_img[mm][lmm, 0, 1] + page_coord[0]) / self.scale_y)) points_co += ' ' img_region.get_Coords().set_points(points_co[:-1]) - + for mm in range(len(polygons_lines_to_be_written_in_xml)): sep_hor = SeparatorRegionType(id=counter.next_region_id, Coords=CoordsType()) page.add_SeparatorRegion(sep_hor) points_co = '' for lmm in range(len(polygons_lines_to_be_written_in_xml[mm])): - points_co += str(int((polygons_lines_to_be_written_in_xml[mm][lmm,0,0] ) / self.scale_x)) + points_co += str(int((polygons_lines_to_be_written_in_xml[mm][lmm, 0, 0]) / self.scale_x)) points_co += ',' - points_co += str(int((polygons_lines_to_be_written_in_xml[mm][lmm,0,1] ) / self.scale_y)) + points_co += str(int((polygons_lines_to_be_written_in_xml[mm][lmm, 0, 1]) / self.scale_y)) points_co += ' ' sep_hor.get_Coords().set_points(points_co[:-1]) for mm in range(len(found_polygons_tables)): @@ -193,15 +228,21 @@ def build_pagexml_no_full_layout(self, found_polygons_text_region, page_coord, o page.add_TableRegion(tab_region) points_co = '' for lmm in range(len(found_polygons_tables[mm])): - points_co += str(int((found_polygons_tables[mm][lmm,0,0] + page_coord[2]) / self.scale_x)) + points_co += str(int((found_polygons_tables[mm][lmm, 0, 0] + page_coord[2]) / self.scale_x)) points_co += ',' - points_co += str(int((found_polygons_tables[mm][lmm,0,1] + page_coord[0]) / self.scale_y)) + points_co += str(int((found_polygons_tables[mm][lmm, 0, 1] + page_coord[0]) / self.scale_y)) points_co += ' ' tab_region.get_Coords().set_points(points_co[:-1]) return pcgts - def build_pagexml_full_layout(self, found_polygons_text_region, found_polygons_text_region_h, page_coord, order_of_texts, id_of_texts, all_found_textline_polygons, all_found_textline_polygons_h, all_box_coord, all_box_coord_h, found_polygons_text_region_img, found_polygons_tables, found_polygons_drop_capitals, found_polygons_marginals, all_found_textline_polygons_marginals, all_box_coord_marginals, slopes, slopes_h, slopes_marginals, cont_page, polygons_lines_to_be_written_in_xml): + def build_pagexml_full_layout(self, found_polygons_text_region, found_polygons_text_region_h, page_coord, + order_of_texts, id_of_texts, all_found_textline_polygons, + all_found_textline_polygons_h, all_box_coord, all_box_coord_h, + found_polygons_text_region_img, found_polygons_tables, found_polygons_drop_capitals, + found_polygons_marginals, all_found_textline_polygons_marginals, + all_box_coord_marginals, slopes, slopes_h, slopes_marginals, cont_page, + polygons_lines_to_be_written_in_xml): self.logger.debug('enter build_pagexml_full_layout') # create the file structure @@ -216,35 +257,48 @@ def build_pagexml_full_layout(self, found_polygons_text_region, found_polygons_t for mm in range(len(found_polygons_text_region)): textregion = TextRegionType(id=counter.next_region_id, type_='paragraph', - Coords=CoordsType(points=self.calculate_polygon_coords(found_polygons_text_region[mm], page_coord))) + Coords=CoordsType( + points=self.calculate_polygon_coords(found_polygons_text_region[mm], + page_coord))) page.add_TextRegion(textregion) - self.serialize_lines_in_region(textregion, all_found_textline_polygons, mm, page_coord, all_box_coord, slopes, counter) + self.serialize_lines_in_region(textregion, all_found_textline_polygons, mm, page_coord, all_box_coord, + slopes, counter) self.logger.debug('len(found_polygons_text_region_h) %s', len(found_polygons_text_region_h)) for mm in range(len(found_polygons_text_region_h)): textregion = TextRegionType(id=counter.next_region_id, type_='header', - Coords=CoordsType(points=self.calculate_polygon_coords(found_polygons_text_region_h[mm], page_coord))) + Coords=CoordsType( + points=self.calculate_polygon_coords(found_polygons_text_region_h[mm], + page_coord))) page.add_TextRegion(textregion) - self.serialize_lines_in_region(textregion, all_found_textline_polygons_h, mm, page_coord, all_box_coord_h, slopes_h, counter) + self.serialize_lines_in_region(textregion, all_found_textline_polygons_h, mm, page_coord, all_box_coord_h, + slopes_h, counter) for mm in range(len(found_polygons_marginals)): marginal = TextRegionType(id=counter.next_region_id, type_='marginalia', - Coords=CoordsType(points=self.calculate_polygon_coords(found_polygons_marginals[mm], page_coord))) + Coords=CoordsType( + points=self.calculate_polygon_coords(found_polygons_marginals[mm], + page_coord))) page.add_TextRegion(marginal) - self.serialize_lines_in_marginal(marginal, all_found_textline_polygons_marginals, mm, page_coord, all_box_coord_marginals, slopes_marginals, counter) + self.serialize_lines_in_marginal(marginal, all_found_textline_polygons_marginals, mm, page_coord, + all_box_coord_marginals, slopes_marginals, counter) for mm in range(len(found_polygons_drop_capitals)): page.add_TextRegion(TextRegionType(id=counter.next_region_id, type_='drop-capital', - Coords=CoordsType(points=self.calculate_polygon_coords(found_polygons_drop_capitals[mm], page_coord)))) + Coords=CoordsType(points=self.calculate_polygon_coords( + found_polygons_drop_capitals[mm], page_coord)))) for mm in range(len(found_polygons_text_region_img)): - page.add_ImageRegion(ImageRegionType(id=counter.next_region_id, Coords=CoordsType(points=self.calculate_polygon_coords(found_polygons_text_region_img[mm], page_coord)))) - + page.add_ImageRegion(ImageRegionType(id=counter.next_region_id, Coords=CoordsType( + points=self.calculate_polygon_coords(found_polygons_text_region_img[mm], page_coord)))) + for mm in range(len(polygons_lines_to_be_written_in_xml)): - page.add_SeparatorRegion(ImageRegionType(id=counter.next_region_id, Coords=CoordsType(points=self.calculate_polygon_coords(polygons_lines_to_be_written_in_xml[mm], [0 , 0, 0, 0])))) - + page.add_SeparatorRegion(ImageRegionType(id=counter.next_region_id, Coords=CoordsType( + points=self.calculate_polygon_coords(polygons_lines_to_be_written_in_xml[mm], [0, 0, 0, 0])))) + for mm in range(len(found_polygons_tables)): - page.add_TableRegion(TableRegionType(id=counter.next_region_id, Coords=CoordsType(points=self.calculate_polygon_coords(found_polygons_tables[mm], page_coord)))) + page.add_TableRegion(TableRegionType(id=counter.next_region_id, Coords=CoordsType( + points=self.calculate_polygon_coords(found_polygons_tables[mm], page_coord)))) return pcgts @@ -260,6 +314,5 @@ def calculate_polygon_coords(self, contour, page_coord): coords += str(int((value_bbox[0][0] + page_coord[2]) / self.scale_x)) coords += ',' coords += str(int((value_bbox[0][1] + page_coord[0]) / self.scale_y)) - coords=coords + ' ' + coords = coords + ' ' return coords[:-1] - diff --git a/tests/base.py b/tests/base.py index 9de35ef..841355d 100644 --- a/tests/base.py +++ b/tests/base.py @@ -10,12 +10,14 @@ import pytest from ocrd_utils import disableLogging, initLogging + def main(fn=None): if fn: sys.exit(pytest.main([fn])) else: unittests_main() + class TestCase(VanillaTestCase): @classmethod @@ -26,6 +28,7 @@ def setUp(self): disableLogging() initLogging() + class CapturingTestCase(TestCase): """ A TestCase that needs to capture stderr/stdout and invoke click CLI. @@ -42,7 +45,7 @@ def invoke_cli(self, cli, args): """ self.capture_out_err() # XXX snapshot just before executing the CLI code = 0 - sys.argv[1:] = args # XXX necessary because sys.argv reflects pytest args not cli args + sys.argv[1:] = args # XXX necessary because sys.argv reflects pytest args not cli args try: cli.main(args=args) except SystemExit as e: diff --git a/tests/test_counter.py b/tests/test_counter.py index 125b7c1..145fbf6 100644 --- a/tests/test_counter.py +++ b/tests/test_counter.py @@ -1,6 +1,7 @@ from tests.base import main from eynollah.eynollah.utils.counter import EynollahIdCounter + def test_counter_string(): c = EynollahIdCounter() assert c.next_region_id == 'region_0001' @@ -11,6 +12,7 @@ def test_counter_string(): assert c.region_id(999) == 'region_0999' assert c.line_id(999, 888) == 'region_0999_line_0888' + def test_counter_init(): c = EynollahIdCounter(region_idx=2) assert c.get('region') == 2 @@ -19,6 +21,7 @@ def test_counter_init(): c.reset() assert c.get('region') == 2 + def test_counter_methods(): c = EynollahIdCounter() assert c.get('region') == 0 @@ -29,5 +32,6 @@ def test_counter_methods(): c.inc('region', -9) assert c.get('region') == 1 + if __name__ == '__main__': main(__file__) diff --git a/tests/test_dpi.py b/tests/test_dpi.py index 6317cd4..7f542e5 100644 --- a/tests/test_dpi.py +++ b/tests/test_dpi.py @@ -3,9 +3,11 @@ from eynollah.eynollah.utils.pil_cv2 import check_dpi from tests.base import main + def test_dpi(): fpath = str(Path(__file__).parent.joinpath('resources', 'kant_aufklaerung_1784_0020.tif')) assert 230 == check_dpi(cv2.imread(fpath)) + if __name__ == '__main__': main(__file__) diff --git a/tests/test_run.py b/tests/test_run.py index eeee964..f93de5e 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -8,6 +8,7 @@ EYNOLLAH_MODELS = environ.get('EYNOLLAH_MODELS', str(testdir.joinpath('..', 'models_eynollah').resolve())) + class TestEynollahRun(TestCase): def test_full_run(self): @@ -20,5 +21,6 @@ def test_full_run(self): print(code, out, err) assert not code + if __name__ == '__main__': main(__file__) diff --git a/tests/test_xml.py b/tests/test_xml.py index c8bac53..2fddd77 100644 --- a/tests/test_xml.py +++ b/tests/test_xml.py @@ -4,11 +4,13 @@ PAGE_2019 = 'http://schema.primaresearch.org/PAGE/gts/pagecontent/2019-07-15' + def test_create_xml(): pcgts = create_page_xml('/path/to/img.tif', 100, 100) xmlstr = to_xml(pcgts) assert 'xmlns:pc="%s"' % PAGE_2019 in xmlstr assert 'Metadata' in xmlstr + if __name__ == '__main__': main([__file__])