import gradio as gr
import tensorflow as tf
from tensorflow.keras.utils import CustomObjectScope
from tensorflow.keras.layers.experimental.preprocessing import RandomHeight

with CustomObjectScope({'RandomHeight': RandomHeight}):
    model_0 = tf.keras.models.load_model('/content/drive/MyDrive/bestmodel_porno_final_meilleure100%2.0.h5')

def classify_image(inp):
    inp = inp.reshape((-1, 224, 224, 3))
    prediction = model_0.predict(inp)
    output  = ""
    if prediction[0][prediction.argmax()] < 0.84:
      output = "bonne image"
    elif prediction.argmax() == 0:
      output = "Rifle violence"
    elif prediction.argmax() == 1:
      output = "guns violence"
    elif prediction.argmax() == 2:
      output = "knife violence"
    elif prediction.argmax() == 3:
      output = "image porno"
    elif prediction.argmax() == 4:
      output = "personne habillée" 
    else:
      output = "tank violence" 
    return output


image = gr.Image(height=224, width=224)

gr.Interface(
    fn=classify_image, inputs=image, outputs="text",live=True, theme="dark-peach",title="API de détection des images violentes",
).launch()