Update app.py

app.py

@@ -1,18 +1,17 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
 from PIL import Image
 from PyPDF2 import PdfReader
 from reportlab.lib.pagesizes import letter
 from reportlab.pdfgen import canvas
 
-# Load Hugging Face OCR model for text extraction
+# Load OCR model for extracting text from images
 ocr_model = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
 
-# Load Bio_ClinicalBERT model for medical analysis
-model_name = "emilyalsentzer/Bio_ClinicalBERT"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
+# Load medical AI model (BioGPT or similar) for prescription validation
+medical_model_name = "microsoft/BioGPT"
+medical_tokenizer = AutoTokenizer.from_pretrained(medical_model_name)
+medical_model = AutoModelForSequenceClassification.from_pretrained(medical_model_name)
 
 # Function to extract text from images or PDFs
 def extract_text(file_path):
@@ -28,62 +27,69 @@ def extract_text(file_path):
     except Exception as e:
         return f"Error processing the file: {e}"
 
-# Function to generate a PDF report
-def create_pdf_report(analysis, output_path):
-    c = canvas.Canvas(output_path, pagesize=letter)
-    c.drawString(100, 750, "Blood Test Report Analysis")
-    c.drawString(100, 730, "---------------------------")
-    y_position = 700
-    for line in analysis.split("\n"):
-        c.drawString(100, y_position, line)
-        y_position -= 20
-    c.save()
-    return output_path
+# Function to validate prescription using the medical model
+def validate_prescription_with_model(extracted_text):
+    # Tokenize and process with the AI model
+    inputs = medical_tokenizer(extracted_text, return_tensors="pt", truncation=True, padding=True)
+    outputs = medical_model(**inputs)
+    logits = outputs.logits
+    predictions = logits.softmax(dim=-1)
 
-# Function to analyze blood test reports using Bio_ClinicalBERT
-def analyze_blood_test(file):
+    # Generate model-driven validation results
+    validation_report = "🔍 Prescription Validation Results:\n"
+    for i, score in enumerate(predictions[0]):
+        token = medical_tokenizer.decode([i]).strip()
+        if token not in ["[PAD]", "[unused1]"]:  # Ignore invalid tokens
+            validation_report += f"- {token}: {score.item():.2f}\n"
+    return validation_report
+
+# Main function to handle prescription analysis
+def analyze_prescription(file):
     try:
-        # Step 1: Extract text from uploaded file
+        # Step 1: Extract text
         extracted_text = extract_text(file)
         if not extracted_text.strip():
             return "No readable text found in the uploaded file.", None
 
-        # Step 2: Tokenize and pass through Bio_ClinicalBERT
-        inputs = tokenizer(extracted_text, return_tensors="pt", truncation=True, padding=True)
-        outputs = model(**inputs)
-
-        # Step 3: Process logits and generate meaningful labels
-        logits = outputs.logits
-        predictions = logits.softmax(dim=-1)
+        # Step 2: Validate prescription using AI model
+        validation_report = validate_prescription_with_model(extracted_text)
 
-        # Generate analysis report
-        analysis_report = "🔍 Analysis Results:\n"
-        for i, score in enumerate(predictions[0]):
-            token = tokenizer.decode([i]).strip()
-            if token not in ["[PAD]", "[unused1]"]:  # Filter out invalid tokens
-                analysis_report += f"- {token}: {score.item():.2f}\n"
+        # Combine the extracted text and validation results
+        full_report = f"Extracted Text:\n{extracted_text}\n\n{validation_report}"
 
-        # Step 4: Generate a PDF report
-        output_pdf = "analysis_report.pdf"
-        create_pdf_report(f"Extracted Text:\n{extracted_text}\n\n{analysis_report}", output_pdf)
+        # Step 3: Generate a PDF report
+        output_pdf = "prescription_validation_report.pdf"
+        create_pdf_report(full_report, output_pdf)
 
-        return analysis_report, output_pdf
+        return full_report, output_pdf
     except Exception as e:
         return f"Error processing file: {e}", None
 
-# Gradio interface setup
+# Function to create a PDF report
+def create_pdf_report(content, output_path):
+    c = canvas.Canvas(output_path, pagesize=letter)
+    c.drawString(100, 750, "Prescription Validation Report")
+    c.drawString(100, 730, "------------------------------")
+    y_position = 700
+    for line in content.split("\n"):
+        c.drawString(100, y_position, line)
+        y_position -= 20
+    c.save()
+    return output_path
+
+# Gradio interface
 interface = gr.Interface(
-    fn=analyze_blood_test,
-    inputs=gr.File(label="Upload your Blood Test Report (PNG, JPG, JPEG, or PDF)"),
+    fn=analyze_prescription,
+    inputs=gr.File(label="Upload your Prescription (PNG, JPG, JPEG, or PDF)"),
     outputs=[
-        gr.Textbox(label="Analysis Results"),
+        gr.Textbox(label="Validation Results"),
         gr.File(label="Download PDF Report")
     ],
-    title="MedAI Analyzer",
+    title="AI-Powered Prescription Validator",
     description=(
-        "Upload your blood test report in image (PNG, JPG, JPEG) or PDF format. "
-        "The app will extract and analyze the values, flag abnormalities, and provide health recommendations. "
-        "You can also download a detailed PDF report of the analysis."
+        "Upload your medical prescription in image (PNG, JPG, JPEG) or PDF format. "
+        "The app will extract the text, analyze it using advanced AI models, and validate the prescription. "
+        "Download a comprehensive PDF report of the validation results."
     ),
     allow_flagging="never"
 )