app.py

import tensorflow as tf
import numpy as np
from PIL import Image
import io
import base64
import json

class MRIClassifier:
    def __init__(self):
        # Load the model
        self.model = tf.keras.models.load_model("./model")
        self.img_size = (224, 224)  # EfficientNet standard size
        print(f"Model loaded with input size: {self.img_size}")
        
    def preprocess(self, image):
        # Resize image to expected input shape
        image = image.resize(self.img_size)
        
        # Convert to RGB if not already
        if image.mode != "RGB":
            image = image.convert("RGB")
            
        # Convert to array and normalize
        image_array = np.array(image) / 255.0
        
        # Add batch dimension
        image_array = np.expand_dims(image_array, axis=0)
        
        return image_array
        
    def predict(self, image_bytes):
        # Convert bytes to image
        image = Image.open(io.BytesIO(image_bytes))
        
        # Preprocess image
        image_array = self.preprocess(image)
        
        # Make prediction
        prediction = self.model.predict(image_array)
        pred_value = float(prediction[0][0])
        
        # Format result
        predicted_class = "VAD-Demented" if pred_value > 0.5 else "Non-Demented"
        
        return [
            {
                "label": predicted_class,
                "score": pred_value
            }
        ]

# Initialize classifier
classifier = MRIClassifier()

def inference(model_inputs):
    # Handle both direct image inputs and base64 encoded images
    if isinstance(model_inputs, dict) and "image" in model_inputs:
        # Base64 encoded image
        image_bytes = base64.b64decode(model_inputs["image"])
    else:
        # Direct image bytes (from API)
        image_bytes = model_inputs
        
    results = classifier.predict(image_bytes)
    return results