services/disease_detection.py

# services/disease_detection.py

import cv2
import numpy as np
import tensorflow as tf
from keras.models import load_model
from services.llama_service import llama2_response
import logging

# Initialize logger for better error tracking
logger = logging.getLogger(__name__)

# Load TensorFlow/Keras models with GPU support if available, otherwise use CPU
try:
    device_name = '/GPU:0' if len(tf.config.list_physical_devices('GPU')) > 0 else '/CPU:0'
    with tf.device(device_name):
        # Load models, ensure no weight_decay is passed to the optimizer
        my_model = load_model('models/Final_Chicken_disease_model.h5', compile=False)  # Disable auto-compiling
        auth_model = load_model('models/auth_model.h5', compile=False)
        print(f"Models loaded successfully on {device_name}.")
except Exception as e:
    logger.error(f"Error loading models: {e}")
    my_model = None  # Ensure there's a fallback in case model loading fails

# Mapping of disease labels to disease names, health status, and treatment recommendations
name_disease = {0: 'Coccidiosis', 1: 'Healthy', 2: 'New Castle Disease', 3: 'Salmonella'}
status_map = {0: 'Critical', 1: 'No issue', 2: 'Critical', 3: 'Critical'}
recommendations = {
    0: 'Panadol',  # Treatment for Coccidiosis
    1: 'No treatment needed',  # Healthy case
    2: 'Percetamol',  # Treatment for New Castle Disease
    3: 'Ponston'  # Treatment for Salmonella
}


def detect_disease(image, db):
    """
    Function to create an instance of PoultryFarmBot and detect disease.
    :param image: Image of poultry feces for disease detection.
    :param db: Database connection to pass into the bot.
    :return: Result of disease detection.
    """
    bot = PoultryFarmBot(db)
    return bot.diagnose_disease(image)


class PoultryFarmBot:
    def __init__(self, db):
        """
        Initializes the PoultryFarmBot with a MongoDB connection.
        :param db: MongoDB connection for future use in logging or storing results.
        """
        self.db = db

    def preprocess_image(self, image):
        """
        Preprocesses the input image for model prediction.
        Resizes the image to 224x224 and expands the dimensions to match the model input format.

        :param image: Input image (numpy array).
        :return: Preprocessed image ready for prediction or None if an error occurs.
        """
        try:
            image_resized = cv2.resize(image, (224, 224))  # Resize image
            image_expanded = np.expand_dims(image_resized, axis=0)  # Add batch dimension
            return image_expanded
        except Exception as e:
            logger.error(f"Error in image preprocessing: {e}")
            return None

    def predict(self, image):
        """
        Predicts the disease from the input image using the pre-trained model.

        :param image: Input image (numpy array).
        :return: Tuple of predicted disease name, health status, and treatment recommendation.
        """
        image_preprocessed = self.preprocess_image(image)

        if image_preprocessed is None:
            return "Image preprocessing failed.", None, None, None

        if my_model is None:
            logger.error("Disease detection model is not loaded.")
            return "Model not loaded.", None, None, None

        try:
            # Predict disease using the loaded model
            indx = my_model.predict(image_preprocessed).argmax()

            # Retrieve the disease name, status, and recommended treatment
            disease_name = name_disease.get(indx, "Unknown disease")
            status = status_map.get(indx, "Unknown condition")
            recommendation = recommendations.get(indx, "No recommendation available")

            return disease_name, status, recommendation
        except Exception as e:
            logger.error(f"Error during disease prediction: {e}")
            return "Prediction failed.", None, None, None

    def generate_disease_response(self, disease_name, status, recommendation):
        """
        Generates a detailed response using Llama 2 model for disease information and recommendations.

        :param disease_name: The name of the detected disease.
        :param status: The health status of the poultry.
        :param recommendation: The treatment recommendation.
        :return: A detailed response with disease information.
        """
        try:
            prompt = (
                f"The disease detected is {disease_name}, classified as {status}. "
                f"Recommended action: {recommendation}. "
                f"Here is some information about {disease_name}: causes, symptoms, and treatment methods "
                "to effectively manage this condition on a poultry farm."
            )
            response = llama2_response(prompt)
            return response.replace(prompt, "").strip()
        except Exception as e:
            logger.error(f"Error generating detailed response: {e}")
            return "Unable to generate detailed response."

    def diagnose_disease(self, image):
        """
        Diagnose a disease using a fecal image.

        :param image: The image of poultry feces to detect a disease.
        :return: Dictionary containing diagnosis result and detailed response if applicable.
        """
        if image is not None and image.size > 0:  # Ensure the image is valid
            disease_name, status, recommendation = self.predict(image)
            if disease_name and status and recommendation:
                detailed_response = self.generate_disease_response(disease_name, status, recommendation)
                return {
                    "disease": disease_name,
                    "status": status,
                    "recommendation": recommendation,
                    "detailed_response": detailed_response
                }
        return {"error": "Please provide a valid image for disease detection."}