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