Update app.py

app.py

@@ -1,83 +1,178 @@
-import gradio as gr
 import os
+import io
 from pdf2image import convert_from_path
-from anthropic import Anthropic
+from openai import OpenAI
 import base64
-import io
+import asyncio
+import gradio as gr
 
-# Initialize Anthropic client with API key
-ANTHROPIC_API_KEY = "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"  # Replace with your Anthropic API key
-client = Anthropic(api_key=ANTHROPIC_API_KEY)
+# We'll use an environment variable for the API key in Spaces
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
 
-def extract_text_from_pdf(pdf_path):
-    """Extract all text from PDF pages using Claude"""
-    try:
-        print(f"Processing PDF...")
-        
-        # Convert PDF to images
-        images = convert_from_path(pdf_path)
-        
-        extracted_texts = []
-        for i, image in enumerate(images):
-            print(f"Processing page {i+1} of {len(images)}...")
+class PDFFormExtractor:
+    def __init__(self, api_key):
+        self.client = OpenAI(api_key=api_key)
+    
+    async def extract_first_form(self, pdf_path):
+        try:
+            if not os.path.exists(pdf_path):
+                raise FileNotFoundError(f"PDF file not found: {pdf_path}")
+            
+            print(f"Processing PDF: {pdf_path}")
+            
+            # Only convert the first page
+            images = convert_from_path(pdf_path, first_page=1, last_page=1)
+            
+            if not images:
+                return "No pages found in the PDF."
+            
+            # Process only the first page
+            image = images[0]
+            print("Processing first page...")
             
-            # Convert image to base64
             img_buffer = io.BytesIO()
             image.save(img_buffer, format='PNG')
             img_base64 = base64.b64encode(img_buffer.getvalue()).decode('utf-8')
             
-            # Use Claude to extract text
-            message = client.messages.create(
-                model="claude-3-opus-20240229",
-                max_tokens=4000,
-                temperature=0,
-                system="Extract ALL text from this image exactly as it appears, preserving all formatting, numbers, and special characters. Include everything you can see, from headers to footers, timestamps to footnotes.",
-                messages=[{
+            # Custom prompt specifically for extracting the Patient Admission Form
+            prompt = """
+You are an expert in document processing and OCR with deep knowledge of Markdown formatting. Your task is to extract the structured content from a PDF page image and convert it into a clean Markdown format.
+
+Extract the Patient Admission Form (KMHIPF002V3) from this image and format it as follows:
+
+```markdown
+# PATIENT ADMISSION FORM
+## DR.KAMAKSHI MEMORIAL HOSPITAL, PALLIKARANAI, CHENNAI.
+
+### PATIENT PROFILE
+*Please paste the sticker within the box*
+
+* UHID: ______
+* Patient Name: ______
+* Age/Gender: ______
+* Doctor Name: ______
+
+### BASIC INFORMATION
+* Date & Time of Admission: [DD/MM/YYYY]
+* Date of Birth: [DD/MM/YYYY]
+
+### IDENTIFICATION
+**ID Proof Already Registered**: □ Yes □ No
+
+**Type of ID**:
+* □ Aadhar
+* □ Passport
+* □ Voter ID
+* □ Driving License
+* □ Others
+
+ID No.: ________________
+Contact No.: ________________
+
+### MEDICAL DETAILS
+* Provisional Diagnosis:
+* Reason for Admission:
+* Plan of Care:
+* Expected Outcome*:
+
+### CONSULTANT DETAILS
+* Primary Consultant Name:
+* Speciality:
+
+### PATIENT REFERENCE INFORMATION
+*(To be filled by Front Office)*
+
+**Reference Via**:
+* □ Doctor
+* □ Hospital
+* □ Ambulance
+* □ DRKMH Employee
+* □ Self / Walk In
+
+**Referrer Details**:
+* Name: ________________
+* Contact No.: ________________
+
+### TYPE OF ADMISSION
+* □ Emergency
+* □ Elective
+* □ MLC
+* □ Surgery
+* □ Medical
+* □ Others: ________________
+
+### TREATMENT TYPE
+* □ In Patient
+* □ Day Care
+* Transfer To: ________________
+
+### CONTACT DETAILS
+**Person to Contact (Next of Kin)**:
+* Name: ________________
+* Relationship with Patient: ________________
+* Address: ________________
+* Pincode: ________________
+* Mobile: ________________
+* Email: ________________
+
+### OFFICIAL USE
+* Front Office Executive Name:
+* Front Office Executive Signature:
+* Advance Amount Paid:
+* ICD-10 Code (For Medical Records Section):
+
+*\\* Subject to change during the course of diseases*
+
+---
+*Form No: KMHIPF002V3* """
+
+            response = self.client.chat.completions.create(
+            model="gpt-4o",
+            messages=[
+                {
+                    "role": "system",
+                    "content": prompt
+                },
+                {
                     "role": "user",
                     "content": [
                         {
                             "type": "text",
-                            "text": "Please extract and transcribe ALL text visible in this image, exactly as it appears. Include every piece of text you can see, maintaining the exact formatting, spacing, and line breaks."
+                            "text": "Extract and format the Patient Admission Form from this image according to the specified markdown format."
                         },
                         {
-                            "type": "image",
-                            "source": {
-                                "type": "base64",
-                                "media_type": "image/png",
-                                "data": img_base64
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/png;base64,{img_base64}"
                             }
                         }
                     ]
-                }]
-            )
-            
-            extracted_texts.append(f"\n=== Page {i + 1} ===\n\n{message.content[0].text}")
+                }
+            ],
+            max_tokens=4096
+        )
         
-        return "\n".join(extracted_texts)
-
-    except Exception as e:
-        return f"Error in text extraction: {str(e)}"
-
-def process_pdf(pdf_file):
-    if pdf_file is None:
-        return "Please upload a PDF file."
+        return response.choices[0].message.content
     
-    try:
-        # In Spaces, pdf_file is the file path
-        extracted_text = extract_text_from_pdf(pdf_file)
-        return extracted_text
     except Exception as e:
-        return f"Error processing PDF: {str(e)}"
-
-# Create Gradio interface
-demo = gr.Interface(
-    fn=process_pdf,
-    inputs=gr.File(label="Upload PDF", file_types=[".pdf"]),
-    outputs=gr.Textbox(label="Extracted Text", lines=25),
-    title="PDF Text Extractor (Powered by Claude)",
-    description="Upload a PDF file to extract all text using Claude's Vision capabilities. Please note that processing may take a few minutes depending on the number of pages.",
-    theme=gr.themes.Soft()
-)
+        print(f"Error in form extraction: {str(e)}")
+        return f"Error: {str(e)}"
+
+def extract_first_form(pdf_file):
+    if OPENAI_API_KEY is None:
+        return "Error: OpenAI API key not found. Please set the OPENAI_API_KEY environment variable."
 
-# Launch the app
-demo.launch(share=True)
+        extractor = PDFFormExtractor(OPENAI_API_KEY)
+        
+        pdf_path = pdf_file.name
+        result = asyncio.run(extractor.extract_first_form(pdf_path))
+        return result
+
+iface = gr.Interface(
+fn=extract_first_form,
+inputs=gr.File(label="Upload PDF with Patient Admission Form"),
+outputs=gr.Textbox(label="Extracted Form in Markdown", lines=30),
+title="Patient Admission Form Extractor",
+description="Upload a PDF file to extract the Patient Admission Form (first form) in markdown format."
+)
+iface.launch()