Update app.py

app.py

@@ -1,19 +1,17 @@
 from flask import Flask, request, jsonify
-from transformers import AutoModelForImageClassification, AutoProcessor
+from transformers import pipeline
 from PIL import Image
 import io
 import fitz  # PyMuPDF
 import os
-import torch
-import torch.nn.functional as F
 from werkzeug.utils import secure_filename
 
 app = Flask(__name__)
 
-# Load model and processor
+# Load model and processor using pipeline
 model_name = "AsmaaElnagger/Diabetic_RetinoPathy_detection"
-model = AutoModelForImageClassification.from_pretrained(model_name)
-processor = AutoProcessor.from_pretrained(model_name)
+classifier = pipeline("image-classification", model=model_name)
+
 
 # PDF to image conversion
 def pdf_to_images_pymupdf(pdf_data):
@@ -30,6 +28,7 @@ def pdf_to_images_pymupdf(pdf_data):
         print(f"Error converting PDF: {e}")
         return None
 
+
 # File classification function (modified for API)
 def classify_file(file_path):
     try:
@@ -37,15 +36,10 @@ def classify_file(file_path):
         if file_ext in ['.jpg', '.jpeg', '.png', '.gif']:
             # Handle image upload
             image = Image.open(file_path).convert("RGB")
-            inputs = processor(images=image, return_tensors="pt")
-            outputs = model(**inputs)
-            probabilities = F.softmax(outputs.logits, dim=-1)[0].tolist()
-            predicted_class_idx = outputs.logits.argmax(-1).item()
-            result_label = model.config.id2label[predicted_class_idx]
-            confidence = probabilities[predicted_class_idx] * 100
+            result = classifier(image)[0]  # Get the top prediction
             return {
-                "prediction": result_label,
-                "confidence": confidence
+                "prediction": result["label"],
+                "confidence": result["score"] * 100,
             }
         elif file_ext == '.pdf':
             # Handle PDF upload
@@ -54,15 +48,10 @@ def classify_file(file_path):
             images = pdf_to_images_pymupdf(pdf_data)
             if images:
                 image = Image.open(io.BytesIO(images[0])).convert("RGB")
-                inputs = processor(images=image, return_tensors="pt")
-                outputs = model(**inputs)
-                probabilities = F.softmax(outputs.logits, dim=-1)[0].tolist()
-                predicted_class_idx = outputs.logits.argmax(-1).item()
-                result_label = model.config.id2label[predicted_class_idx]
-                confidence = probabilities[predicted_class_idx] * 100
+                result = classifier(image)[0]  # Get the top prediction
                 return {
-                    "prediction": result_label,
-                    "confidence": confidence
+                    "prediction": result["label"],
+                    "confidence": result["score"] * 100,
                 }
             else:
                 return {"error": "PDF conversion failed."}
@@ -71,6 +60,7 @@ def classify_file(file_path):
     except Exception as e:
         return {"error": f"An error occurred: {e}"}
 
+
 # API endpoint for file classification
 @app.route('/classify', methods=['POST'])
 def classify():
@@ -81,11 +71,12 @@ def classify():
         return jsonify({"error": "No file selected"}), 400
 
     filename = secure_filename(file.filename)
-    filepath = os.path.join('/tmp', filename) # Save to a temporary location
+    filepath = os.path.join('/tmp', filename)  # Save to a temporary location
     file.save(filepath)
     result = classify_file(filepath)
-    os.remove(filepath) #remove temp file
+    os.remove(filepath)  # remove temp file
     return jsonify(result), 200  # Return JSON response
 
+
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=5000)