Update app.py

app.py

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