fbadine's picture
Update app.py
efafc0d
raw history blame
No virus
2.55 kB
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):
# Add the batch axis
image = tf.expand_dims(image, axis=0)
# Predict
pred = model.predict(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(
"""
<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")
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"),
os.path.join(os.path.curdir, "examples", "sample3.jpg"),
],
inputs=image_input,
outputs=lbl_output,
fn=predict,
cache_examples=True,
)
demo.launch(debug=True, share=False)