Update app.py

app.py

@@ -1,13 +1,21 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification
 
-# Load the breast cancer image classification model
-model = pipeline("image-classification", model="OverDriveLee/Breast_Cancer")
+# Load the feature extractor and model directly
+extractor = AutoFeatureExtractor.from_pretrained("ALM-AHME/beit-large-patch16-224-finetuned-BreastCancer-Classification-BreakHis-AH-60-20-20")
+model = AutoModelForImageClassification.from_pretrained("ALM-AHME/beit-large-patch16-224-finetuned-BreastCancer-Classification-BreakHis-AH-60-20-20")
 
-# Define the prediction function
+# Define the prediction function using the loaded model
 def classify_image(image):
-    results = model(image)
-    return {"prediction": results[0]["label"], "confidence": results[0]["score"]}
+    # Preprocess the image and get the features
+    inputs = extractor(images=image, return_tensors="pt")
+    # Make the prediction using the model
+    outputs = model(**inputs)
+    logits = outputs.logits
+    # 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}
 
 # Define the Gradio interface
 iface = gr.Interface(