app.py

import os
import numpy as np
import gradio as gr
import tensorflow as tf
from tensorflow import keras
from huggingface_hub import from_pretrained_keras

IMAGE_SIZE = (256, 256)

# Load model from HF
model = from_pretrained_keras(
    pretrained_model_name_or_path="fbadine/image-spam-detection"
)

####### To be removed
def get_txt_output(arr):
    txt = f"Image shape: {arr.shape}\n"
    txt += f"Min: {np.min(arr)}, Max: {np.max(arr)}\n"
    txt += f"Image hash: {hash(tuple(arr.reshape(-1)))}\n\n"

    return txt
####### End

# This is the predict function that takes as input an array-like-image and produces
# the probabilities that this image is either spam or ham
def predict(image):
####### To be removed
    txt_output = get_txt_output(image)
####### End
    # Resize image
    resized_image = keras.layers.Resizing(
        IMAGE_SIZE[0],
        IMAGE_SIZE[1],
        interpolation="bilinear",
        crop_to_aspect_ratio=True
    )(image)
    resized_image = tf.expand_dims(resized_image, axis=0)

####### To be removed
    txt_output += get_txt_output(resized_image.numpy())
####### End
    
    # Predict
    pred = model.predict(resized_image)
    prob = float(pred[0][0])

####### To be removed
    txt_output += f"Probability: {prob}"
####### End

    scoring_output = {
        "Spam": prob,
        "Ham": 1 - prob 
    }
    
    #return scoring_output
    return [scoring_output, txt_output]

# Clear Input and outpout
def clear_inputs_and_outputs():
    return [None, None, None]

# Main function
if __name__ == "__main__":
    demo = gr.Blocks()

    with demo:
        gr.Markdown(
            """
            <center><h1>Image Spam Detection</h1></center> \
            This space is a demo of a proof of concept POC Image Spam Detection<br> \
    In this space, you can upload an image to check if it's spam or not or you can use of the provided samples <br><br>
            """
        )

        with gr.Row():
            with gr.Column():
                # Input
                image_input = gr.Image(
                    shape=(256,256),
                    source="upload",
                    label="Upload an Image"
                )
                with gr.Row():
                    clr_btn = gr.Button(value="Clear", variant="secondary")
                    prd_btn = gr.Button(value="Predict")

            with gr.Column():
                # Output
                lbl_output = gr.Label(label="Prediction")
                ####### To be removed
                txt_output = gr.TextArea(label="info_out")
                ####### End

        clr_btn.click(
            fn=clear_inputs_and_outputs,
            inputs=[],
            outputs=[image_input, lbl_output],
        )
        prd_btn.click(
            fn=predict,
            inputs=[image_input],
            #outputs=[lbl_output],
            outputs=[lbl_output, txt_output],
        )

        gr.Examples(
            #examples=[os.path.join(os.path.dirname(__file__), "ShowLetter.jpg")],
            examples=[
                os.path.join(os.path.curdir, "examples", "sample1.jpg"),
                os.path.join(os.path.curdir, "examples", "sample2.jpg"),
            ],
            inputs=image_input,
            outputs=lbl_output,
            fn=predict,
            cache_examples=True,
        )

    demo.launch(debug=True, share=False)