update space

app.py

@@ -1,258 +1,249 @@
-# File: app.py
-
-# Import necessary libraries
 import os
-import logging
-from pymongo import MongoClient, errors
-from datetime import datetime
-from werkzeug.security import generate_password_hash
-from services.utils import PoultryFarmBot
-from services.chatbot import build_chatbot_interface
+import tensorflow as tf
+from keras.models import load_model
+import gradio as gr
+import cv2
+import numpy as np
+from huggingface_hub import login
+from pymongo import MongoClient
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from flask import Flask, render_template, request, redirect, url_for, session
-from celery import Celery
-import time
-import threading
-import gunicorn.app.base
-from six import iteritems
-
-# Setup logging for better monitoring
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-
-# Flask app setup
-app = Flask(__name__)
-app.secret_key = os.environ.get("SECRET_KEY", "default_secret_key")
 
-# Celery configuration
-CELERY_BROKER_URL = os.environ.get("CELERY_BROKER_URL", "redis://localhost:6379/0")
-app.config['CELERY_BROKER_URL'] = CELERY_BROKER_URL
-app.config['CELERY_RESULT_BACKEND'] = CELERY_BROKER_URL
-celery = Celery(app.name, broker=app.config['CELERY_BROKER_URL'])
-celery.conf.update(app.config)
+# Ensure the Hugging Face token is set
+tok = os.environ.get('HF_Token')
+if tok:
+    print("Logging in to Hugging Face with provided token...")
+    login(token=tok, add_to_git_credential=True)
+    print("Login successful.")
+else:
+    print("Warning: Hugging Face token not found in environment variables.")
 
-# MongoDB Setup for logging and audit
+# MongoDB Setup (for inventory, record-keeping, etc.)
+print("Setting up MongoDB connection...")
 MONGO_URI = os.environ.get("MONGO_URI")
-if not MONGO_URI:
-    logger.error("MONGO_URI is not set in the environment variables.")
-    raise ValueError("MONGO_URI environment variable is required but not set.")
-
-# Retry logic for MongoDB connection
-max_retries = 3
-for attempt in range(max_retries):
-    try:
-        logger.info(f"Connecting to MongoDB (Attempt {attempt + 1}/{max_retries}).")
-        client = MongoClient(MONGO_URI, serverSelectionTimeoutMS=5000)  # Timeout after 5 seconds
-        client.server_info()  # Trigger exception if cannot connect
-        db = client.poultry_farm  # Connect to the 'poultry_farm' database
-        enquiries_collection = db.enquiries  # Collection to store farmer enquiries
-        users_collection = db.users  # Collection to store user credentials
-        logs_collection = db.logs  # Collection to store application logs
-        logger.info("Successfully connected to MongoDB.")
-        break
-    except errors.ServerSelectionTimeoutError as e:
-        logger.error(f"Failed to connect to MongoDB (Attempt {attempt + 1}/{max_retries}): {e}")
-        if attempt < max_retries - 1:
-            time.sleep(5)  # Wait for 5 seconds before retrying
-        else:
-            logger.critical("Could not connect to MongoDB after multiple attempts. Exiting application.")
-            raise ConnectionError("Could not connect to MongoDB. Please check the MONGO_URI and ensure the database is running.")
-
-def log_to_db(level, message):
-    try:
-        log_entry = {
-            "level": level,
-            "message": message,
-            "timestamp": datetime.utcnow()
-        }
-        logs_collection.insert_one(log_entry)
-        logger.debug(f"Logged to database: {log_entry}")
-    except Exception as e:
-        logger.error(f"Failed to log to database: {e}")
-
-# Override logger methods to also log to MongoDB
-class MongoHandler(logging.Handler):
-    def emit(self, record):
-        log_entry = self.format(record)
+client = MongoClient(MONGO_URI)
+db = client.poultry_farm  # Database
+print("MongoDB connection established.")
+
+# Check GPU availability for TensorFlow
+print("Checking TensorFlow setup...")
+print("TensorFlow version:", tf.__version__)
+print("Eager execution:", tf.executing_eagerly())
+print("TensorFlow GPU Available:", tf.config.list_physical_devices('GPU'))
+
+# Set TensorFlow to use mixed precision with available GPU
+from tensorflow.keras import mixed_precision
+if len(tf.config.list_physical_devices('GPU')) > 0:
+    # Set mixed precision policy to use float16 for better performance on GPU
+    policy = mixed_precision.Policy('mixed_float16')
+    mixed_precision.set_global_policy(policy)
+    print("Using mixed precision with GPU")
+else:
+    print("Using CPU without mixed precision")
+
+# Load TensorFlow/Keras models with GPU support if available, otherwise use CPU
+try:
+    # Select device based on GPU availability
+    device_name = '/GPU:0' if len(tf.config.list_physical_devices('GPU')) > 0 else '/CPU:0'
+    print(f"Loading models on {device_name}...")
+    with tf.device(device_name):
+        # Load the poultry disease detection model
+        my_model = load_model('models/disease_model.h5', compile=True)
+        print("Disease detection model loaded successfully.")
+        # Load the authentication model
+        auth_model = load_model('models/auth_model.h5', compile=True)
+        print("Authentication model loaded successfully.")
+        print(f"Models loaded successfully on {device_name}.")
+except Exception as e:
+    print(f"Error loading models: {e}")
+
+# Updated Disease names and recommendations based on fecal analysis
+name_disease = {0: 'Coccidiosis', 1: 'Healthy', 2: 'New Castle Disease', 3: 'Salmonella'}
+result = {0: 'Critical', 1: 'No issue', 2: 'Critical', 3: 'Critical'}
+recommend = {
+    0: 'Panadol',
+    1: 'You have no need Medicine',
+    2: 'Percetamol',
+    3: 'Ponston'
+}
+
+class PoultryFarmBot:
+    def __init__(self):
+        self.db = db  # MongoDB database for future use
+        print("PoultryFarmBot initialized with MongoDB connection.")
+
+    # Image Preprocessing for Fecal Disease Detection
+    def preprocess_image(self, image):
         try:
