Update app.py

app.py

@@ -3,15 +3,19 @@ from transformers import AutoModelForImageClassification, AutoProcessor
 from PIL import Image
 import io
 import fitz  # PyMuPDF
-from flask_cors import CORS
+import os
+import torch
+import torch.nn.functional as F
+from werkzeug.utils import secure_filename
 
 app = Flask(__name__)
-CORS(app)
 
+# Load model and processor
 model_name = "AsmaaElnagger/Diabetic_RetinoPathy_detection"
 model = AutoModelForImageClassification.from_pretrained(model_name)
 processor = AutoProcessor.from_pretrained(model_name)
 
+# PDF to image conversion
 def pdf_to_images_pymupdf(pdf_data):
     try:
         pdf_document = fitz.open(stream=pdf_data, filetype="pdf")
@@ -26,36 +30,62 @@ def pdf_to_images_pymupdf(pdf_data):
         print(f"Error converting PDF: {e}")
         return None
 
-@app.route('/classify', methods=['POST'])
-def classify_file():
-    if 'file' not in request.files:
-        return jsonify({'error': 'No file provided'}), 400
-
-    uploaded_file = request.files['file']
-    file_type = uploaded_file.filename.rsplit('.', 1)[-1].lower()
-
+# File classification function (modified for API)
+def classify_file(file_path):
     try:
-        if file_type in ['jpg', 'jpeg', 'png']:
-            image = Image.open(uploaded_file).convert("RGB")
-        elif file_type == 'pdf':
-            pdf_data = uploaded_file.read()
+        file_ext = os.path.splitext(file_path)[-1].lower()
+        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
+            return {
+                "prediction": result_label,
+                "confidence": confidence
+            }
+        elif file_ext == '.pdf':
+            # Handle PDF upload
+            with open(file_path, "rb") as f:
+                pdf_data = f.read()
             images = pdf_to_images_pymupdf(pdf_data)
-            if not images:
-                return jsonify({'error': 'Failed to convert PDF.'}), 500
-            image = Image.open(io.BytesIO(images[0])).convert("RGB")
+            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
+                return {
+                    "prediction": result_label,
+                    "confidence": confidence
+                }
+            else:
+                return {"error": "PDF conversion failed."}
         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 {"error": "Unsupported file type."}
+    except Exception as e:
+        return {"error": f"An error occurred: {e}"}
 
-        return jsonify({'result': result})
+# API endpoint for file classification
+@app.route('/classify', methods=['POST'])
+def classify():
+    if 'file' not in request.files:
+        return jsonify({"error": "No file part"}), 400
+    file = request.files['file']
+    if file.filename == '':
+        return jsonify({"error": "No file selected"}), 400
 
-    except Exception as e:
-        return jsonify({'error': str(e)}), 500
+    filename = secure_filename(file.filename)
+    filepath = os.path.join('/tmp', filename) # Save to a temporary location
+    file.save(filepath)
+    result = classify_file(filepath)
+    os.remove(filepath) #remove temp file
+    return jsonify(result), 200  # Return JSON response
 
 if __name__ == '__main__':
-    app.run(host="0.0.0.0", port=7860)
+    app.run(host='0.0.0.0', port=5000)