tree-diagnosis-ai / model.py
karthikmn's picture
Create model.py
c48ca1e verified
raw
history blame contribute delete
758 Bytes
from transformers import pipeline
from PIL import Image
import io
from utils import explain_prediction, get_treatment_and_fertilizer, check_alert
classifier = pipeline("image-classification", model="huggingface/your-model")
def diagnose_disease(image_bytes):
image = Image.open(io.BytesIO(image_bytes))
prediction = classifier(image)[0]
label = prediction['label']
score = prediction['score']
explanation = explain_prediction(label)
treatment, fertilizer = get_treatment_and_fertilizer(label)
alert = check_alert(label)
return {
"prediction": label,
"confidence": score,
"explanation": explanation,
"treatment": treatment,
"fertilizer": fertilizer,
"alert": alert
}