File size: 2,451 Bytes
acb2a86 27e887d acb2a86 27e887d acb2a86 27e887d acb2a86 27e887d acb2a86 27e887d acb2a86 27e887d acb2a86 27e887d acb2a86 27e887d acb2a86 27e887d acb2a86 27e887d acb2a86 27e887d acb2a86 27e887d acb2a86 27e887d acb2a86 27e887d acb2a86 27e887d acb2a86 27e887d acb2a86 27e887d acb2a86 |
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 |
import gradio as gr
from transformers import AutoModelForImageClassification, AutoProcessor
from PIL import Image
import io
import fitz # PyMuPDF
import os
# 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")
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
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")
inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)
predicted_class_idx = outputs.logits.argmax(-1).item()
result = model.config.id2label[predicted_class_idx]
return result
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")
inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)
predicted_class_idx = outputs.logits.argmax(-1).item()
result = model.config.id2label[predicted_class_idx]
return result
else:
return "PDF conversion failed."
else:
return "Unsupported file type."
except Exception as e:
return f"An error occurred: {e}"
# Gradio UI
demo = gr.Interface(
fn=classify_file,
inputs=gr.File(label="Upload PDF or Image"),
outputs="text",
title="Diabetic Retinopathy Detection",
description="Upload a fundus scan (image or PDF) to detect diabetic retinopathy."
)
# Launch app
demo.launch()
|