Update app.py

app.py

@@ -14,12 +14,22 @@ API_URL = "https://api-inference.huggingface.co/models/riabayonaor/modelo_predic
 META_LLAMA_API_URL = "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-8B-Instruct"
 headers = {"Authorization": f"Bearer {huggingface_api_key}"}
 
+# Lista de clases del modelo
+classes = ["Anthracnose", "Bacterial Wilt", "Belly Rot", "Downy Mildew",
+           "Fresh Cucumber", "Fresh Leaf", "Gummy Stem Blight", "Pythium Fruit Rot"]
+
 def query(image_bytes):
     response = requests.post(API_URL, headers=headers, data=image_bytes)
+    if response.status_code != 200:
+        st.error(f"Error al clasificar la imagen: {response.status_code}")
+        return None
     return response.json()
 
 def llama_query(prompt):
     response = requests.post(META_LLAMA_API_URL, headers=headers, json={"inputs": prompt})
+    if response.status_code != 200:
+        st.error(f"Error al obtener respuesta del modelo Meta Llama: {response.status_code}")
+        return None
     return response.json()
 
 def main():
@@ -42,23 +52,30 @@ def main():
         # Enviar la imagen al modelo de Hugging Face
         predictions = query(img_byte_arr)
 
-        if "error" not in predictions:
-            # Suponiendo que las predicciones están en el formato [{label: "label1", score: 0.95}, {label: "label2", score: 0.05}]
-            top_prediction = max(predictions, key=lambda x: x["score"])
-            st.write(f"Predicción principal: {top_prediction['label']} con confianza {top_prediction['score']:.2f}")
+        if predictions is not None:
+            if "error" not in predictions:
+                # Suponiendo que las predicciones están en el formato [{label: "0", score: 0.95}, {label: "1", score: 0.05}]
+                top_prediction = max(predictions, key=lambda x: x["score"])
+                class_label = classes[int(top_prediction['label'])]
+                st.write(f"Predicción principal: {class_label} con confianza {top_prediction['score']:.2f}")
 
-            # Usar la etiqueta principal para el modelo de Meta Llama
-            prompt = f"Esta enfermedad es {top_prediction['label']}. Explica qué es y sugiere posibles insecticidas o soluciones."
+                # Usar la etiqueta principal para el modelo de Meta Llama
+                prompt = f"Esta enfermedad es {class_label}. Explica qué es y sugiere posibles insecticidas o soluciones."
 
-            # Llamar al modelo Meta Llama
-            llama_response = llama_query(prompt)
-            
-            if "error" not in llama_response:
-                explanation = llama_response[0]["generated_text"]
-                st.write("Explicación y posibles soluciones:")
-                st.write(explanation)
+                # Llamar al modelo Meta Llama
+                llama_response = llama_query(prompt)
+                
+                if llama_response is not None:
+                    if "error" not in llama_response:
+                        explanation = llama_response[0]["generated_text"]
+                        st.write("Explicación y posibles soluciones:")
+                        st.write(explanation)
+                    else:
+                        st.write("No se pudo obtener una respuesta del modelo Meta Llama.")
+                else:
+                    st.write("No se pudo obtener una respuesta del modelo Meta Llama.")
             else:
-                st.write("No se pudo obtener una respuesta del modelo Meta Llama.")
+                st.write("No se pudo clasificar la imagen.")
         else:
             st.write("No se pudo clasificar la imagen.")