Spaces:
Sleeping
Sleeping
File size: 1,166 Bytes
8439d88 8be214a 8439d88 c59422d 8be214a 8439d88 c2cc89d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 |
from flask import Flask, request, jsonify
from huggingface_hub import from_pretrained_keras
import numpy as np
from PIL import Image
import io
import tensorflow as tf
app = Flask(__name__)
# Load the model
# model = from_pretrained_keras("MissingBreath/recycle-garbage-model")
# model = from_pretrained_keras("./recycle-garbage-model")
model = tf.keras.models.load_model('_9217')
# Class labels
# class_labels = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
@app.route('/classify', methods=['POST'])
def classify():
file = request.files['image']
if file:
img = Image.open(io.BytesIO(file.read()))
img = img.resize((128, 128))
img_array = np.array(img) / 255.0
img_array = np.expand_dims(img_array, axis=0)
predictions = model.predict(img_array)
predicted_class_idx = np.argmax(predictions)
# predicted_class = class_labels[predicted_class_idx]
# return jsonify({'prediction': predicted_class})
return jsonify({'prediction': predicted_class_idx})
else:
return jsonify({'error': 'No image provided'}), 400
if __name__ == '__main__':
app.run(debug=False) |