#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@Contact :   liuyuqi.gov@msn.cn
@Time    :   2024/06/04
@License :   Copyright © 2017-2022 liuyuqi. All Rights Reserved.
@Desc    :   
"""

import numpy as np
import tensorflow as tf
from deepexplain.tensorflow import DeepExplain
from tensorflow.keras import backend as K
from tensorflow.keras.models import Sequential, Model
import gradio as gr

model = tf.keras.applications.MobileNet()

def saliency(model, x, y):
    y = y.reshape(1, 1, 1, 1000)
    with DeepExplain(session=K.get_session()) as de:
        input_tensor = model.layers[0].input
        fModel = Model(inputs=input_tensor, outputs = model.layers[-3].output)
        target_tensor = fModel(input_tensor)

        attributions_gradin = de.explain('grad*input', target_tensor, input_tensor, x, ys=y)
        sal = np.sum(np.abs(attributions_gradin.squeeze()), axis=-1)
        sal = (sal - sal.min()) / (sal.max() - sal.min())
        return sal

inp = gr.inputs.ImageUpload()
out = gr.outputs.Label(label_names='imagenet1000', max_label_words=1, word_delimiter=",")

io = gr.Interface(inputs=inp,
                      outputs=out,
                      model=model,
                      model_type='keras',
                      saliency=saliency)

io.launch();