File size: 994 Bytes
1b672d5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
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])}
|