-            log_to_db(record.levelname, log_entry)
+            print("Preprocessing image for disease detection...")
+            # Resize the image to match model input size (224x224)
+            image_check = cv2.resize(image, (224, 224))
+            # Add batch dimension to the image array
+            image_check = np.expand_dims(image_check, axis=0)
+            print("Image preprocessing successful.")
+            return image_check
         except Exception as e:
-            logger.error(f"Failed to emit log to MongoDB: {e}")
-
-mongo_handler = MongoHandler()
-mongo_handler.setLevel(logging.INFO)
-formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
-mongo_handler.setFormatter(formatter)
-logger.addHandler(mongo_handler)
-
-# Hugging Face token setup
-tok = os.environ.get('HF_TOKEN')
-if tok:
-    # Log in to Hugging Face using the token from environment variables
-    logger.info("Logging in to Hugging Face.")
-else:
-    logger.warning("Hugging Face token not found in environment variables.")
+            print(f"Error in image preprocessing: {e}")
+            return None
+
+    # Predict Disease from Fecal Image
+    def predict(self, image):
+        print("Starting disease prediction...")
+        # Preprocess the image before prediction
+        image_check = self.preprocess_image(image)
+        if image_check is None:
+            print("Image preprocessing failed.")
+            return "Image preprocessing failed.", None, None, None
+
+        # Predict using the fecal disease detection model
+        try:
+            print("Running model prediction...")
+            indx = my_model.predict(image_check).argmax()
+            print(f"Prediction complete. Predicted index: {indx}")
+            name = name_disease.get(indx, "Unknown disease")
+            status = result.get(indx, "unknown condition")
+            recom = recommend.get(indx, "no recommendation available")
+
+            # Generate additional information about the disease using Llama 2
+            detailed_response = self.generate_disease_response(name, status, recom)
+            print("Generated detailed response using Llama 2.")
+            return detailed_response, name, status, recom
+        except Exception as e:
+            print(f"Error during prediction: {e}")
+            return "Error during prediction.", None, None, None
+
+    # Generate a detailed response using Llama 2 for disease information and recommendations
+    def generate_disease_response(self, disease_name, status, recommendation):
+        print("Generating detailed disease response...")
+        # Create a prompt for Llama 3 to generate detailed disease information
+        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)
+        # Post-process to remove the prompt if accidentally included in the response
+        final_response = response.replace(prompt, "").strip()
+        print("Detailed disease response generated.")
+        return final_response
+
+    # Diagnose Disease Using Fecal Image
+    def diagnose_disease(self, image):
+        print("Diagnosing disease from provided image...")
+        # Ensure image is valid and has elements
+        if image is not None and image.size > 0:
+            return self.predict(image)
+        print("Invalid image provided.")
+        return "Please provide an image of poultry fecal matter for disease detection.", None, None, None
 
 # Initialize the bot instance
