Spaces:
Sleeping
Sleeping
from flask import Flask, request, jsonify, send_file | |
from tensorflow.keras.models import load_model, Model | |
from PIL import Image | |
import numpy as np | |
import os | |
import cv2 | |
import tensorflow as tf | |
import firebase_admin | |
from firebase_admin import credentials, db | |
from datetime import datetime | |
app = Flask(__name__) | |
# β 1. Initialize Firebase | |
cred = credentials.Certificate("glaucoma-4b682-firebase-adminsdk-fbsvc-cd31fbe99d.json") # Path to your service account JSON | |
firebase_admin.initialize_app(cred, { | |
'databaseURL': 'https://glaucoma-4b682-default-rtdb.firebaseio.com/' | |
}) | |
results_ref = db.reference('results') # Will save results here | |
# β 2. Load the Model | |
model = load_model('mobilenet_glaucoma_model.h5', compile=False) | |
# β 3. Preprocess Image | |
def preprocess_image(img): | |
img = img.resize((224, 224)) | |
img = np.array(img) / 255.0 | |
img = np.expand_dims(img, axis=0) | |
return img | |
# β 4. Grad-CAM Generation | |
def make_gradcam(img_array, model, last_conv_layer_name='Conv2D_1'): | |
"""Generate Grad-CAM for the given image and model.""" | |
last_conv_layer = model.get_layer(last_conv_layer_name) | |
grad_model = Model(inputs=model.inputs, outputs=[last_conv_layer.output, model.output]) | |
with tf.GradientTape() as tape: | |
conv_outputs, predictions = grad_model(img_array) | |
loss = predictions[:, 0] | |
grads = tape.gradient(loss, conv_outputs) | |
pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2)) | |
conv_outputs = conv_outputs[0] | |
for i in range(conv_outputs.shape[-1]): | |
conv_outputs[..., i] *= pooled_grads[i] | |
heatmap = tf.reduce_mean(conv_outputs, axis=-1).numpy() | |
heatmap = np.maximum(heatmap, 0) | |
heatmap /= np.max(heatmap) | |
return heatmap | |
# β 5. Save Grad-CAM Overlay | |
def save_gradcam_image(original_img, heatmap, filename='gradcam.png', output_dir='results'): | |
"""Save the Grad-CAM overlay image and return its path.""" | |
if not os.path.exists(output_dir): | |
os.makedirs(output_dir) | |
img = np.array(original_img.resize((224, 224))) | |
heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0])) | |
heatmap = np.uint8(255 * heatmap) | |
heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET) | |
overlay = cv2.addWeighted(img, 0.6, heatmap, 0.4, 0) | |
filepath = os.path.join(output_dir, filename) | |
cv2.imwrite(filepath, overlay) | |
return filepath | |
def home(): | |
return "Glaucoma Detection Flask API is running!" | |
def test_file(): | |
"""Check if the Firebase service account JSON is present and readable.""" | |
filepath = "glaucoma-4b682-firebase-adminsdk-fbsvc-cd31fbe99d.json" | |
if os.path.exists(filepath): | |
return f"β Service account file found at: {filepath}" | |
else: | |
return "β Service account JSON NOT found." | |
def predict(): | |
"""Perform prediction and save results to Firebase.""" | |
if 'file' not in request.files: | |
return jsonify({'error': 'No file uploaded'}), 400 | |
file = request.files['file'] | |
if file.filename == '': | |
return jsonify({'error': 'No file selected'}), 400 | |
try: | |
img = Image.open(file.stream).convert('RGB') | |
img_array = preprocess_image(img) | |
prediction = model.predict(img_array)[0] | |
glaucoma_prob = 1 - prediction[0] | |
normal_prob = prediction[0] | |
result = 'Glaucoma' if glaucoma_prob > normal_prob else 'Normal' | |
confidence = float(glaucoma_prob) if result == 'Glaucoma' else float(normal_prob) | |
# Grad-CAM | |
heatmap = make_gradcam(img_array, model, last_conv_layer_name='Conv2D_1') | |
gradcam_filename = f"gradcam_{int(datetime.now().timestamp())}.png" | |
save_gradcam_image(img, heatmap, filename=gradcam_filename) | |
# Save to Firebase | |
results_ref.push({ | |
'image_filename': file.filename, | |
'prediction': result, | |
'confidence': confidence, | |
'gradcam_filename': gradcam_filename, | |
'timestamp': datetime.now().isoformat() | |
}) | |
return jsonify({ | |
'prediction': result, | |
'confidence': confidence, | |
'normal_probability': float(normal_prob), | |
'glaucoma_probability': float(glaucoma_prob), | |
'gradcam_image': gradcam_filename | |
}) | |
except Exception as e: | |
return jsonify({'error': str(e)}), 500 | |
def results(): | |
"""List all results from the Firebase database.""" | |
results_data = results_ref.get() | |
if not results_data: | |
results_data = [] | |
return jsonify(results_data) | |
def get_gradcam(filename): | |
"""Serve the Grad-CAM overlay image.""" | |
filepath = os.path.join('results', filename) | |
if os.path.exists(filepath): | |
return send_file(filepath, mimetype='image/png') | |
else: | |
return jsonify({'error': 'File not found'}), 404 | |
if __name__ == '__main__': | |
app.run(host='0.0.0.0', port=7860) | |