smart_trash / app.py
salim4n's picture
ready to test
75c8efc
import streamlit as st
from tensorflow.keras.models import load_model
import numpy as np
from PIL import Image
import cv2
from tensorflow.keras.preprocessing.image import img_to_array, load_img
@st.cache_data()
def load():
model_path = "best_model.h5"
model = load_model(model_path, compile=False)
return model
# chargement du model
model = load()
def predict(upload):
img = Image.open(upload)
img = np.asarray(img)
img_resize = cv2.resize(img, (224, 224))
img_resize = np.expand_dims(img_resize, axis=0)
pred = model.predict(img_resize)
rec = pred[0][0]
return rec
def draw():
#rectangle sur la prediction
img = cv2.imread(upload)
img = cv2.resize(img, (224, 224))
img = cv2.rectangle(img, (0, 0), (224, 224), (0, 255, 0), 3)
cv2.imwrite('output.png', img)
st.title("Poubelle Intelligente")
upload = st.file_uploader("Charger Image", type=["pnj", "jpeg", "jpg"])
c1, c2 = st.columns(2)
if upload:
rec = predict(upload)
prob_rec = predict(upload) * 100
prob_org = (1 - rec) * 100
c1.image(Image.open(upload))
if prob_rec > 50:
c2.write(f"Je suis certains à {prob_rec:.2f} % que ceci est recyclable")
else:
c2.write(f"Je suis certains à {prob_org:.2f} % que ceci ne soit pas recyclable")