sojaML / app.py
monabarreiro's picture
Update app.py
1da0099 verified
Raw
History Blame Contribute Delete
1.3 kB
from transformers import pipeline
import gradio as gr
from transformers import AutoConfig
label_map = {
"0": "Necrosis/rust ",
"1": "Muerto",
"2": "Antracnosis/Soy Disease Classifier",
"3": "Healthy",
"4": "Marchitez bacteriana"}
pipe = pipeline(
"image-classification",
model="sbaner24/vit-base-patch16-224-Soybean_11-46"
)
def classify_image(image):
results = pipe(image)
return {label_map.get(r['label'], r['label']): float(r['score']) for r in results}
def predict(image):
img_tensor = preprocess(image).unsqueeze(0)
with torch.no_grad():
outputs = model(img_tensor)
probs = torch.nn.functional.softmax(outputs, dim=1)
return {class_names[i]: float(probs[0][i]) for i in range(len(class_names))}
with gr.Blocks() as demo:
gr.Markdown("Clasificador de enfermedades de soja")
gr.Markdown("Sube una imagen de una hoja de soja para detectar posibles enfermedades.")
gr.HTML(""" <a href="http://localhost:8081/enfSoja" target="_blank">Mostrar Enfermedades</a>""")
with gr.Row():
inp = gr.Image(type="pil", label="Sub铆 una imagen")
out = gr.Textbox(label="Predicci贸n")
btn = gr.Button("Predecir")
btn.click(fn=classify_image, inputs=inp, outputs=out)
demo.launch(debug=True)