Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import tensorflow as tf | |
| import numpy as np | |
| import cv2 | |
| from PIL import Image | |
| from huggingface_hub import from_pretrained_keras | |
| def resize_image(img_in,input_height,input_width): | |
| return cv2.resize( img_in, ( input_width,input_height) ,interpolation=cv2.INTER_NEAREST) | |
| def visualize_model_output(prediction, img): | |
| unique_classes = np.unique(prediction[:,:,0]) | |
| rgb_colors = {'0' : [255, 255, 255], | |
| '1' : [255, 0, 0], | |
| '2' : [255, 125, 0], | |
| '3' : [255, 0, 125], | |
| '4' : [125, 125, 125], | |
| '5' : [125, 125, 0], | |
| '6' : [0, 125, 255], | |
| '7' : [0, 125, 0], | |
| '8' : [125, 125, 125], | |
| '9' : [0, 125, 255], | |
| '10' : [125, 0, 125], | |
| '11' : [0, 255, 0], | |
| '12' : [0, 0, 255], | |
| '13' : [0, 255, 255], | |
| '14' : [255, 125, 125], | |
| '15' : [255, 0, 255]} | |
| output = np.zeros(prediction.shape) | |
| print(output.shape,'shape of output') | |
| for unq_class in unique_classes: | |
| print(unq_class,'unq_class') | |
| rgb_class_unique = rgb_colors[str(int(unq_class))] | |
| output[:,:,0][prediction[:,:,0]==unq_class] = rgb_class_unique[0] | |
| output[:,:,1][prediction[:,:,0]==unq_class] = rgb_class_unique[1] | |
| output[:,:,2][prediction[:,:,0]==unq_class] = rgb_class_unique[2] | |
| img = resize_image(img, output.shape[0], output.shape[1]) | |
| output = output.astype(np.int32) | |
| img = img.astype(np.int32) | |
| print(output,'shapes') | |
| print(img,'shapes2') | |
| added_image = cv2.addWeighted(img,0.5,output,0.1,0) | |
| print(added_image) | |
| return added_image | |
| def do_prediction(model_name, img): | |
| img_org = np.copy(img) | |
| model = from_pretrained_keras(model_name) | |
| match model_name: | |
| # numerical output | |
| case "SBB/eynollah-column-classifier": | |
| num_col=model.layers[len(model.layers)-1].output_shape[1] | |
| return "Found {} columns".format(num_col), None | |
| # bitmap output | |
| case "SBB/eynollah-binarization" | "SBB/eynollah-page-extraction" | "SBB/eynollah-textline" | "SBB/eynollah-textline_light" | "SBB/eynollah-enhancement" | "SBB/eynollah-tables" | "SBB/eynollah-main-regions" | "SBB/eynollah-main-regions-aug-rotation" | "SBB/eynollah-main-regions-aug-scaling" | "SBB/eynollah-main-regions-ensembled" | "SBB/eynollah-full-regions-1column" | "SBB/eynollah-full-regions-3pluscolumn": | |
| img_height_model=model.layers[len(model.layers)-1].output_shape[1] | |
| img_width_model=model.layers[len(model.layers)-1].output_shape[2] | |
| n_classes=model.layers[len(model.layers)-1].output_shape[3] | |
| if img.shape[0] < img_height_model: | |
| img = resize_image(img, img_height_model, img.shape[1]) | |
| if img.shape[1] < img_width_model: | |
| img = resize_image(img, img.shape[0], img_width_model) | |
| marginal_of_patch_percent = 0.1 | |
| margin = int(marginal_of_patch_percent * img_height_model) | |
| width_mid = img_width_model - 2 * margin | |
| height_mid = img_height_model - 2 * margin | |
| img = img / float(255.0) | |
| img = img.astype(np.float16) | |
| 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) | |
| nxf = int(nxf) + 1 if nxf > int(nxf) else int(nxf) | |
| nyf = int(nyf) + 1 if nyf > int(nyf) else int(nyf) | |
| for i in range(nxf): | |
| for j in range(nyf): | |
| if i == 0: | |
| index_x_d = i * width_mid | |
| index_x_u = index_x_d + img_width_model | |
| else: | |
| index_x_d = i * width_mid | |
| index_x_u = index_x_d + img_width_model | |
| if j == 0: | |
| index_y_d = j * height_mid | |
| index_y_u = index_y_d + img_height_model | |
| else: | |
| index_y_d = j * height_mid | |
| index_y_u = index_y_d + img_height_model | |
| if index_x_u > img_w: | |
| index_x_u = img_w | |
| index_x_d = img_w - img_width_model | |
| if index_y_u > img_h: | |
| index_y_u = img_h | |
| index_y_d = img_h - img_height_model | |
| img_patch = img[index_y_d:index_y_u, index_x_d:index_x_u, :] | |
| label_p_pred = model.predict(img_patch.reshape(1, img_patch.shape[0], img_patch.shape[1], img_patch.shape[2]), | |
| verbose=0) | |
| seg = np.argmax(label_p_pred, axis=3)[0] | |
| seg_color = np.repeat(seg[:, :, np.newaxis], 3, axis=2) | |
| if i == 0 and j == 0: | |
| seg_color = seg_color[0 : seg_color.shape[0] - margin, 0 : seg_color.shape[1] - margin, :] | |
| #seg = seg[0 : seg.shape[0] - margin, 0 : seg.shape[1] - margin] | |
| #mask_true[index_y_d + 0 : index_y_u - margin, index_x_d + 0 : index_x_u - margin] = seg | |
| prediction_true[index_y_d + 0 : index_y_u - margin, index_x_d + 0 : index_x_u - margin, :] = seg_color | |
| elif i == nxf - 1 and j == nyf - 1: | |
| seg_color = seg_color[margin : seg_color.shape[0] - 0, margin : seg_color.shape[1] - 0, :] | |
| #seg = seg[margin : seg.shape[0] - 0, margin : seg.shape[1] - 0] | |
| #mask_true[index_y_d + margin : index_y_u - 0, index_x_d + margin : index_x_u - 0] = seg | |
| prediction_true[index_y_d + margin : index_y_u - 0, index_x_d + margin : index_x_u - 0, :] = seg_color | |
| elif i == 0 and j == nyf - 1: | |
| seg_color = seg_color[margin : seg_color.shape[0] - 0, 0 : seg_color.shape[1] - margin, :] | |
| #seg = seg[margin : seg.shape[0] - 0, 0 : seg.shape[1] - margin] | |
| #mask_true[index_y_d + margin : index_y_u - 0, index_x_d + 0 : index_x_u - margin] = seg | |
| prediction_true[index_y_d + margin : index_y_u - 0, index_x_d + 0 : index_x_u - margin, :] = seg_color | |
| elif i == nxf - 1 and j == 0: | |
| seg_color = seg_color[0 : seg_color.shape[0] - margin, margin : seg_color.shape[1] - 0, :] | |
| #seg = seg[0 : seg.shape[0] - margin, margin : seg.shape[1] - 0] | |
| #mask_true[index_y_d + 0 : index_y_u - margin, index_x_d + margin : index_x_u - 0] = seg | |
| prediction_true[index_y_d + 0 : index_y_u - margin, index_x_d + margin : index_x_u - 0, :] = seg_color | |
| elif i == 0 and j != 0 and j != nyf - 1: | |
| seg_color = seg_color[margin : seg_color.shape[0] - margin, 0 : seg_color.shape[1] - margin, :] | |
| #seg = seg[margin : seg.shape[0] - margin, 0 : seg.shape[1] - margin] | |
| #mask_true[index_y_d + margin : index_y_u - margin, index_x_d + 0 : index_x_u - margin] = seg | |
| prediction_true[index_y_d + margin : index_y_u - margin, index_x_d + 0 : index_x_u - margin, :] = seg_color | |
| elif i == nxf - 1 and j != 0 and j != nyf - 1: | |
| seg_color = seg_color[margin : seg_color.shape[0] - margin, margin : seg_color.shape[1] - 0, :] | |
| #seg = seg[margin : seg.shape[0] - margin, margin : seg.shape[1] - 0] | |
| #mask_true[index_y_d + margin : index_y_u - margin, index_x_d + margin : index_x_u - 0] = seg | |
| prediction_true[index_y_d + margin : index_y_u - margin, index_x_d + margin : index_x_u - 0, :] = seg_color | |
| elif i != 0 and i != nxf - 1 and j == 0: | |
| seg_color = seg_color[0 : seg_color.shape[0] - margin, margin : seg_color.shape[1] - margin, :] | |
| #seg = seg[0 : seg.shape[0] - margin, margin : seg.shape[1] - margin] | |
| #mask_true[index_y_d + 0 : index_y_u - margin, index_x_d + margin : index_x_u - margin] = seg | |
| prediction_true[index_y_d + 0 : index_y_u - margin, index_x_d + margin : index_x_u - margin, :] = seg_color | |
| elif i != 0 and i != nxf - 1 and j == nyf - 1: | |
| seg_color = seg_color[margin : seg_color.shape[0] - 0, margin : seg_color.shape[1] - margin, :] | |
| #seg = seg[margin : seg.shape[0] - 0, margin : seg.shape[1] - margin] | |
| #mask_true[index_y_d + margin : index_y_u - 0, index_x_d + margin : index_x_u - margin] = seg | |
| prediction_true[index_y_d + margin : index_y_u - 0, index_x_d + margin : index_x_u - margin, :] = seg_color | |
| else: | |
| seg_color = seg_color[margin : seg_color.shape[0] - margin, margin : seg_color.shape[1] - margin, :] | |
| #seg = seg[margin : seg.shape[0] - margin, margin : seg.shape[1] - margin] | |
| #mask_true[index_y_d + margin : index_y_u - margin, index_x_d + margin : index_x_u - margin] = seg | |
| prediction_true[index_y_d + margin : index_y_u - margin, index_x_d + margin : index_x_u - margin, :] = seg_color | |
| prediction_true = prediction_true.astype(np.uint8) | |
| ''' | |
| img = img / float(255.0) | |
| image = resize_image(image, 224,448) | |
| prediction = model.predict(image.reshape(1,224,448,image.shape[2])) | |
| prediction = tf.squeeze(tf.round(prediction)) | |
| prediction = np.argmax(prediction,axis=2) | |
| prediction = np.repeat(prediction[:, :, np.newaxis]*255, 3, axis=2) | |
| print(prediction.shape) | |
| ''' | |
| #prediction_true = prediction_true * -1 | |
| #prediction_true = prediction_true + 1 | |
| return "No numerical output", visualize_model_output(prediction_true,img_org) | |
| # catch-all (we should not reach this) | |
| case _: | |
| return None, None | |
| title = "Welcome to the Eynollah Demo page! 👁️" | |
| description = """ | |
| <div class="row" style="display: flex"> | |
| <div class="column" style="flex: 50%; font-size: 17px"> | |
| This Space demonstrates the functionality of various Eynollah models developed at <a rel="nofollow" href="https://huggingface.co./SBB">SBB</a>. | |
| <br><br> | |
| The Eynollah suite introduces an <u>end-to-end pipeline</u> to extract layout, text lines and reading order for historic documents, where the output can be used as an input for OCR engines. | |
| Please keep in mind that with this demo you can just use <u>one of the 13 sub-modules</u> of the whole Eynollah system <u>at a time</u>. | |
| </div> | |
| <div class="column" style="flex: 5%; font-size: 17px"></div> | |
| <div class="column" style="flex: 45%; font-size: 17px"> | |
| <strong style="font-size: 19px">Resources for more information:</strong> | |
| <ul> | |
| <li>The GitHub Repo can be found <a rel="nofollow" href="https://github.com/qurator-spk/eynollah">here</a></li> | |
| <li>Associated Paper: <a rel="nofollow" href="https://doi.org/10.1145/3604951.3605513">Document Layout Analysis with Deep Learning and Heuristics</a></li> | |
| <li>The full Eynollah pipeline can be viewed <a rel="nofollow" href="https://huggingface.co./spaces/SBB/eynollah-demo-test/blob/main/eynollah-flow.png">here</a></li> | |
| </ul> | |
| </li> | |
| </div> | |
| </div> | |
| """ | |
| iface = gr.Interface( | |
| title=title, | |
| description=description, | |
| fn=do_prediction, | |
| inputs=[ | |
| gr.Dropdown([ | |
| "SBB/eynollah-binarization", | |
| "SBB/eynollah-enhancement", | |
| "SBB/eynollah-page-extraction", | |
| "SBB/eynollah-column-classifier", | |
| "SBB/eynollah-tables", | |
| "SBB/eynollah-textline", | |
| "SBB/eynollah-textline_light", | |
| "SBB/eynollah-main-regions", | |
| "SBB/eynollah-main-regions-aug-rotation", | |
| "SBB/eynollah-main-regions-aug-scaling", | |
| "SBB/eynollah-main-regions-ensembled", | |
| "SBB/eynollah-full-regions-1column", | |
| "SBB/eynollah-full-regions-3pluscolumn" | |
| ], label="Select one model of the Eynollah suite 👇", info=""), | |
| gr.Image() | |
| ], | |
| outputs=[ | |
| gr.Textbox(label="Output of model (numerical or bitmap) ⬇️"), | |
| gr.Image() | |
| ], | |
| #examples=[['example-1.jpg']] | |
| ) | |
| iface.launch() |