import tensorflow as tf
from PIL import Image
import numpy as np
import requests
from io import BytesIO

# Load model
model = tf.keras.models.load_model('model.keras')

# Class names
CLASS_NAMES = ['Esophagitis', 'GERD', 'Normal', 'Ulcer']

def predict_from_url(image_url):
    response = requests.get(image_url)
    image = Image.open(BytesIO(response.content))
    return predict_image(image)

def predict_from_path(image_path):
    image = Image.open(image_path)
    return predict_image(image)

def predict_image(image):
    # Preprocess
    image = image.convert('RGB')
    image = image.resize((224, 224))
    image_array = np.array(image) / 255.0
    image_array = np.expand_dims(image_array, axis=0)
    
    # Predict
    predictions = model.predict(image_array)[0]
    
    # Format results
    results = {
        CLASS_NAMES[i]: float(predictions[i]) 
        for i in range(len(CLASS_NAMES))
    }
    
    predicted_class = CLASS_NAMES[np.argmax(predictions)]
    confidence = float(np.max(predictions))
    
    return {
        'predicted_class': predicted_class,
        'confidence': confidence,
        'all_predictions': results
    }

# Example usage
if __name__ == "__main__":
    result = predict_from_path('test_image.jpg')
    print(f"Prediction: {result['predicted_class']}")
    print(f"Confidence: {result['confidence']:.2%}")
    print(f"All probabilities: {result['all_predictions']}")