Update app.py

app.py

@@ -1,6 +1,6 @@
 from flask import Flask, request, jsonify
 from flask_cors import CORS
-from transformers import Dinov2ForImageClassification, AutoProcessor # Import the correct model class.
+from transformers import Dinov2ForImageClassification, AutoProcessor
 from PIL import Image
 import io
 import fitz
@@ -11,8 +11,7 @@ CORS(app)
 
 # Load model and processor
 model_name = "AsmaaElnagger/Diabetic_RetinoPathy_detection"
-# model = AutoModelForImageClassification.from_pretrained(model_name) # REMOVE THIS LINE
-model = Dinov2ForImageClassification.from_pretrained(model_name)  # USE THE CORRECT MODEL CLASS
+model = Dinov2ForImageClassification.from_pretrained(model_name)
 processor = AutoProcessor.from_pretrained(model_name)
 
 def pdf_to_images_pymupdf(pdf_data):
@@ -39,7 +38,6 @@ def classify_file():
 
     try:
         if file_type in ['jpg', 'jpeg', 'png', 'gif']:
-            # Handle image upload
             img_data = uploaded_file.read()
             image = Image.open(io.BytesIO(img_data)).convert("RGB")
             inputs = processor(images=image, return_tensors="pt")
@@ -50,12 +48,10 @@ def classify_file():
             return jsonify({'result': result})
 
         elif file_type == 'pdf':
-            # Handle PDF upload
             pdf_data = uploaded_file.read()
             images = pdf_to_images_pymupdf(pdf_data)
 
             if images:
-                # Process the first image in the pdf, you may need to loop through all images.
                 image = Image.open(io.BytesIO(images[0])).convert("RGB")
                 inputs = processor(images=image, return_tensors="pt")
                 outputs = model(**inputs)
@@ -73,4 +69,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)