Dyno1307's picture
Update src/app.py
ffcec64 verified
from flask import Flask, render_template, request
from tensorflow.keras.models import load_model
import numpy as np
from PIL import Image
import io
import cv2
import os
import tensorflow as tf
import time
app = Flask(__name__)
# Define the labels for classification
CLASSIFICATION_LABELS = ['Crown and Root Rot', 'Healthy Wheat', 'Leaf Rust', 'Wheat Loose Smut']
# Lazy load model to avoid startup freeze
model = None
def get_model():
global model
if model is None:
model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models/classification_model.h5')
model = load_model(model_path)
return model
@app.route('/')
def index():
return render_template('index.html')
@app.route('/predict', methods=['POST'])
def predict():
if 'file' not in request.files:
return 'No file part'
file = request.files['file']
if file.filename == '':
return 'No selected file'
if file:
# Load model only when needed
classification_model = get_model()
# Read and preprocess the image
img = Image.open(io.BytesIO(file.read()))
img_np = np.array(img)
img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
img_resized = cv2.resize(img_np, (224, 224))
img_preprocessed = tf.keras.applications.vgg19.preprocess_input(
np.reshape(img_resized, (1, 224, 224, 3))
)
# Run prediction
prediction = classification_model.predict(img_preprocessed)
label_index = np.argmax(prediction)
label = CLASSIFICATION_LABELS[label_index]
# Save the output image
timestamp = str(int(time.time()))
output_image_filename = f'output_{timestamp}.jpg'
output_image_path = os.path.join('static', output_image_filename)
cv2.imwrite(output_image_path, img_bgr)
return render_template('result.html', image_path=output_image_filename, label=label)
@app.route('/health')
def health():
return "App is running fine ✅"
if __name__ == '__main__':
port = int(os.environ.get('PORT', 5000))
app.run(host='0.0.0.0', port=port)