Spaces:
Build error
Build error
File size: 2,779 Bytes
167c1f9 e299395 167c1f9 6732d7a 4df4c38 167c1f9 e299395 e03bfe0 e299395 167c1f9 4df4c38 167c1f9 f06afa9 167c1f9 e299395 4df4c38 167c1f9 4df4c38 e299395 4df4c38 e299395 4df4c38 167c1f9 4df4c38 e299395 4df4c38 e299395 4df4c38 167c1f9 4df4c38 6732d7a e299395 4df4c38 6732d7a 4df4c38 e299395 4df4c38 e299395 4df4c38 167c1f9 e299395 167c1f9 4df4c38 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
from flask import Flask, request, jsonify
from transformers import pipeline
from PIL import Image
import io
import fitz # PyMuPDF
import os
from werkzeug.utils import secure_filename
app = Flask(__name__)
# Load model and processor using pipeline
model_name = "AsmaaElnagger/Diabetic_RetinoPathy_detection"
classifier = pipeline("image-classification", model=model_name)
# PDF to image conversion
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
# File classification function (modified for API)
def classify_file(file_path):
try:
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")
result = classifier(image)[0] # Get the top prediction
return {
"prediction": result["label"],
"confidence": result["score"] * 100,
}
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 images:
image = Image.open(io.BytesIO(images[0])).convert("RGB")
result = classifier(image)[0] # Get the top prediction
return {
"prediction": result["label"],
"confidence": result["score"] * 100,
}
else:
return {"error": "PDF conversion failed."}
else:
return {"error": "Unsupported file type."}
except Exception as e:
return {"error": f"An error occurred: {e}"}
# 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
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=5000)
|