from keras.models import load_model
import cv2
import numpy as np
import sys
import re

IMAGE_SIZE = (256,256)
def iou(y_true, y_pred):
    from keras import backend as K
    y_true = K.cast(K.greater(y_true, 0.5), dtype='float32')
    y_pred = K.cast(K.greater(y_pred, 0.5), dtype='float32')
    inter = K.sum(K.sum(K.squeeze(y_true * y_pred, axis=3), axis=2), axis=1)
    union = K.sum(K.sum(K.squeeze(K.clip(y_true + y_pred, 0, 1), axis=3), axis=2), axis=1)
    return K.mean((inter + K.epsilon()) / (union + K.epsilon()))
def threshold_slow(T, image):
    # grab the image dimensions
    h = image.shape[0]
    w = image.shape[1]

    # loop over the image, pixel by pixel
    for y in range(0, h):
        for x in range(0, w):
            # threshold the pixel
            image[y, x] = 255 if image[y, x] >= T else 0

    # return the thresholded image
    return image
model = load_model('model2.h5', custom_objects={'iou': iou})
name = 1
index = []
while name < 156:
    if not cv2.imread('textlocalize/validation/Input/'+str(name)+'.jpg') is None :
        index.append(str(name))
    if not cv2.imread('textlocalize/validation/Input/0'+str(name)+'.jpg') is None:
        index.append("0"+str(name))
    if not cv2.imread('textlocalize/validation/Input/00'+str(name)+'.jpg') is None:
        index.append("00"+str(name))
    name += 1
for start in index:
    print(start)
    test_im = cv2.imread('textlocalize/validation/Input/'+str(start)+'.jpg')
    if not test_im is None:
        true_size = test_im.shape

        imshow_size = (512,round(true_size[0]*512/true_size[1]))
        # cv2.imshow('Input',cv2.resize(test_im, imshow_size))

        test_im = cv2.cvtColor(test_im, cv2.COLOR_BGR2RGB)
        test_im = cv2.resize(test_im, (IMAGE_SIZE[1], IMAGE_SIZE[0]))
        test_im = test_im/255.
        test_im = np.expand_dims(test_im, axis=0)
        segmented = model.predict(test_im)
        #segmented = np.around(segmented)
        IMAGE_out = (true_size[1],true_size[0])
        # segmented = (segmented[0, :, :, 0]*255).astype('uint8')
        segmented30 = threshold_slow(0.3,segmented[0, :, :, 0])
        segmented40 = threshold_slow(0.4,segmented[0, :, :, 0])
        segmented50 = threshold_slow(0.5,segmented[0, :, :, 0])
        output30 = cv2.resize(segmented30, (IMAGE_out[0],IMAGE_out[1]))
        output40 = cv2.resize(segmented40, (IMAGE_out[0],IMAGE_out[1]))
        output50 = cv2.resize(segmented50, (IMAGE_out[0],IMAGE_out[1]))
        status = cv2.imwrite('textlocalize/testoutput30/'+str(start)+'.jpg',output30)
        status = cv2.imwrite('textlocalize/testoutput40/'+str(start)+'.jpg',output40)
        status = cv2.imwrite('textlocalize/testoutput50/'+str(start)+'.jpg',output50)
    name += 1
# cv2.imshow('Output',cv2.resize(segmented, imshow_size))
cv2.waitKey()