app.py

import keras
from keras.models import load_model
import gradio as gr
import cv2
import numpy as np
import torch  # Import torch here
from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import login
import os
import spaces

# Ensure the HF token is set
tok = os.getenv('HF_Token')
login(token=tok)

# Load Keras models
my_model = load_model('models/Final_Chicken_disease_model.h5', compile=True)
auth_model = load_model('models/auth_model.h5', compile=True)

# Load the tokenizer and model, ensuring they run on CPU
llama_tokenizer = AutoTokenizer.from_pretrained('meta-llama/Meta-Llama-3-8B-Instruct')
llama_model = AutoModelForCausalLM.from_pretrained(
    'meta-llama/Meta-Llama-3-8B-Instruct',
    device_map=None,  # Do not use device_map, force CPU
    torch_dtype=torch.float32  # Ensure model uses float32 for compatibility
).to('cpu')  # Force the model to CPU

# Dictionaries for disease names, results, and recommendations
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 of Medicine', 2: 'Paracetamol', 3: 'Ponston'}


def predict(image):
    # Preprocess the image for the authentication model
    image_check = cv2.resize(image, (224, 224))
    image_check = np.expand_dims(image_check, axis=0)  # Add batch dimension
    indx = auth_model.predict(image_check).argmax()

    if indx == 0:  # If the image is recognized as a chicken disease image
        # Preprocess the image for the disease prediction model
        image = cv2.resize(image, (224, 224))
        image = np.expand_dims(image, axis=0)  # Add batch dimension
        indx = my_model.predict(image).argmax()

        name = name_disease.get(indx)
        status = result.get(indx)
        recom = recommend.get(indx)
    else:  # If the image is not recognized as a chicken disease image
        name = 'Unknown Image'
        status = 'N/A'
        recom = 'N/A'

    return f"Chicken is {status}, the disease it has is {name}, the recommended medication is {recom}"


@spaces.GPU(duration=200)
def chat_response(user_input):
    inputs = llama_tokenizer(user_input, return_tensors='pt')
    outputs = llama_model.generate(inputs['input_ids'], max_length=500, do_sample=True)
    response = llama_tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response


def combined_interface(image, text):
    if image is not None:
        return predict(image)
    elif text:
        return chat_response(text)
    else:
        return "Please provide an image or ask a question."


example_image_path = 'disease.jpg'  # Replace with actual file path

# Check if the example image exists and is a file
if not os.path.isfile(example_image_path):
    print(f"Warning: Example image file not found: {example_image_path}")
    example_image_path = None  # Disable the example

# Only include the example if the path is valid
examples = [[example_image_path, ''], ['', 'What should I do if my chicken is sick?']] if example_image_path else None

# Create Gradio interface
interface = gr.Interface(
    fn=combined_interface,
    inputs=[gr.Image(label='Upload Image'), gr.Textbox(label='Ask a question')],
    outputs=gr.Textbox(label="Response"),
    # examples=examples  # Use examples only if valid
)

# Launch the interface
interface.launch(debug=True)