app.py

import gradio as gr
import joblib
from collections import Counter
from sklearn.preprocessing import OneHotEncoder
import numpy as np

# Load the models
model_dt = joblib.load("model_dt.pkl")
model_rf = joblib.load("model_rf.pkl")
model_nb = joblib.load("model_nb.pkl")
model_knn = joblib.load("model_knn.pkl")

# Initialize OneHotEncoder and fit it on the full list of symptoms (this should match the model's training data)
symptom_list = ["coughing", "diarrhoea", "fever", "frothing", "loss_of_appetite", "anorexia", "abdominal_pain", "bellowing"]
encoder = OneHotEncoder(handle_unknown='ignore', sparse_output=False)

# Fit the encoder on the complete list of possible symptoms (or the training dataset feature set)
encoder.fit(np.array(symptom_list).reshape(-1, 1))  # Assuming symptoms were the only categorical feature during training

# Treatment plan mapping based on predicted disease
def treatment_plan(disease):
    treatment_plans = {
        "peri_weaning_diarrhoea": "Consult a veterinarian for proper treatment.",
        "mastitis": "Administer antibiotics and consult a vet.",
        "bloat": "Ensure the cattle has a proper diet and avoid overfeeding.",
        "foot_rot": "Apply topical antiseptics and bandages.",
        # Add more diseases and their respective treatment plans here
    }
    return treatment_plans.get(disease, "Consult a veterinarian for proper treatment.")

# Unified prediction function for all models with averaging
def get_predictions(symptoms):
    # Apply One-Hot Encoding to symptoms (now that encoder is fitted)
    encoded_symptoms = encoder.transform(np.array(symptoms).reshape(-1, 1)).flatten()
    
    # Ensure the input matches the number of features the model expects
    expected_features = 93  # Adjust this based on the model's training setup (e.g., if using other features)
    if len(encoded_symptoms) != expected_features:
        raise ValueError(f"Expected {expected_features} features, but got {len(encoded_symptoms)} features.")
    
    # Get predictions from all four models
    prediction_dt = model_dt.predict([encoded_symptoms])[0]
    prediction_rf = model_rf.predict([encoded_symptoms])[0]
    prediction_nb = model_nb.predict([encoded_symptoms])[0]
    prediction_knn = model_knn.predict([encoded_symptoms])[0]
    
    # Calculate average prediction (for simplicity, using most common prediction)
    predictions = [prediction_dt, prediction_rf, prediction_nb, prediction_knn]
    count = Counter(predictions)
    most_common_disease = count.most_common(1)[0][0]
    
    return prediction_dt, prediction_rf, prediction_nb, prediction_knn, most_common_disease

# Create Gradio Interface
def predict_disease(cattle_id, breed, symptom1, symptom2, symptom3, symptom4, symptom5, ward):
    symptoms = [symptom1, symptom2, symptom3, symptom4, symptom5]
    
    # Get predictions from models and calculate average prediction
    prediction_dt, prediction_rf, prediction_nb, prediction_knn, average_prediction = get_predictions(symptoms)
    
    # Get the treatment plan based on the disease predicted
    treatment = treatment_plan(average_prediction)
    
    # Format the predictions and treatment for output
    prediction_results = f"Predictions: \nDecision Tree: {prediction_dt} \nRandom Forest: {prediction_rf} \nNaive Bayes: {prediction_nb} \nK-Nearest Neighbour: {prediction_knn} \nAverage Prediction: {average_prediction}"
    
    return prediction_results, treatment

# Define the Gradio interface with inputs for Cattle ID, breed, symptoms, and ward
iface = gr.Interface(fn=predict_disease, 
                     inputs=[gr.Textbox(label="Enter Cattle ID"), 
                             gr.Dropdown(["Gir", "Girolando", "Holstein Friesian", "Jersey", "Kankrej", "Rathi", "Red Sindhi", "Sahiwal", "Tharparkar"], label="Select Cow Breed"),
                             gr.Dropdown(symptom_list, label="Select Symptom 1"),
                             gr.Dropdown(symptom_list, label="Select Symptom 2"),
                             gr.Dropdown(symptom_list, label="Select Symptom 3"),
                             gr.Dropdown(symptom_list, label="Select Symptom 4"),
                             gr.Dropdown(symptom_list, label="Select Symptom 5"),
                             gr.Dropdown(["Ward A", "Ward B", "Ward C"], label="Assign Ward")],
                     outputs=[gr.Textbox(label="Prediction Results"), gr.Textbox(label="Suggested Treatment")])

# Launch the interface
iface.launch()