Update app.py

app.py

@@ -1,17 +1,15 @@
 from flask import Flask, request, jsonify
-from transformers import AutoModelForImageClassification, AutoProcessor
+from transformers import pipeline
 from PIL import Image
 import io
 import fitz  # PyMuPDF
 from flask_cors import CORS
 
 app = Flask(__name__)
-CORS(app)  # Allow cross-origin requests
+CORS(app)
 
-# Load model and processor
-model_name = "AsmaaElnagger/Diabetic_RetinoPathy_detection"
-model = AutoModelForImageClassification.from_pretrained(model_name)
-processor = AutoProcessor.from_pretrained(model_name)
+# Using the pipeline to automatically load the model and processor
+pipe = pipeline("image-classification", model="AsmaaElnagger/Diabetic_RetinoPathy_detection")
 
 def pdf_to_images_pymupdf(pdf_data):
     try:
@@ -40,12 +38,8 @@ def classify_file():
             # Handle image upload
             img_data = uploaded_file.read()
             image = Image.open(io.BytesIO(img_data)).convert("RGB")
-            inputs = processor(images=image, return_tensors="pt")
-            outputs = model(**inputs)
-            logits = outputs.logits
-            predicted_class_idx = logits.argmax(-1).item()
-            result = model.config.id2label[predicted_class_idx]
-            return jsonify({'result': result})
+            result = pipe(image)  # Use pipeline for classification
+            return jsonify({'result': result[0]['label']})
 
         elif file_type == 'pdf':
             # Handle PDF upload
@@ -53,14 +47,10 @@ def classify_file():
             images = pdf_to_images_pymupdf(pdf_data)
 
             if images:
-                # Process the first image in the pdf, you may need to loop through all images.
+                # Process the first image in the pdf
                 image = Image.open(io.BytesIO(images[0])).convert("RGB")
-                inputs = processor(images=image, return_tensors="pt")
-                outputs = model(**inputs)
-                logits = outputs.logits
-                predicted_class_idx = logits.argmax(-1).item()
-                result = model.config.id2label[predicted_class_idx]
-                return jsonify({'result': result})
+                result = pipe(image)  # Use pipeline for classification
+                return jsonify({'result': result[0]['label']})
             else:
                 return jsonify({'error': 'PDF conversion failed.'}), 500
 
@@ -71,4 +61,4 @@ def classify_file():
         return jsonify({'error': f'An error occurred: {e}'}), 500
 
 if __name__ == '__main__':
-    app.run(host="0.0.0.0", port=7860, debug=True)  # Make it accessible from any network interface
+    app.run(host="0.0.0.0", port=7860, debug=True)