app.py

from flask import Flask, render_template, request, jsonify
from flask_cors import CORS
import tensorflow as tf
import numpy as np
import json
import io
import base64
from PIL import Image
import os
from tensorflow.keras.layers import Layer

# ✅ Define custom Cast layer
class Cast(Layer):
    def call(self, inputs):
        return tf.cast(inputs, tf.float32)

app = Flask(__name__)
CORS(app)

# Configure upload settings
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max file size

# Load model and class labels
try:
    model = tf.keras.models.load_model(
        'pneumonia_model.h5',
        compile=False,
        custom_objects={'Cast': Cast}  # ✅ properly reference Cast class
    )
    print("✅ Model loaded successfully!")
    
    with open('class_labels.json', 'r') as f:
        class_labels = json.load(f)
    print(f"✅ Class labels loaded: {class_labels}")
    
except Exception as e:
    print(f"❌ Could not load model files: {e}")
    model = None
    class_labels = ["NORMAL", "PNEUMONIA"]

def preprocess_image(image):
    try:
        if image.mode != 'RGB':
            image = image.convert('RGB')
        img_size = (224, 224)
        image = image.resize(img_size)
        img_array = np.array(image)
        img_array = img_array.astype('float32') / 255.0
        img_array = np.expand_dims(img_array, axis=0)
        return img_array
    except Exception as e:
        print(f"Error preprocessing image: {e}")
        return None

def predict_pneumonia(image):
    if model is None:
        return {
            'predicted_class': 'Unknown',
            'confidence': 0.0,
            'error': 'Model not loaded'
        }
    
    try:
        img_array = preprocess_image(image)
        if img_array is None:
            return {
                'predicted_class': 'Unknown',
                'confidence': 0.0,
                'error': 'Image preprocessing failed'
            }
        
        prediction = model.predict(img_array)
        predicted_class_idx = int(prediction[0] > 0.5)
        predicted_class = class_labels[predicted_class_idx]
        confidence = float(max(prediction[0], 1 - prediction[0]))
        
        return {
            'predicted_class': predicted_class,
            'confidence': confidence,
            'probability_normal': float(1 - prediction[0]),
            'probability_pneumonia': float(prediction[0])
        }
        
    except Exception as e:
        print(f"Prediction error: {e}")
        return {
            'predicted_class': 'Unknown',
            'confidence': 0.0,
            'error': str(e)
        }

@app.route('/')
def home():
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    try:
        if 'image' not in request.files:
            return jsonify({'error': 'No image uploaded'}), 400
        file = request.files['image']
        if file.filename == '':
            return jsonify({'error': 'No image selected'}), 400
        allowed_extensions = {'png', 'jpg', 'jpeg', 'gif', 'bmp'}
        if not ('.' in file.filename and 
                file.filename.rsplit('.', 1)[1].lower() in allowed_extensions):
            return jsonify({'error': 'Invalid file type. Please upload PNG, JPG, JPEG, GIF, or BMP'}), 400
        image = Image.open(file.stream)
        result = predict_pneumonia(image)
        return jsonify(result)
    except Exception as e:
        return jsonify({'error': f'Processing error: {str(e)}'}), 500

@app.route('/predict_base64', methods=['POST'])
def predict_base64():
    try:
        data = request.get_json()
        if 'image' not in data:
            return jsonify({'error': 'No image data provided'}), 400
        image_data = data['image']
        if image_data.startswith('data:image'):
            image_data = image_data.split(',')[1]
        image_bytes = base64.b64decode(image_data)
        image = Image.open(io.BytesIO(image_bytes))
        result = predict_pneumonia(image)
        return jsonify(result)
    except Exception as e:
        return jsonify({'error': f'Processing error: {str(e)}'}), 500

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)