predict / app.py
aniket-manhas's picture
Upload folder using huggingface_hub
c473c6d verified
from flask import Flask, request, jsonify
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
import numpy as np
import io
from PIL import Image
import os
import requests
# Constants
MODEL_URL = "https://storage.googleapis.com/potato-model-detection/potato_disease_model.h5"
MODEL_PATH = "best_model.keras"
# Download model if not already present
def download_model():
if not os.path.exists(MODEL_PATH):
print("Downloading model...")
response = requests.get(MODEL_URL)
if response.status_code == 200:
with open(MODEL_PATH, 'wb') as f:
f.write(response.content)
print("Model downloaded successfully.")
else:
raise Exception(f"Failed to download model: {response.status_code}")
# Call download before loading the model
download_model()
model = load_model(MODEL_PATH)
class_names = ["Early Blight", "Late Blight", "Healthy"]
# Initialize Flask app
app = Flask(__name__)
# Image preprocessing function
def preprocess_image(img_bytes):
img = Image.open(io.BytesIO(img_bytes)).resize((256, 256))
img_array = image.img_to_array(img) / 255.0
img_array = np.expand_dims(img_array, axis=0)
return img_array
@app.route("/predict", methods=["POST"])
def predict():
if 'file' not in request.files:
return jsonify({"error": "No file part in request"}), 400
file = request.files['file']
if file.filename == '':
return jsonify({"error": "No file selected"}), 400
try:
img_bytes = file.read()
processed = preprocess_image(img_bytes)
prediction = model.predict(processed)
class_index = np.argmax(prediction[0])
result = class_names[class_index]
return jsonify({"prediction": result})
except Exception as e:
return jsonify({"error": str(e)}), 500
@app.route("/", methods=["GET"])
def home():
return "Potato Disease Detection API is running!"
if __name__ == "__main__":
app.run(host='0.0.0.0', port=5000)