from flask import Flask, request, jsonify
from flask_cors import CORS
from transformers import AutoModelForImageClassification, AutoProcessor
from PIL import Image
import io
import fitz
import torch
import os

app = Flask(__name__)
CORS(app)

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:
        pdf_document = fitz.open(stream=pdf_data, filetype="pdf")
        images = []
        for page_num in range(pdf_document.page_count):
            page = pdf_document.load_page(page_num)
            pix = page.get_pixmap()
            img_data = pix.tobytes("jpeg")
            images.append(img_data)
        return images
    except Exception as e:
        print(f"Error converting PDF: {e}")
        return None

@app.route('/classify', methods=['POST', 'OPTIONS'])
def classify_file():
    if request.method == 'OPTIONS':
        response = jsonify()
        response.headers.add('Access-Control-Allow-Origin', '*')
        response.headers.add('Access-Control-Allow-Methods', 'POST, OPTIONS')
        response.headers.add('Access-Control-Allow-Headers', 'Content-Type')
        return response

    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()

    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")
            with torch.no_grad():
                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':
            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")
                with torch.no_grad():
                    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
        else:
            return jsonify({'error': 'Unsupported file type'}), 400

    except Exception as e:
        return jsonify({'error': f'An error occurred: {e}'}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)