Update app.py

app.py

@@ -17,6 +17,7 @@ from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, Tabl
 from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
 import io
 from dotenv import load_dotenv
+import fitz  # PyMuPDF for PDF processing
 
 # Load environment variables
 load_dotenv()
@@ -30,12 +31,12 @@ logger = logging.getLogger(__name__)
 
 # Configuration and Constants
 class Config:
-    GEMINI_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent"
-    GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+    GEMINI_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash-exp:generateContent"
+    GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")  # Load from .env
     MAX_RETRIES = 3
     TIMEOUT = 30
     MAX_IMAGE_SIZE = (1600, 1600)
-    ALLOWED_MIME_TYPES = ["image/jpeg", "image/png"]
+    ALLOWED_MIME_TYPES = ["image/jpeg", "image/png", "application/pdf"]
     MAX_FILE_SIZE = 5 * 1024 * 1024  # 5MB
 
 # Custom Exceptions
@@ -45,6 +46,9 @@ class APIError(Exception):
 class ImageProcessingError(Exception):
     pass
 
+class PDFProcessingError(Exception):
+    pass
+
 # Initialize session state
 def init_session_state():
     if 'processing_history' not in st.session_state:
@@ -179,22 +183,20 @@ class PDFGenerator:
         doc.build(elements)
         return buffer.getvalue()
 
-
 class ImageProcessor:
     @staticmethod
-    def validate_image(uploaded_file) -> tuple[bool, str]:
+    def validate_file(uploaded_file) -> tuple[bool, str]:
         try:
             if uploaded_file.size > Config.MAX_FILE_SIZE:
                 return False, f"File size exceeds {Config.MAX_FILE_SIZE // (1024*1024)}MB limit"
             
-            image = Image.open(uploaded_file)
-            if image.format.upper() not in ['JPEG', 'PNG']:
-                return False, "Unsupported image format. Please upload JPEG or PNG"
+            if uploaded_file.type not in Config.ALLOWED_MIME_TYPES:
+                return False, "Unsupported file type. Please upload JPEG, PNG, or PDF."
             
-            return True, "Image validation successful"
+            return True, "File validation successful"
         except Exception as e:
-            logger.error(f"Image validation error: {str(e)}")
-            return False, f"Image validation failed: {str(e)}"
+            logger.error(f"File validation error: {str(e)}")
+            return False, f"File validation failed: {str(e)}"
 
     @staticmethod
     def preprocess_image(image: Image.Image) -> Image.Image:
@@ -214,14 +216,23 @@ class DocumentProcessor:
     def __init__(self):
         self.image_processor = ImageProcessor()
     
-    def process_document(self, image: Image.Image) -> Dict[str, Any]:
+    def process_document(self, uploaded_file) -> Dict[str, Any]:
         try:
-            processed_image = self.image_processor.preprocess_image(image)
-            image_base64 = self.encode_image(processed_image)
+            if uploaded_file.type.startswith("image/"):
+                # Process image
+                image = Image.open(uploaded_file)
+                processed_image = self.image_processor.preprocess_image(image)
+                image_base64 = self.encode_image(processed_image)
+                extracted_text = self.extract_text(image_base64)
+            elif uploaded_file.type == "application/pdf":
+                # Process PDF
+                extracted_text = self.extract_text_from_pdf(uploaded_file)
+            else:
+                raise ValueError("Unsupported file type.")
             
             results = {
-                "document_type": self.classify_document(image_base64),
-                "extracted_text": self.extract_text(image_base64),
+                "document_type": self.classify_document(extracted_text),
+                "extracted_text": extracted_text,
                 "structured_data": None
             }
             
@@ -241,8 +252,22 @@ class DocumentProcessor:
         image.save(buffered, format="JPEG", quality=95)
         return base64.b64encode(buffered.getvalue()).decode('utf-8')
 
