| import tensorflow as tf | |
| import numpy as np | |
| import cv2 | |
| import os | |
| class TBClassifier: | |
| def __init__(self, model_path="tuberculosis_cnn_saved_model"): | |
| self.model = tf.saved_model.load(model_path) | |
| self.infer = self.model.signatures['serving_default'] | |
| self.class_names = ['Normal', 'Tuberculosis'] | |
| self.img_size = 256 | |
| def preprocess(self, image): | |
| if len(image.shape) == 3: | |
| image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) | |
| image = cv2.resize(image, (self.img_size, self.img_size)) / 255.0 | |
| image = image.reshape(1, self.img_size, self.img_size, 1).astype(np.float32) | |
| return image | |
| def predict(self, image): | |
| try: | |
| image = self.preprocess(image) | |
| input_key = list(self.infer.structured_input_signature[1].keys())[0] | |
| inputs = {input_key: tf.convert_to_tensor(image)} | |
| outputs = self.infer(**inputs) | |
| output_key = list(outputs.keys())[0] | |
| pred = outputs[output_key].numpy()[0][0] | |
| class_id = 1 if pred > 0.5 else 0 | |
| confidence = pred if class_id == 1 else 1 - pred | |
| return { | |
| "prediction": self.class_names[class_id], | |
| "confidence": float(confidence) | |
| } | |
| except Exception as e: | |
| return {"error": str(e)} | |