Update app.py

app.py

@@ -1,17 +1,15 @@
 from flask import Flask, request, jsonify
-from flask_cors import CORS
-from transformers import Dinov2ForImageClassification, AutoProcessor
+from transformers import AutoModelForImageClassification, AutoProcessor
 from PIL import Image
 import io
-import fitz
-import torch
+import fitz  # PyMuPDF
+from flask_cors import CORS
 
 app = Flask(__name__)
 CORS(app)
 
-# Load model and processor
 model_name = "AsmaaElnagger/Diabetic_RetinoPathy_detection"
-model = Dinov2ForImageClassification.from_pretrained(model_name)
+model = AutoModelForImageClassification.from_pretrained(model_name)
 processor = AutoProcessor.from_pretrained(model_name)
 
 def pdf_to_images_pymupdf(pdf_data):
@@ -34,39 +32,30 @@ def classify_file():
         return jsonify({'error': 'No file provided'}), 400
 
     uploaded_file = request.files['file']
-    file_type = uploaded_file.filename.rsplit('.', 1)[1].lower()
+    file_type = uploaded_file.filename.rsplit('.', 1)[-1].lower()
 
     try:
-        if file_type in ['jpg', 'jpeg', 'png', 'gif']:
-            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})
-
+        if file_type in ['jpg', 'jpeg', 'png']:
+            image = Image.open(uploaded_file).convert("RGB")
         elif file_type == 'pdf':
             pdf_data = uploaded_file.read()
             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)
-                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
-
+            if not images:
+                return jsonify({'error': 'Failed to convert PDF.'}), 500
+            image = Image.open(io.BytesIO(images[0])).convert("RGB")
         else:
             return jsonify({'error': 'Unsupported file type'}), 400
 
+        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})
+
     except Exception as e:
-        return jsonify({'error': f'An error occurred: {e}'}), 500
+        return jsonify({'error': str(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)