Update app.py

app.py

@@ -15,15 +15,20 @@ def classify_image(image):
     # Get the predicted label and confidence
     predicted_label = logits.argmax(dim=1).item()
     confidence = logits.softmax(dim=1).max().item()
-    return {"prediction": predicted_label, "confidence": confidence}
+
+    # Map predicted label to "benigno" or "maligno"
+    class_names = ["benigno", "maligno"]
+    predicted_class = class_names[predicted_label]
+
+    return {"prediction": predicted_class, "confidence": confidence}
 
 # Define the Gradio interface
 iface = gr.Interface(
     fn=classify_image,
     inputs=gr.inputs.Image(),
     outputs="json",
-    title="Breast Cancer Image Classification",
-    description="This application classifies breast cancer images into different classes."
+    title="Classificação de Imagens de Câncer de Mama",
+    description="Este aplicativo classifica imagens de câncer de mama em diferentes classes."
 )
 
 # Launch the Gradio interface