import tensorflow as tf
import numpy as np
import cv2
from PIL import Image
import io

class InferenceHandler:
    def __init__(self):
        self.model = tf.saved_model.load("pneumonia_cnn_saved_model")
        self.class_names = ['PNEUMONIA', 'NORMAL']
        self.infer = self.model.signatures['serving_default']

    def preprocess(self, image):
        img = np.array(image.convert('L'))  # Grayscale
        img = cv2.resize(img, (150, 150))
        img = img / 255.0
        img = img.reshape(1, 150, 150, 1).astype(np.float32)
        return img

    def __call__(self, inputs):
        # inputs: dict with 'image' key (e.g., uploaded image)
        image = Image.open(io.BytesIO(inputs['image']))
        img_array = self.preprocess(image)
        prediction = self.infer(tf.convert_to_tensor(img_array))['dense_1'].numpy()
        class_id = (prediction > 0.5).astype("int32")[0][0]
        return {"prediction": self.class_names[class_id], "probability": float(prediction[0][0])}