Update app.py

app.py

@@ -3,11 +3,14 @@ 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")
@@ -15,59 +18,57 @@ def pdf_to_images_pymupdf(pdf_data):
         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")  # Or "png"
+            img_data = pix.tobytes("jpeg")
             images.append(img_data)
         return images
     except Exception as e:
         print(f"Error converting PDF: {e}")
         return None
 
-def classify_file(file):
-    # Get the file type
-    file_type = file.name.split('.')[-1].lower()
-
+# File classification
+def classify_file(file_path):
     try:
-        if file_type in ['jpg', 'jpeg', 'png', 'gif']:
+        file_ext = os.path.splitext(file_path)[-1].lower()
+
+        if file_ext in ['.jpg', '.jpeg', '.png', '.gif']:
             # Handle image upload
-            image = Image.open(file).convert("RGB")
+            image = Image.open(file_path).convert("RGB")
             inputs = processor(images=image, return_tensors="pt")
             outputs = model(**inputs)
-            logits = outputs.logits
-            predicted_class_idx = logits.argmax(-1).item()
+            predicted_class_idx = outputs.logits.argmax(-1).item()
             result = model.config.id2label[predicted_class_idx]
             return result
 
-        elif file_type == 'pdf':
+        elif file_ext == '.pdf':
             # Handle PDF upload
-            pdf_data = file.read()
+            with open(file_path, "rb") as f:
+                pdf_data = f.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)
-                logits = outputs.logits
-                predicted_class_idx = logits.argmax(-1).item()
+                predicted_class_idx = outputs.logits.argmax(-1).item()
                 result = model.config.id2label[predicted_class_idx]
                 return result
             else:
-                return 'PDF conversion failed.'
+                return "PDF conversion failed."
 
         else:
-            return 'Unsupported file type'
+            return "Unsupported file type."
 
     except Exception as e:
-        return f'An error occurred: {e}'
+        return f"An error occurred: {e}"
 
-# Create the Gradio interface
+# Gradio UI
 demo = gr.Interface(
-    fn=classify_file, 
-    inputs=gr.File(label="Upload PDF or Image"), 
-    outputs="text", 
-    title="Diabetic Retinopathy Detection", 
-    description="Upload an image or PDF for classification."
+    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 the Gradio app
+# Launch app
 demo.launch()