Update app.py

app.py

@@ -20,12 +20,9 @@ def predict(image):
         logits = model(**inputs).logits
     probs = F.softmax(logits, dim=1)[0].tolist()
 
-    # Prepare output dict
+    # Return dictionary: class name -> confidence (no extra keys!)
     labels = [model.config.id2label[i] for i in range(len(probs))]
-    result = {labels[i]: float(f"{probs[i]:.6f}") for i in range(len(labels))}
-    predicted_idx = int(torch.argmax(logits, dim=1)[0])
-    result["predicted_label"] = model.config.id2label[predicted_idx]
-    return result
+    return {labels[i]: float(f"{probs[i]:.6f}") for i in range(len(labels))}
 
 # Gradio Interface
 def main():