-    def classify_document(self, image_base64: str) -> str:
-        prompt = """
+    @staticmethod
+    def extract_text_from_pdf(uploaded_file) -> str:
+        try:
+            pdf_bytes = uploaded_file.read()
+            pdf_document = fitz.open(stream=pdf_bytes, filetype="pdf")
+            text = ""
+            for page_num in range(len(pdf_document)):
+                page = pdf_document.load_page(page_num)
+                text += page.get_text()
+            return text
+        except Exception as e:
+            logger.error(f"PDF processing error: {str(e)}")
+            raise PDFProcessingError(f"Failed to process PDF: {str(e)}")
+
+    def classify_document(self, text: str) -> str:
+        prompt = f"""
         Analyze this medical document and classify it into one of the following categories:
         - Lab Report
         - Patient Chart
@@ -251,8 +276,11 @@ class DocumentProcessor:
         - Medical Certificate
         - Other (specify)
         Provide only the category name.
+
+        Document Text:
+        {text}
         """
-        response = GeminiAPI.call_api(prompt, image_base64)
+        response = GeminiAPI.call_api(prompt)
         return response["candidates"][0]["content"]["parts"][0]["text"].strip()
 
     def extract_text(self, image_base64: str) -> str:
@@ -318,6 +346,9 @@ class DocumentProcessor:
             self.correct_medicine_name(med) for med in structured_data.get('medications', [])
         ]
         
+        # Improve symptoms extraction
+        structured_data['symptoms'] = self.extract_symptoms(text)
+        
         return structured_data
 
     @staticmethod
@@ -341,6 +372,17 @@ class DocumentProcessor:
         medication['name'] = response["candidates"][0]["content"]["parts"][0]["text"].strip()
         return medication
 
+    @staticmethod
+    def extract_symptoms(text: str) -> list[str]:
+        """Extract symptoms from the text."""
+        prompt = f"""
+        Extract all symptoms mentioned in the following medical text. Return only a list of symptoms:
+        {text}
+        """
+        response = GeminiAPI.call_api(prompt)
+        symptoms = response["candidates"][0]["content"]["parts"][0]["text"].strip().split("\n")
+        return [symptom.strip() for symptom in symptoms if symptom.strip()]
+
     @staticmethod
     def parse_json_response(response: Dict[str, Any]) -> Dict[str, Any]:
         try:
@@ -420,7 +462,7 @@ def main():
     setup_page()
     
     st.title("🏥 Advanced Medical Document Processor")
-    st.markdown("Upload medical documents for automated processing and analysis.")
+    st.markdown("Upload medical documents (images or PDFs) for automated processing and analysis.")
     
     # Sidebar
     with st.sidebar:
@@ -441,30 +483,32 @@ def main():
     # Main content
     uploaded_file = st.file_uploader(
         "Choose a medical document",
-        type=['png', 'jpg', 'jpeg'],
-        help="Upload a clear image of a medical document (max 5MB)"
+        type=['png', 'jpg', 'jpeg', 'pdf'],
+        help="Upload a clear image or PDF of a medical document (max 5MB)"
     )
     
     if uploaded_file:
         try:
-            # Validate image
-            is_valid, message = ImageProcessor.validate_image(uploaded_file)
+            # Validate file
+            is_valid, message = ImageProcessor.validate_file(uploaded_file)
             if not is_valid:
                 st.error(message)
                 return
             
-            # Display image
-            image = Image.open(uploaded_file)
-            col1, col2 = st.columns([1, 2])
-            
-            with col1:
-                st.image(image, caption="Uploaded Document", use_column_width=True)
+            # Display file
+            if uploaded_file.type.startswith("image/"):
+                image = Image.open(uploaded_file)
+                col1, col2 = st.columns([1, 2])
+                with col1:
+                    st.image(image, caption="Uploaded Document", use_column_width=True)
+            elif uploaded_file.type == "application/pdf":
+                st.info("PDF file uploaded. Processing...")
             
             # Process document
             if st.button("🔍 Process Document"):
                 with st.spinner("Processing document..."):
                     processor = DocumentProcessor()
-                    results = processor.process_document(image)
+                    results = processor.process_document(uploaded_file)
                     
                     # Generate PDF
                     pdf_bytes = PDFGenerator.create_pdf(results['structured_data'])
@@ -486,7 +530,7 @@ def main():
                     })
                 
                 # Display results
-                with col2:
+                with col2 if uploaded_file.type.startswith("image/") else st:
                     st.success("Document processed successfully!")
                     st.markdown(f"**Document Type:** {results['document_type']}")