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" ) # 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): # Resize image resized_image = keras.layers.Resizing( IMAGE_SIZE[0], IMAGE_SIZE[1], interpolation="bilinear", crop_to_aspect_ratio=True )(image) # Add the batch axis resized_image = tf.expand_dims(resized_image, axis=0) # Predict pred = model.predict(resized_image) prob = float(pred[0][0]) scoring_output = { "Spam": prob, "Ham": 1 - prob } return scoring_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( """

Image Spam Detection

\ This space is a demo of a proof of concept POC Image Spam Detection
\ In this space, you can upload an image to check if it's spam or not or you can use of the provided samples

""" ) 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") 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], ) 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)