Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import tensorflow as tf | |
| import numpy as np | |
| import cv2 | |
| 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 do_prediction(img): | |
| model = from_pretrained_keras("vahidrezanezhad/sbb_binarization") | |
| 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 prediction_true * 255 | |
| iface = gr.Interface(fn=do_prediction, inputs=gr.Image(), outputs=gr.Image()) | |
| iface.launch() | |