File size: 1,681 Bytes
7bf97d3 3b24276 | 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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | # app.py
import io, os
from flask import Flask, request, render_template, jsonify
from PIL import Image
import torch
from model_utils import load_model, predict
app = Flask(__name__)
# Paths (adjust if you moved model files)
MODEL_DIR = os.path.join(os.path.dirname(__file__), "model")
CHECKPOINT_PATH = os.path.join(MODEL_DIR, "newplant_model_final.pth")
LABELS_PATH = os.path.join(MODEL_DIR, "labels.txt")
REMEDIES_PATH = os.path.join(MODEL_DIR, "remedies.json")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)
# Load model and metadata once at startup
model, labels, remedies = load_model(CHECKPOINT_PATH, LABELS_PATH, REMEDIES_PATH, device)
@app.route("/")
def index():
return render_template("index.html")
@app.route("/predict", methods=["POST"])
def predict_route():
if "file" not in request.files:
return jsonify({"error": "no file part"}), 400
file = request.files["file"]
if file.filename == "":
return jsonify({"error": "empty filename"}), 400
try:
img_bytes = file.read()
pil_img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
top_label, confidence, topk = predict(model, pil_img, labels, device, topk=5)
remedy = remedies.get(top_label, None)
response = {
"label": top_label,
"confidence": confidence,
"remedies": remedy,
"topk": [{"label": l, "confidence": float(c)} for l,c in topk]
}
return jsonify(response)
except Exception as e:
return jsonify({"error": str(e)}), 500
if __name__ == "__main__":
app.run(host="0.0.0.0", port=7860, debug=True)
|