import cv2
import numpy as np
import logging
from tensorflow.keras.models import load_model

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class TBImageProcessor:
    """Processes real TB saliva images using a trained CNN model"""

    def __init__(self, model_path="tb_cnn_model.h5"):
        try:
            self.model = load_model(model_path)
            logger.info("TB Image Processor Model Loaded Successfully.")
        except Exception as e:
            logger.error(f"Failed to load TB Image Model: {e}")
            self.model = None

    def process_image(self, image_path):
        """Analyze the TB image and return risk assessment."""
        if not self.model:
            return {"error": "Model not loaded. Cannot process image."}

        try:
            image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
            image = cv2.resize(image, (128, 128))  # Resizing for CNN input
            image = np.expand_dims(image, axis=[0, -1]) / 255.0  # Normalize

            prediction = self.model.predict(image)
            confidence = float(prediction[0][0])
            result = "TB Detected" if confidence > 0.5 else "No TB"
            
            return {
                "result": result,
                "confidence": confidence
            }
        except Exception as e:
            logger.error(f"Error processing image: {e}")
            return {"error": "Image processing failed."}