utils.py

# utils.py

import tensorflow as tf
from keras.models import load_model
import cv2
import logging
import numpy as np
from transformers import AutoModelForCausalLM, AutoTokenizer
from datetime import datetime
from werkzeug.security import generate_password_hash, check_password_hash

logger = logging.getLogger(__name__)

def load_model_with_device(model_path, device_name):
    """
    Load a Keras model on the specified device.

    Args:
        model_path (str): Path to the model file.
        device_name (str): Device to load the model on ('/GPU:0' or '/CPU:0').

    Returns:
        Model object if successful, None otherwise.
    """
    try:
        logger.info(f"Loading model from '{model_path}' on {device_name}.")
        with tf.device(device_name):
            model = load_model(model_path, compile=True)
            logger.info(f"Model '{model_path}' loaded successfully on {device_name}.")
            return model
    except Exception as e:
        logger.error(f"Error loading model '{model_path}': {e}")
        return None

def llama3_response(user_input, tokenizer, model):
    """
    Generate a response using the Llama 3.2 model.

    Args:
        user_input (str): Input prompt for the Llama 3.2 model.
        tokenizer (AutoTokenizer): Tokenizer to preprocess the input.
        model (AutoModelForCausalLM): Pre-trained language model for generating responses.

    Returns:
        str: Generated response from the model.
    """
    try:
        logger.info("Generating response using Llama 3.2 model.")
        # Tokenize the input prompt
        inputs = tokenizer(user_input, return_tensors="pt", truncation=True, max_length=150, padding=True)
        # Generate a response using the Llama 3.2 model
        outputs = model.generate(
            inputs["input_ids"],
            max_length=150,
            do_sample=True,
            temperature=0.7,
            pad_token_id=tokenizer.pad_token_id,
            attention_mask=inputs["attention_mask"]
        )
        # Decode the generated response
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        logger.info("Response generation successful.")
        return response
    except Exception as e:
        logger.error(f"Error generating response: {str(e)}")
        return f"Error generating response: {str(e)}"

class PoultryFarmBot:
    def __init__(self, db):
        """
        Initialize the PoultryFarmBot with a MongoDB database connection.

        Args:
            db (MongoClient): MongoDB database instance.
        """
        self.db = db  # MongoDB database for future use
        logger.info("PoultryFarmBot initialized with MongoDB connection.")

    def preprocess_image(self, image):
        """
        Preprocess the input image for disease detection.

        Args:
            image (numpy.ndarray): Input image to preprocess.

        Returns:
            numpy.ndarray: Preprocessed image ready for model input.
        """
        try:
            logger.info("Preprocessing image for disease detection.")
            # Resize the image to match the model input size (224, 224)
            image_check = cv2.resize(image, (224, 224))
            # Add a batch dimension to the image
            image_check = np.expand_dims(image_check, axis=0)
            logger.info("Image preprocessing successful.")
            return image_check
        except Exception as e:
            logger.error(f"Error in image preprocessing: {e}")
            return None

    def predict_disease(self, image, model, name_disease, result, recommend):
        """
        Predict the disease from the given poultry fecal image.

        Args:
            image (numpy.ndarray): Input image to predict.
            model (tf.keras.Model): Loaded Keras model for disease prediction.
            name_disease (dict): Dictionary mapping class indices to disease names.
            result (dict): Dictionary mapping class indices to status.
            recommend (dict): Dictionary mapping class indices to recommendations.

        Returns:
            tuple: Detailed response, disease name, status, and recommendation.
        """
        logger.info("Starting disease prediction.")
        image_check = self.preprocess_image(image)
        if image_check is None:
            logger.warning("Image preprocessing failed.")
            return "Image preprocessing failed.", None, None, None

        try:
            logger.info("Running model prediction.")
            indx = model.predict(image_check).argmax()
            name = name_disease.get(indx, "Unknown disease")
            status = result.get(indx, "unknown condition")
            recom = recommend.get(indx, "no recommendation available")
            logger.info(f"Disease prediction successful: {name}, Status: {status}, Recommendation: {recom}")
            detailed_response = self.generate_disease_response(name, status, recom)
            return detailed_response, name, status, recom
        except Exception as e:
            logger.error(f"Error during disease prediction: {e}")
            return "Error during prediction.", None, None, None

    def generate_disease_response(self, disease_name, status, recommendation):
        """
        Generate a detailed response about the detected disease.

        Args:
            disease_name (str): Name of the detected disease.
            status (str): Status of the disease (e.g., critical or no issue).
            recommendation (str): Recommended action.

        Returns:
            str: Detailed response generated by Llama 3.2 model.
        """
        logger.info(f"Generating detailed response for disease: {disease_name}")
        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."
        )
        return prompt

    def log_enquiry(self, enquiry_type, content, response, user_id, enquiries_collection):
        """
        Log a farmer's enquiry in the database.

        Args:
            enquiry_type (str): Type of the enquiry ('image' or 'text').
            content (str): The content of the enquiry.
            response (str): The response given by the system.
            user_id (str): The ID of the user making the enquiry.
            enquiries_collection (MongoClient): Collection to store farmer enquiries.
        """
        enquiry = {
            "user_id": user_id,
            "enquiry_type": enquiry_type,
            "content": content,
            "response": response,
            "timestamp": datetime.utcnow()
        }
        logger.info(f"Logging enquiry: {enquiry}")
        enquiries_collection.insert_one(enquiry)

    def authenticate_user(self, username, password, users_collection):
        """
        Authenticate a user with username and password.

        Args:
            username (str): Username of the user.
            password (str): Password of the user.
            users_collection (MongoClient): Collection to store user credentials.

        Returns:
            dict: User information if authentication is successful, None otherwise.
        """
        logger.info(f"Authenticating user: {username}")
        user = users_collection.find_one({"username": username})
        if user and check_password_hash(user['password'], password):
            logger.info("Authentication successful.")
            return user
        logger.warning("Authentication failed.")
        return None

    def register_user(self, username, password, users_collection):
        """
        Register a new user with username and password.

        Args:
            username (str): Username of the new user.
            password (str): Password of the new user.
            users_collection (MongoClient): Collection to store user credentials.

        Returns:
            bool: True if registration is successful, False otherwise.
        """
        logger.info(f"Registering user: {username}")
        if users_collection.find_one({"username": username}):
            logger.warning("Username already exists.")
            return False
        hashed_password = generate_password_hash(password)
        users_collection.insert_one({"username": username, "password": hashed_password})
        logger.info("User registration successful.")
        return True