lp128396's picture
Update app.py
54fb183 verified
from flask import Flask, request, jsonify
from flask_cors import CORS
import tensorflow as tf
import numpy as np
from PIL import Image
import io
app = Flask(__name__)
CORS(app)
interpreter = tf.lite.Interpreter(model_path="final_resnet.tflite")
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
class_names = {
0: 'FreshApple', 1: 'FreshBanana', 2: 'FreshGrape', 3: 'FreshGuava',
4: 'FreshJujube', 5: 'FreshOrange', 6: 'FreshPomegranate', 7: 'FreshStrawberry',
8: 'RottenApple', 9: 'RottenBanana', 10: 'RottenGrape', 11: 'RottenGuava',
12: 'RottenJujube', 13: 'RottenOrange', 14: 'RottenPomegranate', 15: 'RottenStrawberry'
}
def preprocess_image(img_bytes):
img = Image.open(io.BytesIO(img_bytes)).resize((224, 224))
img = np.array(img) / 255.0
return np.expand_dims(img, axis=0)
@app.route('/')
def home():
return "API is running!"
@app.route('/favicon.ico')
def favicon():
return '', 204 # or serve a real favicon if needed
@app.route('/predict', methods=['POST'])
def predict():
file = request.files.get('image')
if not file:
return jsonify({'error': 'No image provided'}), 400
img = preprocess_image(file.read()).astype(np.float32)
interpreter.set_tensor(input_details[0]['index'], img)
interpreter.invoke()
pred = interpreter.get_tensor(output_details[0]['index'])[0]
sorted_indices = np.argsort(pred)[::-1]
top1, top2 = sorted_indices[:2]
top1_label = class_names[top1]
top2_label = class_names[top2]
top1_conf = pred[top1] * 100
top2_conf = pred[top2] * 100
result = {}
# Extract categories
top1_cat = "Fresh" if "Fresh" in top1_label else "Rotten"
top2_cat = "Fresh" if "Fresh" in top2_label else "Rotten"
# Format prediction output (without fruit name)
prediction_text = f"{top1_cat} ({top1_conf:.1f}%)"
# Rule 1: If prediction > 80%
if top1_conf >= 80:
result['prediction'] = prediction_text
result['message'] = "✅ Safe to eat!" if top1_cat == "Fresh" else "⚠️ Not safe to eat!"
# Rule 3: < 80%, top two different categories
elif top1_cat != top2_cat:
result['prediction'] = f"{top1_cat} ({top1_conf:.1f}%) vs {top2_cat} ({top2_conf:.1f}%)"
if top1_cat == "Rotten":
result['message'] = "⚠️ Do not eat this fruit."
else:
result['message'] = "🍃 Seems fresh, but be cautious."
# Rule 4: < 80%, same category (e.g., FreshApple vs FreshBanana)
elif top1_cat == top2_cat:
result['prediction'] = prediction_text
if top1_cat == "Rotten":
result['message'] = "⚠️ Do not eat this fruit."
else:
result['message'] = "🍃 Seems fresh, but be cautious."
# Fallback message (if none matched)
else:
result['prediction'] = prediction_text
result['message'] = "⚠️ Unable to confidently predict freshness."
# Optional extra info
result['confidence'] = round(top1_conf, 2)
result['fruitType'] = top1_label
result['fresh'] = 'Fresh' in top1_label
return jsonify(result)