AzizWazir's picture
Update app.py
71a27aa verified
import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
import joblib
# Load the pre-trained model and scaler
model = joblib.load("water_quality_model.pkl") # Ensure the model file exists
scaler = joblib.load("scaler.pkl") # Ensure the scaler file exists
# Define water quality categories
def classify_quality(score):
if score >= 0.8:
return "Excellent"
elif score >= 0.6:
return "Good"
elif score >= 0.4:
return "Fair"
else:
return "Poor"
# Define the prediction function
def predict_water_quality(pH, turbidity, dissolved_oxygen, nitrate_levels):
# Prepare the input data
data = np.array([[pH, turbidity, dissolved_oxygen, nitrate_levels]])
scaled_data = scaler.transform(data)
# Predict the quality score
score = model.predict(scaled_data)[0]
quality = classify_quality(score)
# Generate insights (graphical output)
plt.figure(figsize=(6, 4))
categories = ["Excellent", "Good", "Fair", "Poor"]
scores = [0.8, 0.6, 0.4, 0.2]
colors = ["green", "blue", "orange", "red"]
plt.bar(categories, scores, color=colors, alpha=0.7)
plt.axhline(y=score, color="black", linestyle="--", label=f"Predicted Score: {score:.2f} ({quality})")
plt.legend()
plt.title("Water Quality Prediction")
plt.xlabel("Categories")
plt.ylabel("Score")
plt.tight_layout()
# Save the graph to a file
graph_path = "water_quality_graph.png"
plt.savefig(graph_path)
plt.close()
return quality, graph_path
# Define the Gradio interface
interface = gr.Interface(
fn=predict_water_quality,
inputs=[
gr.inputs.Number(label="pH"),
gr.inputs.Number(label="Turbidity (NTU)"),
gr.inputs.Number(label="Dissolved Oxygen (mg/L)"),
gr.inputs.Number(label="Nitrate Levels (mg/L)"),
],
outputs=[
gr.outputs.Textbox(label="Water Quality"),
gr.outputs.Image(label="Water Quality Graph")
],
title="Water Quality Prediction",
description="Enter water parameters to predict the quality of water. Results include the classification and a visual representation of quality."
)
# Launch the application
if __name__ == "__main__":
interface.launch()