qurator-spk
/
sbb_binarization
mirror of https://github.com/qurator-spk/sbb_binarization.git
1
0
Fork 0
Code Issues Releases Wiki Activity

Add new file

Browse Source
pull/1/head
Rezanezhad, Vahid 5 years ago
parent ad908882f3
commit 9cc72d1918
1 changed files with 509 additions and 0 deletions
Show Stats Download Patch File Download Diff File

509
sbb_binarize/sbb_binarize.py
Unescape Escape View File

@ -0,0 +1,509 @@
#! /usr/bin/env python3
__version__= '1.0'
import argparse
import sys
import os
import numpy as np
import warnings
import xml.etree.ElementTree as et
import pandas as pd
from tqdm import tqdm
import csv
import cv2
import seaborn as sns
import matplotlib.pyplot as plt
from keras.models import load_model
import tensorflow as tf
from keras import backend as K
from skimage.filters import threshold_otsu
import keras.losses
with warnings.catch_warnings():
warnings.simplefilter("ignore")
__doc__=\
"""
Tool to load model and binarize a given image.
"""
class sbb_binarize:
def __init__(self,image,model, patches='false',save=None, ground_truth=None,weights_dir=None ):
self.image=image
self.patches=patches
self.save=save
self.model_dir=model
self.ground_truth=ground_truth
self.weights_dir=weights_dir
def resize_image(self,img_in,input_height,input_width):
return cv2.resize( img_in, ( input_width,input_height) ,interpolation=cv2.INTER_NEAREST)
def color_images(self,seg):
ann_u=range(self.n_classes)
if len(np.shape(seg))==3:
seg=seg[:,:,0]
seg_img=np.zeros((np.shape(seg)[0],np.shape(seg)[1],3)).astype(np.uint8)
colors=sns.color_palette("hls", self.n_classes)
for c in ann_u:
c=int(c)
segl=(seg==c)
seg_img[:,:,0][seg==c]=c
seg_img[:,:,1][seg==c]=c
seg_img[:,:,2][seg==c]=c
return seg_img
def otsu_copy_binary(self,img):
img_r=np.zeros((img.shape[0],img.shape[1],3))
img1=img[:,:,0]
#print(img.min())
#print(img[:,:,0].min())
#blur = cv2.GaussianBlur(img,(5,5))
#ret3,th3 = cv2.threshold(blur,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
retval1, threshold1 = cv2.threshold(img1, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
img_r[:,:,0]=threshold1
img_r[:,:,1]=threshold1
img_r[:,:,2]=threshold1
#img_r=img_r/float(np.max(img_r))*255
return img_r
def otsu_copy(self,img):
img_r=np.zeros((img.shape[0],img.shape[1],3))
#img1=img[:,:,0]
#print(img.min())
#print(img[:,:,0].min())
#blur = cv2.GaussianBlur(img,(5,5))
#ret3,th3 = cv2.threshold(blur,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
_, threshold1 = cv2.threshold(img[:,:,0], 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
_, threshold2 = cv2.threshold(img[:,:,1], 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
_, threshold3 = cv2.threshold(img[:,:,2], 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
img_r[:,:,0]=threshold1
img_r[:,:,1]=threshold2
img_r[:,:,2]=threshold3
###img_r=img_r/float(np.max(img_r))*255
return img_r
def otsu_org(self,img):
binary_global = img > threshold_otsu(img)
binary_global=binary_global*255
#plt.imshow(binary_sauvola*255,cmap=plt.cm.gray)
#plt.imshow(binary_global)
#plt.show()
#print(np.unique(binary_global))
binary_global=np.repeat(binary_global[:, :, np.newaxis], 3, axis=2)
plt.imshow(binary_global)
plt.show()
print(binary_global.shape)
return binary_global
def soft_dice_loss(self,y_true, y_pred, epsilon=1e-6):
axes = tuple(range(1, len(y_pred.shape)-1))
numerator = 2. * K.sum(y_pred * y_true, axes)
denominator = K.sum(K.square(y_pred) + K.square(y_true), axes)
return 1.00 - K.mean(numerator / (denominator + epsilon)) # average over classes and batch
def weighted_categorical_crossentropy(self,weights=None):
def loss(y_true, y_pred):
labels_floats = tf.cast(y_true, tf.float32)
per_pixel_loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=labels_floats,logits=y_pred)
if weights is not None:
weight_mask = tf.maximum(tf.reduce_max(tf.constant(
np.array(weights, dtype=np.float32)[None, None, None])
* labels_floats, axis=-1), 1.0)
per_pixel_loss = per_pixel_loss * weight_mask[:, :, :, None]
return tf.reduce_mean(per_pixel_loss)
return self.loss
def IoU(self,Yi,y_predi):
## mean Intersection over Union
## Mean IoU = TP/(FN + TP + FP)
IoUs = []
Nclass = np.unique(Yi)
for c in Nclass:
TP = np.sum( (Yi == c)&(y_predi==c) )
FP = np.sum( (Yi != c)&(y_predi==c) )
FN = np.sum( (Yi == c)&(y_predi != c))
IoU = TP/float(TP + FP + FN)
if self.n_classes>2:
print("class {:02.0f}: #TP={:6.0f}, #FP={:6.0f}, #FN={:5.0f}, IoU={:4.3f}".format(c,TP,FP,FN,IoU))
IoUs.append(IoU)
if self.n_classes>2:
mIoU = np.mean(IoUs)
print("_________________")
print("Mean IoU: {:4.3f}".format(mIoU))
return mIoU
elif self.n_classes==2:
mIoU = IoUs[1]
print("_________________")
print("IoU: {:4.3f}".format(mIoU))
return mIoU
def start_new_session_and_model(self):
config = tf.ConfigProto()
config.gpu_options.allow_growth=True
self.session =tf.Session(config=config)# tf.InteractiveSession()
#keras.losses.custom_loss = self.weighted_categorical_crossentropy
def load_model(self,model_name):
self.model = load_model(self.model_dir+'/'+model_name , compile=False)
#if self.weights_dir!=None:
# print('man burdayammmmaaa')
# self.model.load_weights(self.weights_dir)
self.img_height=self.model.layers[len(self.model.layers)-1].output_shape[1]
self.img_width=self.model.layers[len(self.model.layers)-1].output_shape[2]
self.n_classes=self.model.layers[len(self.model.layers)-1].output_shape[3]
def end_session(self):
self.session.close()
del self.model
del self.session
def predict(self,model_name):
#self.start_new_session_and_model(model_name)
self.load_model(model_name)
if self.patches=='true' or self.patches=='True':
print(self.patches,'gadaaiikk')
#def textline_contours(img,input_width,input_height,n_classes,model):
img=cv2.imread(self.image)
if img.shape[0]<self.img_height:
img=cv2.resize( img, ( img.shape[1],self.img_width) ,interpolation=cv2.INTER_NEAREST)
if img.shape[1]<self.img_width:
img=cv2.resize( img, ( self.img_height,img.shape[0]) ,interpolation=cv2.INTER_NEAREST)
margin=True
if margin:
kernel = np.ones((5,5),np.uint8)
width=self.img_width
height=self.img_height
#offset=int(.1*width)
offset=int(0.1*width)
width_mid=width-2*offset
height_mid=height-2*offset
#img= cv2.medianBlur(img, 5)
#img = cv2.GaussianBlur(img,(5,5),0)
#img= cv2.medianBlur(img, 5)
#img= cv2.medianBlur(img, 5)
#img= cv2.medianBlur(img, 5)
#img= cv2.medianBlur(img, 5)
#img=self.otsu_copy_binary(img)
#img=self.otsu_org(img)
img=img.astype(np.uint8)
#for i in range(10):
# img= cv2.medianBlur(img, 3)
img=img/255.0
img_h=img.shape[0]
img_w=img.shape[1]
prediction_true=np.zeros((img_h,img_w,3))
mask_true=np.zeros((img_h,img_w))
nxf=img_w/float(width_mid)
nyf=img_h/float(height_mid)
if nxf>int(nxf):
nxf=int(nxf)+1
else:
nxf=int(nxf)
if nyf>int(nyf):
nyf=int(nyf)+1
else:
nyf=int(nyf)
for i in range(nxf):
for j in range(nyf):
if i==0:
index_x_d=i*width_mid
index_x_u=index_x_d+width#(i+1)*width
elif i>0:
index_x_d=i*width_mid
index_x_u=index_x_d+width#(i+1)*width
if j==0:
index_y_d=j*height_mid
index_y_u=index_y_d+height#(j+1)*height
elif j>0:
index_y_d=j*height_mid
index_y_u=index_y_d+height#(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_p_pred=self.model.predict(
img_patch.reshape(1,img_patch.shape[0],img_patch.shape[1],img_patch.shape[2]))
#print(np.unique(label_p_pred))
th3=label_p_pred[0,:,:,1]
th3=th3*255
th3=th3.astype(np.uint8)
#print(np.unique(th3))
ret3,th3 = cv2.threshold(th3,30,250,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
seg=np.argmax(label_p_pred,axis=3)[0]
seg_color=self.color_images(seg)
seg_color=seg_color[offset:seg_color.shape[0]-offset,offset:seg_color.shape[1]-offset,:]
seg=seg[offset:seg.shape[0]-offset,offset:seg.shape[1]-offset]
th3=th3[offset:th3.shape[0]-offset,offset:th3.shape[1]-offset]
mask_true[index_y_d+offset:index_y_u-offset,index_x_d+offset:index_x_u-offset]=seg
prediction_true[index_y_d+offset:index_y_u-offset,index_x_d+offset:index_x_u-offset,:]=seg_color
y_predi = mask_true
#print(np.unique(mask_true))
#find_contours(mask_true)
#y_testi = label[:,:,0]#np.argmax(label.reshape(1,label.shape[0],label.shape[1],label.shape[2]), axis=3)
#y_predi=cv2.erode(y_predi,kernel,iterations=3)
y_predi=cv2.resize( y_predi, ( img.shape[1],img.shape[0]) ,interpolation=cv2.INTER_NEAREST)
return y_predi
if not margin:
kernel = np.ones((5,5),np.uint8)
width=self.img_width
height=self.img_height
#img = cv2.medianBlur(img,5)
img=self.otsu_copy_binary(img)
#img=cv2.bilateralFilter(img,9,75,75)
img = cv2.GaussianBlur(img,(5,5),0)
img=img/255.0
img_h=img.shape[0]
img_w=img.shape[1]
prediction_true=np.zeros((img_h,img_w,3))
mask_true=np.zeros((img_h,img_w))
nxf=img_w/float(width)
nyf=img_h/float(height)
if nxf>int(nxf):
nxf=int(nxf)+1
else:
nxf=int(nxf)
if nyf>int(nyf):
nyf=int(nyf)+1
else:
nyf=int(nyf)
print(nxf,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_p_pred=self.model.predict(img_patch.reshape(1,img_patch.shape[0],img_patch.shape[1],img_patch.shape[2]) )
seg=np.argmax(label_p_pred,axis=3)[0]
seg_color=self.color_images(seg)
###seg_color=color_images_diva(seg,n_classes)
mask_true[index_y_d:index_y_u,index_x_d:index_x_u]=seg
prediction_true[index_y_d:index_y_u,index_x_d:index_x_u,:]=seg_color
#mask_true=color_images(mask_true,n_classes)
y_predi = mask_true
#print(np.unique(mask_true))
#find_contours(mask_true)
#y_testi = label[:,:,0]#np.argmax(label.reshape(1,label.shape[0],label.shape[1],label.shape[2]), axis=3)
#y_predi=cv2.erode(y_predi,kernel,iterations=3)
y_predi=cv2.resize( y_predi, ( img.shape[1],img.shape[0]) ,interpolation=cv2.INTER_NEAREST)
#self.end_session()
return y_predi
#def extract_page(img,input_width,input_height,n_classes,model):
if self.patches=='false' or self.patches=='False':
img=cv2.imread(self.image,0)
img_org_height=img.shape[0]
img_org_width=img.shape[1]
#kernel = np.ones((5,5),np.uint8)
width=self.img_width
height=self.img_height
#for _ in range(1):
#img = cv2.medianBlur(img,5)
img=self.otsu_org(img)
#img=img.astype(np.uint8)
img=img.astype(np.uint8)
#img = cv2.medianBlur(img,5)
#img=img.astype(np.uint8)
#img = cv2.GaussianBlur(img,(5,5),0)
#img=self.otsu_copy_binary(img)
img=img.astype(np.uint8)
img=img/255.0
img=self.resize_image(img,self.img_height,self.img_width)
label_p_pred=self.model.predict(
img.reshape(1,img.shape[0],img.shape[1],img.shape[2]))
seg=np.argmax(label_p_pred,axis=3)[0]
print(np.shape(seg),np.unique(seg))
plt.imshow(seg*255)
plt.show()
seg_color=self.color_images(seg)
print(np.unique(seg_color))
#imgs = seg_color#/np.max(seg_color)*255#np.repeat(seg_color[:, :, np.newaxis], 3, axis=2)
y_predi=cv2.resize( seg_color, ( img_org_width,img_org_height) ,interpolation=cv2.INTER_NEAREST)
return y_predi
def run(self):
self.start_new_session_and_model()
models_n=os.listdir(self.model_dir)
img_last=0
for model_in in models_n:
res=self.predict(model_in)
if self.ground_truth!=None:
gt_img=cv2.imread(self.ground_truth)
print(np.shape(gt_img),np.shape(res))
#self.IoU(gt_img[:,:,0],res)
#print(np.unique(res))
img_fin=np.zeros((res.shape[0],res.shape[1],3) )
res[:,:][res[:,:]==0]=2
res=res-1
res=res*255
img_fin[:,:,0]=res
img_fin[:,:,1]=res
img_fin[:,:,2]=res
img_fin=img_fin.astype(np.uint8)
img_fin=(res[:,:]==0)*255
img_last=img_last+img_fin
kernel = np.ones((5,5),np.uint8)
img_last[:,:][img_last[:,:]>0]=255
img_last=(img_last[:,:]==0)*255
#img_fin= cv2.medianBlur(img_fin, 5)
if self.save is not None:
cv2.imwrite('./'+self.save,img_last)
plt.imshow(img_last)
plt.show()
def main():
parser=argparse.ArgumentParser()
parser.add_argument('-i','--image', dest='inp1', default=None, help='directory of alto files which have to be transformed.')
parser.add_argument('-p','--patches', dest='inp3', default=False, help='use patches of image for prediction or should image resize be applied to be fit for model. this parameter should be true or false')
parser.add_argument('-s','--save', dest='inp4', default=False, help='save prediction with agive name here. The name and format should be given (0045.tif).')
parser.add_argument('-m','--model', dest='inp2', default=None, help='model directory and name should be provided here.')
parser.add_argument('-gt','--groundtruth', dest='inp5', default=None, help='ground truth directory if you want to see the iou of prediction.')
parser.add_argument('-mw','--model_weights', dest='inp6', default=None, help='previous model weights which are saved.')
options=parser.parse_args()
possibles=globals()
possibles.update(locals())
x=sbb_binarize(options.inp1,options.inp2,options.inp3,options.inp4,options.inp5,options.inp6)
x.run()
if __name__=="__main__":
main()
Write Preview
Loading…
Cancel
Save