-logger.info("Initializing PoultryFarmBot instance.")
-bot = PoultryFarmBot(db)
-logger.info("PoultryFarmBot instance initialized.")
-
-# Global model and tokenizer variables
-model = None
-tokenizer = None
-model_loading_event = threading.Event()
-
-# Celery task to load the model and tokenizer
-@celery.task
-def load_model_and_tokenizer():
-    global model, tokenizer
-    try:
-        logger.info("Loading Llama 3.2 model and tokenizer.")
-        model_name = "meta-llama/Llama-3.2-3B"
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForCausalLM.from_pretrained(model_name)
-        if tokenizer.pad_token is None:
-            logger.info("Adding padding token to tokenizer.")
-            tokenizer.add_special_tokens({'pad_token': '[PAD]'})
-            model.resize_token_embeddings(len(tokenizer))
-        logger.info("Model and tokenizer loaded successfully.")
-        model_loading_event.set()  # Set the event to indicate that the model is loaded
-    except Exception as e:
-        logger.error(f"Failed to load Llama 3.2 model or tokenizer: {e}")
-        model_loading_event.set()  # Set the event even if loading fails to prevent indefinite waiting
-        raise RuntimeError("Could not load the Llama 3.2 model or tokenizer. Please check the configuration.")
-
-# Authentication function
-def authenticate_user(username, password):
-    """
-    Authenticate a user with username and password.
-
-    Args:
-        username (str): Username for authentication.
-        password (str): Password for authentication.
-
-    Returns:
-        bool: True if authentication is successful, False otherwise.
-    """
-    logger.info(f"Authenticating user: {username}")
-    user = bot.authenticate_user(username, password)
-    if user:
-        logger.info(f"Authentication successful for user: {username}")
-        return True, user
-    else:
-        logger.warning(f"Authentication failed for user: {username}")
-        return False, None
-
-# Registration function
-def register_user(username, password):
-    """
-    Register a new user with username and password.
-
-    Args:
-        username (str): Username for registration.
-        password (str): Password for registration.
-
-    Returns:
-        bool: True if registration is successful, False otherwise.
-    """
-    logger.info(f"Registering user: {username}")
+print("Initializing PoultryFarmBot instance...")
+bot = PoultryFarmBot()
+
+# Load Llama 2 model and tokenizer for text generation
+print("Loading Llama 3.2 model and tokenizer...")
+model_name = "meta-llama/Llama-3.2-1B"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+print("Llama 3.2 model and tokenizer loaded successfully.")
+
+# Set the padding token to EOS token or add a new padding token
+if tokenizer.pad_token is None:
+    print("Adding pad token to tokenizer...")
+    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+    model.resize_token_embeddings(len(tokenizer))
+    print("Pad token added and model resized.")
+
+# Define Llama 2 response generation
+def llama2_response(user_input):
     try:
-        if users_collection.find_one({"username": username}):
-            logger.warning("Username already exists: %s", username)
-            return False
-        hashed_password = generate_password_hash(password)
-        users_collection.insert_one({"username": username, "password": hashed_password})
-        logger.info("User registered successfully: %s", username)
-        return True
+        print("Generating response using Llama 2...")
+        # Tokenize user input for the Llama 2 model
+        inputs = tokenizer(user_input, return_tensors="pt", truncation=True, max_length=150, padding=True)
+        # Generate a response using the Llama 2 model
+        outputs = model.generate(
+            inputs["input_ids"],
+            max_length=150,
+            do_sample=True,
+            temperature=0.7,
+            pad_token_id=tokenizer.pad_token_id,  # Use the newly set padding token
+            attention_mask=inputs["attention_mask"]
+        )
+
+        # Decode and return the response
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        print("Response generated successfully.")
+        return response
     except Exception as e:
-        logger.error(f"Failed to register user: {e}")
-        return False
-
-# Flask routes for AdminLTE authentication
-@app.route('/', methods=['GET', 'POST'])
-def login():
-    if request.method == 'POST':
-        username = request.form['username']
-        password = request.form['password']
-        logger.info(f"Login attempt for user: {username}")
-        success, user = authenticate_user(username, password)
-        if success:
-            logger.info("Authentication successful for user: %s", username)
-            session['username'] = username
-
-            # Start loading the model asynchronously after successful login
-            logger.info("Starting model loading asynchronously.")
-            load_model_and_tokenizer.apply_async()
-
-            return redirect(url_for('chatbot'))
-        else:
-            logger.warning("Authentication failed for user: %s", username)
-            return render_template('login.html', error="Invalid username or password.")
-    return render_template('login.html')
-
-@app.route('/register', methods=['GET', 'POST'])
-def register():
-    if request.method == 'POST':
-        username = request.form['username']
-        password = request.form['password']
-        logger.info(f"Registration attempt for user: {username}")
-        success = register_user(username, password)
-        if success:
-            logger.info("Registration successful for user: %s", username)
-            return redirect(url_for('login'))
+        print(f"Error generating response: {e}")
+        return f"Error generating response: {str(e)}"
+
+# Main chatbot function: handles both generative AI and disease detection
+def chatbot_response(image, text):
+    print("Received user input for chatbot response...")
+    # If an image is provided, perform disease detection
+    if image is not None:
+        print("Image provided, attempting disease diagnosis...")
+        diagnosis, name, status, recom = bot.diagnose_disease(image)
+        if name and status and recom:
+            print("Diagnosis complete, returning detailed response.")
+            return diagnosis
         else:
