You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
eynollah/sbb_newspapers_org_image/plot.py

170 lines
9.1 KiB
Python

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np
import os.path
import cv2
from scipy.ndimage import gaussian_filter1d
from .utils import crop_image_inside_box
from .utils.rotate import rotyate_image_different
from .utils.resize import resize_image
class EynollahPlotter():
"""
Class collecting all the plotting and image writing methods
"""
def __init__(
self,
*,
dir_of_all,
dir_of_deskewed,
dir_of_layout,
dir_of_cropped_images,
image_filename,
image_filename_stem,
image_org=None,
scale_x=1,
scale_y=1,
):
self.dir_of_all = dir_of_all
self.dir_of_layout = dir_of_layout
self.dir_of_cropped_images = dir_of_cropped_images
self.dir_of_deskewed = dir_of_deskewed
self.image_filename = image_filename
self.image_filename_stem = image_filename_stem
# XXX TODO hacky these cannot be set at init time
self.image_org = image_org
self.scale_x = scale_x
self.scale_y = scale_y
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']
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]
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']
values_indexes = [0, 1, 2, 3, 4]
plt.figure(figsize=(80, 40))
plt.rcParams["font.size"] = "40"
plt.subplot(1, 2, 1)
plt.imshow(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]
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"))
def save_plot_of_layout(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", "Header", "Marginalia", "Drop capital", "Image", "Separator"]
values_indexes = [0, 1, 2, 8, 4, 5, 6]
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]
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"))
def save_plot_of_layout_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", "Header", "Marginalia", "Drop capital", "Image", "Separator"]
values_indexes = [0, 1, 2, 8, 4, 5, 6]
plt.figure(figsize=(80, 40))
plt.rcParams["font.size"] = "40"
plt.subplot(1, 2, 1)
plt.imshow(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]
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"))
def save_plot_of_textlines(self, textline_mask_tot_ea, image_page):
if self.dir_of_all is not None:
values = np.unique(textline_mask_tot_ea[:, :])
pixels = ["Background", "Textlines"]
values_indexes = [0, 1]
plt.figure(figsize=(80, 40))
plt.rcParams["font.size"] = "40"
plt.subplot(1, 2, 1)
plt.imshow(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]
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"))
def save_deskewed_image(self, slope_deskew):
if self.dir_of_all is not None:
cv2.imwrite(os.path.join(self.dir_of_all, self.image_filename_stem + "_org.png"), self.image_org)
if self.dir_of_deskewed is not None:
img_rotated = rotyate_image_different(self.image_org, slope_deskew)
cv2.imwrite(os.path.join(self.dir_of_deskewed, self.image_filename_stem + "_deskewed.png"), img_rotated)
def save_page_image(self, image_page):
if self.dir_of_all is not None:
cv2.imwrite(os.path.join(self.dir_of_all, self.image_filename_stem + "_page.png"), image_page)
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.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.gca().invert_yaxis()
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.legend(loc='best')
plt.savefig(os.path.join(self.dir_of_all, self.image_filename_stem+'_rotation_angle.png'))
def write_images_into_directory(self, img_contoures, image_page):
if self.dir_of_cropped_images is not None:
index = 0
for cont_ind in img_contoures:
x, y, w, h = cv2.boundingRect(cont_ind)
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))
path = os.path.join(self.dir_of_cropped_images, self.image_filename_stem + "_" + str(index) + ".jpg")
cv2.imwrite(path, croped_page)
index += 1