from layers import BilinearUpSampling2D
from tensorflow.keras.models import load_model
from utils import load_images, predict
import matplotlib.pyplot as plt
import numpy as np
import gradio as gr
from huggingface_hub import from_pretrained_keras
import os
import sys

print('Loading model...')
model = from_pretrained_keras("mostafapasha/ribs-segmentation-model", compile=False)
print('Successfully loaded model...')
examples = ['examples/VinDr_RibCXR_val_008.png', 'examples/VinDr_RibCXR_val_013.png']


def segment(img_arr):
    if np.ndim(img_arr) != 2:
        img_arr = img_arr[:, :, 1]
    image = img_arr[np.newaxis, :, :, np.newaxis]
    threshold = 0.5
    logits = model(image)
    prob = tf.sigmoid(logits)
    pred = tf.cast(prob > threshold, dtype=tf.float32)
    pred = np.array(pred.numpy())[0,:,:,0]
    return pred
iface = gr.Interface(fn=segment, inputs=gr.inputs.Image(shape=(512, 512)), outputs="image", flagging_dir=os.path.join(sys.path[0], "flagged")).launch(share=True)