Update app.py

app.py

@@ -1,75 +1,73 @@
-from flask import Flask, request, jsonify 
-from transformers import AutoModelForImageClassification, AutoProcessor
-from PIL import Image
-import io
-import fitz  # PyMuPDF
-from flask_cors import CORS
-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")  # Or "png"
-            images.append(img_data)
-        return images
-    except Exception as e:
-        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()
-
-    try:
-        if file_type in ['jpg', 'jpeg', 'png', 'gif']:
-            # Handle image upload
-            img_data = uploaded_file.read()
-            image = Image.open(io.BytesIO(img_data)).convert("RGB")
-            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 jsonify({'result': result})
-
-        elif file_type == 'pdf':
-            # Handle PDF upload
-            pdf_data = uploaded_file.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()
-                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__':
-    # For Hugging Face Spaces, we need to bind to 0.0.0.0 and use the default port 8080
-    app.run(debug=False, host="0.0.0.0", port=int(os.environ.get("PORT", 8080)))
+import gradio as gr
+from transformers import AutoModelForImageClassification, AutoProcessor
+from PIL import Image
+import io
+import fitz  # PyMuPDF
+
+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")  # Or "png"
+            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()
+
+    try:
+        if file_type in ['jpg', 'jpeg', 'png', 'gif']:
+            # Handle image upload
+            image = Image.open(file).convert("RGB")
+            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 result
+
+        elif file_type == 'pdf':
+            # Handle PDF upload
+            pdf_data = file.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()
+                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}'
+
+# Create the Gradio interface
+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."
+)
+
+# Launch the Gradio app
+demo.launch()