Spaces:
Sleeping
Sleeping
from flask import Flask, request, render_template, jsonify | |
from flask_cors import CORS | |
import numpy as np | |
from PIL import Image | |
import io | |
from tensorflow.keras.models import load_model | |
# Load the model | |
model = load_model('ecosort.h5') | |
# Define a dictionary to map class numbers to class names | |
class_mapping = { | |
0: 'battery', | |
1: 'biological', | |
2: 'brown-glass', | |
3: 'cardboard', | |
4: 'clothes', | |
5: 'green-glass', | |
6: 'metal', | |
7: 'paper', | |
8: 'plastic', | |
9: 'shoes', | |
10: 'trash', | |
11: 'white-glass' | |
} | |
app = Flask(__name__) | |
CORS(app) | |
def home(): | |
return render_template('index.html') | |
def predict(): | |
if 'file' not in request.files: | |
return "No file uploaded" | |
file = request.files['file'] | |
if file.filename == '': | |
return "No selected file" | |
img = Image.open(io.BytesIO(file.read())) | |
if img is None: | |
return "Invalid image file" | |
# Preprocess the image | |
img = img.resize((224, 224)) | |
img_array = np.asarray(img) / 255.0 | |
# Make predictions using your model | |
prediction = model.predict(np.expand_dims(img_array, axis=0)) | |
predicted_class = np.argmax(prediction) | |
# Get the class name from the dictionary | |
class_name = class_mapping.get(predicted_class, 'Unknown Class') | |
# Return the prediction result as JSON | |
return jsonify({'prediction': class_name}) | |
if __name__ == '__main__': | |
app.run(debug=True) | |