import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image
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('bestmodel.h5')

    
print("TensorFlow version:", tf.__version__)
print("GPU Available:", tf.config.list_physical_devices('GPU'))

# Configure TensorFlow to use memory growth
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)


def classify_image(inp):
    # Convert to PIL Image if we have a numpy array
    image = None
    if isinstance(inp, np.ndarray):
        image = Image.fromarray(inp)
    else:
        image = inp
    
    # Resize image to 224x224
    image = image.resize((224, 224), Image.Resampling.LANCZOS)
    
    # Convert to numpy array and ensure correct shape
    image_array = np.array(image)
    
    # Handle grayscale images
    if len(image_array.shape) == 2:
        image_array = np.stack([image_array] * 3, axis=-1)
    
    # Add batch dimension and ensure correct shape
    inp = image_array.reshape((-1, 224, 224, 3))
    # 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,title="API de détection des images violentes",
).launch()