## Daniel Buscombe, Marda Science LLC 2023
import gradio as gr
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from skimage.transform import resize
from skimage.io import imsave
from skimage.filters import threshold_otsu
from skimage.measure import EllipseModel, CircleModel, ransac

##========================================================
def fromhex(n):
    """hexadecimal to integer"""
    return int(n, base=16)

##========================================================
def label_to_colors(
    img,
    mask,
    alpha,  # =128,
    colormap,  # =class_label_colormap, #px.colors.qualitative.G10,
    color_class_offset,  # =0,
    do_alpha,  # =True
):
    """
    Take MxN matrix containing integers representing labels and return an MxNx4
    matrix where each label has been replaced by a color looked up in colormap.
    colormap entries must be strings like plotly.express style colormaps.
    alpha is the value of the 4th channel
    color_class_offset allows adding a value to the color class index to force
    use of a particular range of colors in the colormap. This is useful for
    example if 0 means 'no class' but we want the color of class 1 to be
    colormap[0].
    """

    colormap = [
        tuple([fromhex(h[s : s + 2]) for s in range(0, len(h), 2)])
        for h in [c.replace("#", "") for c in colormap]
    ]

    cimg = np.zeros(img.shape[:2] + (3,), dtype="uint8")
    minc = np.min(img)
    maxc = np.max(img)

    for c in range(minc, maxc + 1):
        cimg[img == c] = colormap[(c + color_class_offset) % len(colormap)]

    cimg[mask == 1] = (0, 0, 0)

    if do_alpha is True:
        return np.concatenate(
            (cimg, alpha * np.ones(img.shape[:2] + (1,), dtype="uint8")), axis=2
        )
    else:
        return cimg

##====================================
def standardize(img):
    # standardization using adjusted standard deviation

    N = np.shape(img)[0] * np.shape(img)[1]
    s = np.maximum(np.std(img), 1.0 / np.sqrt(N))
    m = np.mean(img)
    img = (img - m) / s
    del m, s, N
    #
    if np.ndim(img) == 2:
        img = np.dstack((img, img, img))

    return img

############################################################
############################################################

#load model
filepath = './saved_model'
model = tf.keras.models.load_model(filepath, compile = True)
model.compile

#segmentation
def segment(input_img, dims=(1024, 1024)):

    w = input_img.shape[0]
    h = input_img.shape[1]       
    img = standardize(input_img)
    
    img = resize(img, dims, preserve_range=True, clip=True) 
    img = np.expand_dims(img,axis=0)
    
    est_label = model.predict(img) 

    #Test Time Augmentation
    est_label2 = np.flipud(model.predict((np.flipud(img)), batch_size=1))
    est_label3 = np.fliplr(model.predict((np.fliplr(img)), batch_size=1))
    est_label4 = np.flipud(np.fliplr(model.predict((np.flipud(np.fliplr(img))))))

    #soft voting - sum the softmax scores to return the new TTA estimated softmax scores
    est_label = est_label + est_label2 + est_label3 + est_label4
    est_label /= 4
    
    pred = np.squeeze(est_label, axis=0)
    pred = resize(pred, (w, h), preserve_range=True, clip=True)
    
    bias=.1
    thres_land = threshold_otsu(pred[:,:,1])-bias
    print("Coin threshold: %f" % (thres_land))
    mask = (pred[:,:,1]<=thres_land).astype('uint8')
    
    imsave("greyscale.png", mask*255)

    class_label_colormap = [
        "#3366CC",
        "#DC3912",
        "#FF9900",
    ]
    # add classes
    class_label_colormap = class_label_colormap[:2]

    color_label = label_to_colors(
        mask,
        input_img[:, :, 0] == 0,
        alpha=128,
        colormap=class_label_colormap,
        color_class_offset=0,
        do_alpha=False,
    )
   
    imsave("color.png", color_label)
   
    #overlay plot
    plt.clf()
    plt.imshow(input_img,cmap='gray')
    plt.imshow(color_label, alpha=0.4)
    plt.axis("off")
    plt.margins(x=0, y=0)

    ############################################################
    dst = 1-mask.squeeze()
    points = np.array(np.nonzero(dst)).T
    points = np.column_stack((points[:,1], points[:,0]))

    # print("Fitting ellipse to coin to compute diameter ....")
    # model_robust, inliers = ransac(points, EllipseModel, min_samples=100,residual_threshold=2, max_trials=3)
    # r=np.max([model_robust.params[2] , model_robust.params[3]])
    # x=model_robust.params[0]
    # y=model_robust.params[1]
    # a_over_b = model_robust.params[2] / model_robust.params[3] ##a/b

    print("Fitting circle to coin to compute diameter ....")
    model_robust, inliers = ransac(points, CircleModel, min_samples=100,residual_threshold=2, max_trials=100)
    r=model_robust.params[2]
    x=model_robust.params[0]
    y=model_robust.params[1]

    print('diameter of coin = %f pixels' % (r*2))
    print('image scaling (assuming quarter dollar) = %f mm/pixel' % (24.26 / r*2))

    plt.plot(x, y, 'ko')
    plt.plot(np.arange(x-r, x+r, int(r*2)), np.arange(y-r, y+r, int(r*2)),'m')

    plt.savefig("overlay.png", dpi=300, bbox_inches="tight")

    return 'diameter of coin = %f pixels' % (r*2), 'image scaling (assuming quarter dollar) = %f mm/pixel' % (24.26 / r*2), color_label, plt , "greyscale.png", "color.png", "overlay.png"

title = "Find and measure coins in images of sand!" 
description = "This model demonstration segments beach sediment imagery into two classes: a) background, and b) coin, then measuring the coin. Allows upload of imagery and download of label imagery only one at a time. This model is part of the Doodleverse https://github.com/Doodleverse" 
    

examples = [['examples/20191011_091052.jpg'],
['examples/20191010_135020.jpg'],
['examples/IMG_20210922_170908944.jpg'],
['examples/IMG_20211121_120533257_HDR.jpg'],
['examples/20210208_172834.jpg'],
['examples/20220101_165359.jpg']] 

inp = gr.Image()
out1 = gr.Image(type='numpy')
out2 = gr.Plot(type='matplotlib')
out3 = gr.File()
out4 = gr.File()
out5 = gr.File()

Segapp = gr.Interface(segment, inp, ["text", "text", out1, out2, out3, out4, out5], title = title, description = description, examples=examples, theme="grass")
#, allow_flagging='manual', flagging_options=["bad", "ok", "good", "perfect"], flagging_dir="flagged")

Segapp.launch(enable_queue=True)