import gradio as gr
import tensorflow as tf
model_0 = tf.keras.models.load_model('bestmodel_porno_final_meilleure100%2.0.h5')
def classify_image(inp):
    inp = inp.reshape((-1, 224, 224, 3))
    prediction = model_0.predict(inp)
    if 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.inputs.Image(shape=(224, 224))
label = gr.outputs.Label(num_top_classes=3)

gr.Interface(
    fn=classify_image, inputs=image, outputs=label, interpretation="default"
).launch()