-            logger.warning("Registration failed for user: %s", username)
-            return render_template('register.html', error="Username already exists. Please choose a different one.")
-    return render_template('register.html')
-
-@app.route('/chatbot')
-def chatbot():
-    if 'username' not in session:
-        logger.warning("Access to chatbot attempted without login. Redirecting to login page.")
-        return redirect(url_for('login'))
-    username = session['username']
-
-    # Wait until the model is loaded with a timeout mechanism
-    logger.info("Waiting for model to load.")
-    if not model_loading_event.wait(timeout=100):  # Timeout after 100 seconds
-        logger.error("Model loading timed out for user: %s", username)
-        return "Model loading timed out. Please try again later."
-
-    if model is None or tokenizer is None:
-        logger.error("Model failed to load for user: %s", username)
-        return "Model failed to load. Please try again later."
-
-    logger.info("Launching Gradio chatbot interface for user: %s", username)
-    chatbot_interface = build_chatbot_interface({'username': username})
-    return chatbot_interface.launch(inline=True, share=True)
-
-# Custom Gunicorn application to run Flask safely in a production environment like Hugging Face Space Docker
-class StandaloneApplication(gunicorn.app.base.BaseApplication):
-    def __init__(self, app, options=None):
-        self.application = app
-        self.options = options or {}
-        super().__init__()
-
-    def load_config(self):
-        config = {key: value for key, value in iteritems(self.options) if key in self.cfg.settings and value is not None}
-        for key, value in config.items():
-            self.cfg.set(key.lower(), value)
-
-    def load(self):
-        return self.application
-
-# Launch the Flask application using Gunicorn in a production-safe manner
+            print("Diagnosis failed or incomplete, returning diagnostic message only.")
+            return diagnosis  # Return only the diagnostic message if no disease found
+    else:
+        print("No image provided, using Llama 3.2 for text response...")
+        # Use Llama 2 for more accurate responses to user text queries
+        return llama2_response(text)
+
+# Gradio interface styling and layout with ChatGPT-like theme
+print("Setting up Gradio interface...")
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="green", neutral_hue="slate")) as chatbot_interface:
+    gr.Markdown("# 🐔 Poultry Management Chatbot")
+    gr.Markdown(
+        "This chatbot can help you manage your poultry with conversational AI.  Upload an image of poultry fecal matter for disease detection or just ask questions!"
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            # Image input for poultry feces (optional)
+            fecal_image = gr.Image(
+                label="Upload Image of Poultry Feces (Optional)",
+                type="numpy",
+                elem_id="image-upload",
+                show_label=True,
+            )
+        with gr.Column(scale=2):
+            # Text input for user questions
+            user_input = gr.Textbox(
+                label="Type your question or chat with the assistant",
+                placeholder="Ask a question about poultry management...",
+                lines=3,
+                elem_id="user-input",
+            )
+
+    # Output box to display chatbot responses
+    output_box = gr.Textbox(
+        label="Response",
+        placeholder="The response will appear here...",
+        interactive=False,
+        lines=10,
+        elem_id="output-box",
+    )
+
+    # Submit button to trigger response generation
+    submit_button = gr.Button(
+        "Submit",
+        variant="primary",
+        elem_id="submit-button"
+    )
+    submit_button.click(
+        fn=chatbot_response,
+        inputs=[fecal_image, user_input],
+        outputs=[output_box]
+    )
+
+# Launch the Gradio interface
 if __name__ == "__main__":
-    try:
-        logger.info("Launching Flask server for AdminLTE authentication.")
-        options = {
-            'bind': '%s:%s' % ('0.0.0.0', '5000'),
-            'workers': 2,  # Set number of worker processes
-            'threads': 4,  # Set number of threads per worker
-            'timeout': 120,  # Set timeout for workers
-        }
-        StandaloneApplication(app, options).run()
-    except Exception as e:
-        logger.error(f"Failed to launch Flask server: {e}")
-        raise RuntimeError("Could not launch the Flask server. Please check the application setup.")
+    print("Launching Gradio interface...")
+    chatbot_interface.queue().launch(debug=True)