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)}