Spaces:
Runtime error
Runtime error
from flask import Flask, request, jsonify | |
from flask_cors import CORS | |
import tensorflow as tf | |
import numpy as np | |
from io import BytesIO | |
from PIL import Image | |
app = Flask(__name__) | |
CORS(app) | |
# Load models for each plant | |
models = { | |
'potato': tf.keras.models.load_model('./potato_model.h5'), | |
'tomato': tf.keras.models.load_model('./tomato_model.h5'), | |
'grape': tf.keras.models.load_model('./grape_model.h5'), | |
'corn': tf.keras.models.load_model('./corn_model.h5'), | |
'pepper': tf.keras.models.load_model('./pepper_model.h5') | |
} | |
# Function to load and preprocess image | |
def load_image(file): | |
img = Image.open(BytesIO(file)) | |
img = img.resize((224, 224)) | |
img_array = tf.keras.preprocessing.image.img_to_array(img) | |
img_array = np.expand_dims(img_array, axis=0) # Convert single image to a batch | |
img_array = img_array / 255.0 # Normalize the image as done in preprocessing | |
return img_array | |
# Class labels for each plant model | |
class_labels = { | |
'potato': ["Early_Blight", "Healthy", "Late_Blight"], | |
'tomato': ["Bacterial_spot", "Early_blight", "Late_blight", "Leaf_Mold", "Septoria_leaf_spot", "Spider_mites Two-spotted_spider_mite", "Target_Spot", "_Yellow_Leaf_Curl_Virus", "_mosaic_virus", "healthy"], | |
'grape': ["Black Rot", "ESCA", "Healthy", "Leaf Blight"], | |
'corn': ["Blight", "Common_Rust", "Gray_Leaf_Spot", "Healthy"], | |
'pepper': ["Bacterial_spot", "healthy"] | |
} | |
# Prediction function | |
def predict(file, plant): | |
img = load_image(file) | |
model = models[plant] | |
predictions = model.predict(img) | |
predicted_class_idx = np.argmax(predictions, axis=1)[0] | |
confidence = predictions[0][predicted_class_idx] | |
return class_labels[plant][predicted_class_idx], confidence | |
# Check if the file is an allowed image type | |
def allowed_file(filename): | |
allowed_extensions = {'png', 'jpg', 'jpeg'} | |
return '.' in filename and filename.rsplit('.', 1)[1].lower() in allowed_extensions | |
# Routes for each plant | |
def predict_potato(): | |
return predict_route('potato') | |
def predict_tomato(): | |
return predict_route('tomato') | |
def predict_grape(): | |
return predict_route('grape') | |
def predict_corn(): | |
return predict_route('corn') | |
def predict_pepper(): | |
return predict_route('pepper') | |
# Common function to handle prediction route | |
def predict_route(plant): | |
if 'file' not in request.files: | |
return jsonify({"error": "No file part"}), 400 | |
file = request.files['file'] | |
if file.filename == '': | |
return jsonify({"error": "No selected file"}), 400 | |
if not allowed_file(file.filename): | |
return jsonify({"error": "Unsupported file type"}), 400 | |
try: | |
predicted_class, confidence = predict(file.read(), plant) | |
return jsonify({"predicted_class": predicted_class, "confidence": confidence}) | |
except Exception as e: | |
return jsonify({"error": str(e)}), 500 | |
# Run the Flask app | |
#if __name__ == '__main__': | |
# app.run(debug=True) | |