import cv2
from fastai.vision.all import *
import numpy as np
import gradio as gr
from scipy import ndimage

fnames = get_image_files("./albumentations/original")
def label_func(fn): return "./albumentations/labelled/"f"{fn.stem}.png"
codes = np.loadtxt('labels.txt', dtype=str)
w, h = 768, 1152
img_size = (w,h)
im_size = (h,w)

dls = SegmentationDataLoaders.from_label_func(
    ".", bs=3, fnames = fnames, label_func = label_func, codes = codes,
    item_tfms=Resize(img_size)
)

learn = unet_learner(dls, resnet34)
learn.load('learn')

def segmentImage(img_path):
    img = cv2.imread(img_path, 0)
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            if img[i][j] > 0:
                img[i][j] = 1
    kernel = np.ones((3,3), np.uint8)
    img = cv2.erode(img, kernel, iterations=1)
    img = cv2.dilate(img, kernel, iterations=1)
    img = ndimage.binary_fill_holes(img).astype(int)
    labels, nlabels = ndimage.label(img)
    colors = np.random.randint(0, 255, (nlabels + 1, 3))
    colors[0] = 0
    img_color = colors[labels]
    return img_color

def predict_segmentation(img):
    gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    resized_img = cv2.resize(gray_img, im_size)
    pred = learn.predict(resized_img)
    color_pred = pred[0].show(ctx=None, cmap='gray', alpha=None)
    color_pred_array = color_pred_array.astype(np.uint8)
    output_image = Image.fromarray(color_pred_array)
    # Save the image to a temporary file
    temp_file = 'temp.png'
    output_image.save(temp_file)
    # Call the segmentImage function
    segmented_image = segmentImage(temp_file)
    return output_image, segmented_image

input_image = gr.inputs.Image()
output_image1 = gr.outputs.Image(type='pil')
output_image2 = gr.outputs.Image(type='pil')
app = gr.Interface(fn=predict_segmentation, inputs=input_image, outputs=[output_image1, output_image2], title='Microstructure Segmentation', description='Segment the input image into grain and background.')
app.launch()