Update app.py

app.py

@@ -1,50 +1,59 @@
 import streamlit as st
-from transformers import pipeline
-from PIL import Image
-import pytesseract
-import PyPDF2
-import pdfplumber
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
 import torch
+import easyocr
+from PIL import Image
+import fitz  # PyMuPDF
 
-# Load the BART model for summarization and NLI
+# Load EasyOCR Reader
+reader = easyocr.Reader(['en'])
+
+# Load ClinicalBERT Model
 @st.cache_resource
-def load_model():
-    return pipeline("zero-shot-classification", model="facebook/bart-large-mnli", device=0 if torch.cuda.is_available() else -1)
+def load_clinicalbert():
+    model_name = "emilyalsentzer/Bio_ClinicalBERT"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name)
+    return pipeline("text-classification", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
 
-classifier = load_model()
+clinicalbert_analyzer = load_clinicalbert()
 
-# OCR for Image using Tesseract
+# OCR function using EasyOCR
 def extract_text_from_image(image):
-    return pytesseract.image_to_string(image)
+    image = Image.open(image).convert('RGB')
+    result = reader.readtext(image, detail=0)
+    extracted_text = " ".join(result)
+    return extracted_text
 
-# Extract text from PDF using pdfplumber
+# PDF text extraction using PyMuPDF
 def extract_text_from_pdf(pdf_file):
     text = ""
-    with pdfplumber.open(pdf_file) as pdf:
-        for page in pdf.pages:
-            text += page.extract_text()
+    with fitz.open(pdf_file) as pdf:
+        for page in pdf:
+            text += page.get_text()
     return text
 
-# Summarize, interpret and give actionable insights
+# Analysis function using ClinicalBERT
 def analyze_report(text):
-    # Provide a summary
-    summary = classifier(text, candidate_labels=["summary"], multi_label=False)['labels'][0]
-
-    # Interpretation of results
-    interpretation = classifier(text, candidate_labels=["interpretation", "normal", "abnormal"], multi_label=True)
+    # Define prompts for analysis
+    summary_prompt = f"Summarize the following medical report:\n{text}"
+    interpretation_prompt = f"Interpret the following lab results:\n{text}"
+    recommendation_prompt = f"Provide actionable recommendations based on this medical report:\n{text}"
 
-    # Recommendations
-    recommendations = classifier(text, candidate_labels=["follow-up", "Holistic/OTC treatment", "dietary change", "medication"], multi_label=True)
+    # Use ClinicalBERT for text analysis
+    summary = clinicalbert_analyzer(summary_prompt)[0]['label']
+    interpretation = clinicalbert_analyzer(interpretation_prompt)[0]['label']
+    recommendations = clinicalbert_analyzer(recommendation_prompt)[0]['label']
 
     return {
         "summary": summary,
-        "interpretation": interpretation['labels'],
-        "recommendations": recommendations['labels']
+        "interpretation": interpretation,
+        "recommendations": recommendations
     }
 
 # Streamlit UI
-st.title("Medical Lab Report Analyzer")
-st.write("Upload your medical lab report (PDF/Image) for insights.")
+st.title("Clinical Lab Report Analyzer")
+st.write("Upload your medical lab report (PDF/Image) to get a summary, interpretation, and actionable recommendations.")
 
 uploaded_file = st.file_uploader("Choose a PDF/Image file", type=["pdf", "png", "jpg", "jpeg"])
 
@@ -57,12 +66,11 @@ if uploaded_file:
             extracted_text = extract_text_from_pdf(uploaded_file)
     else:
         with st.spinner("Extracting text from Image..."):
-            image = Image.open(uploaded_file)
-            extracted_text = extract_text_from_image(image)
+            extracted_text = extract_text_from_image(uploaded_file)
 
     # Analyze the extracted text
     if extracted_text:
-        with st.spinner("Analyzing report..."):
+        with st.spinner("Analyzing the medical report using ClinicalBERT..."):
             result = analyze_report(extracted_text)
         
         # Display the results
@@ -70,11 +78,9 @@ if uploaded_file:
         st.write(result['summary'])
 
         st.subheader("Interpretation of Results")
-        for label in result['interpretation']:
-            st.write(f"- {label.capitalize()}")
+        st.write(result['interpretation'])
 
         st.subheader("Actionable Recommendations")
-        for rec in result['recommendations']:
-            st.write(f"- {rec.capitalize()}")
+        st.write(result['recommendations'])
     else:
         st.error("No text could be extracted. Please try with a different file.")