Spaces:
Sleeping
Sleeping
File size: 675 Bytes
3f98cbd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
import numpy as np
from tensorflow.keras.preprocessing.image import img_to_array
IMG_SIZE = 128
# Removed 'trash' from the list
CLASS_NAMES = ['cardboard', 'glass', 'metal', 'paper', 'plastic']
def preprocess_image(img):
img = img.resize((IMG_SIZE, IMG_SIZE))
img = img_to_array(img)
img = img / 255.0
img = np.expand_dims(img, axis=0) # Add batch dimension
return img
def predict_image(model, image):
processed = preprocess_image(image)
prediction = model.predict(processed)
class_idx = np.argmax(prediction)
confidence = float(np.max(prediction))
label = CLASS_NAMES[class_idx]
return label, confidence, class_idx, processed
|