import gradio as gr
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
import requests
from PIL import Image
import numpy as np
import cv2
processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
model = VisionEncoderDecoderModel.from_pretrained("aico/TrOCR-MNIST")

def _group_rectangles(rec):
    """
    Uion intersecting rectangles.
    Args:
        rec - list of rectangles in form [x, y, w, h]
    Return:
        list of grouped ractangles 
    """
    tested = [False for i in range(len(rec))]
    final = []
    i = 0
    while i < len(rec):
        if not tested[i]:
            j = i+1
            while j < len(rec):
                if not tested[j] and intersect_area(rec[i], rec[j]):
                    rec[i] = union(rec[i], rec[j])
                    tested[j] = True
                    j = i
                j += 1
            final += [rec[i]]
        i += 1

    return final

def process_image(image):
    bounding_boxes = []
    generated_text_list = []
    #boundingBoxes_2 = []
    #print(np.shape(image))
    #print(image)
    #dim = (28,28)
    #resized = cv2.resize(image, dim, interpolation = cv2.INTER_AREA)
    #rint(image.astype('uint8'))
    #cv2.imwrite("image.png",image.astype('uint8'),(28, 28))
    #mask = np.zeros(np.shape(image), dtype=np.uint8)
    thresh = cv2.threshold(image, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
    #gray = cv2.cvtColor(thresh, cv2.COLOR_BGR2GRAY)
    
    cnts = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    cnts = cnts[0] if len(cnts) == 2 else cnts[1]
    (cnts, _) = contours.sort_contours(cnts, method="left-to-right")
    dim = (28, 28)
    for c in cnts:
        area = cv2.contourArea(c)
        #print(area)
        #if area < 120:
        bounding_boxes.append(cv2.boundingRect(c))
        print("for loop bb: ",bounding_boxes)
    
    boundingBoxes_filter =  [i for i in bounding_boxes if i != (0 , 0, 128, 128)]
    
    boundingBoxes = _group_rectangles(boundingBoxes_filter)
    #print(boundingBoxes)
    #
    #print(boundingBoxes_2)
    for (x, y, w, h) in boundingBoxes:
        print(x,y,w,h)
        ROI = thresh[y:y+h, x:x+w]
        ROI2 = cv2.bitwise_not(ROI)
        borderoutput = cv2.copyMakeBorder(ROI2, 30, 30, 30, 30, cv2.BORDER_CONSTANT, value=[0, 0, 0])
        
        resized = cv2.resize(borderoutput, dim, interpolation = cv2.INTER_AREA)
        cv2.imwrite('ROI_{}.png'.format(x), resized)
        #imageinv = cv2.bitwise_not(resized)
        img = Image.fromarray(resized.astype('uint8')).convert("RGB")
        
        pixel_values = processor(img, return_tensors="pt").pixel_values
        generated_ids = model.generate(pixel_values)
        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
        print(generated_text)
        generated_text_list.append(generated_text)
    #img = Image.fromarray(image.astype('uint8')).convert("RGB")
    #img = Image.open("image.png").convert("RGB")
    #print(img)
    
    # prepare image
    #pixel_values = processor(img, return_tensors="pt").pixel_values

    # generate (no beam search)
    #generated_ids = model.generate(pixel_values)

    # decode
    #generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return ''.join(generated_text_list)
    #return generated_text

title = "Interactive demo: Single Digits MNIST"
description = "Aico - University Utrecht"
iface = gr.Interface(fn=process_image, 
                     inputs="sketchpad", 
                     outputs="label",
                     title = title,
                     description = description)
iface.launch(debug=True)