pneumonia-cnn-model / handler.py
sukhmani1303's picture
Upload handler.py with huggingface_hub
1b672d5 verified
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])}