Spaces:
Sleeping
Sleeping
| 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) | |
| def home(): | |
| return "API is running!" | |
| def favicon(): | |
| return '', 204 # or serve a real favicon if needed | |
| 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) | |