Update tabs/report_analyzer.py

tabs/report_analyzer.py

@@ -2,16 +2,17 @@ import gradio as gr
 import torch
 import fitz
 import pytesseract
-import pickle
 import re
+import os
+import google.generativeai as genai
 from PIL import Image, ImageEnhance, ImageFilter
 from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Load FLAN-T5 model for summarization
-gpt_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
-gpt_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
+# Configure Gemini (PaLM) API
+genai.configure(api_key=os.getenv("PALM_API_KEY"))
+model = genai.GenerativeModel("gemini-pro")
 
-# Load English to Hindi translation model
+# Translation model (e.g., for Hindi)
 translation_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-hi")
 translation_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-hi")
 translator = pipeline("translation", model=translation_model, tokenizer=translation_tokenizer)
@@ -96,6 +97,22 @@ def preprocess_image(image_path):
     image = ImageEnhance.Contrast(image).enhance(2)
     return image
 
+def summarize_with_gemini(cleaned_lines):
+    prompt = f"""
+You are a medical assistant. Summarize this lab report in clear, simple language:
+1. Summary in 2–3 lines
+2. Explain abnormal values
+3. List health concerns (if any) in bullet points
+
+Data:
+{chr(10).join(cleaned_lines[:6])}
+"""
+    try:
+        response = model.generate_content(prompt)
+        return response.text.strip() if response and response.text else "(No summary returned)"
+    except Exception as e:
+        return f"(Gemini summarization failed: {e})"
+
 def ocr_and_explain(file, language):
     if not file:
         return "Please upload a valid report.", ""
@@ -132,25 +149,19 @@ def ocr_and_explain(file, language):
                         f"<i>Reference Range: {values['low']}-{values['high']} {values['unit']}</i><br><br>"
                     )
                     rule_lines.append(html_line)
-
                     cleaned_lines.append(f"{term}: {value:.2f} {values['unit']} → {status} (Normal: {values['low']}-{values['high']} {values['unit']})")
                 except:
                     continue
 
     rule_explanation = "\n".join(rule_lines) if rule_lines else "No known lab terms detected."
-    simplified_prompt = "You are a medical assistant. Summarize the following lab test results:\n\n" + "\n".join(cleaned_lines[:6])
 
-    try:
-        inputs = gpt_tokenizer(simplified_prompt, return_tensors="pt", truncation=True)
-        outputs = gpt_model.generate(**inputs, max_new_tokens=300)
-        gpt_summary = gpt_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    except Exception as e:
-        gpt_summary = f"(FLAN-T5 summarization failed: {e})"
+    # 🔁 Gemini summary
+    gpt_summary = summarize_with_gemini(cleaned_lines)
 
     final_output = (
         "<h4 style='color:#ffa500;'>📌 Rule-Based Results:</h4><br>" +
         rule_explanation +
-        "<hr><h4 style='color:#77dd77;'>🧠 Simplified Summary:</h4><br>" +
+        "<hr><h4 style='color:#77dd77;'>🧠 Gemini Summary:</h4><br>" +
         gpt_summary
     )
 
@@ -183,7 +194,8 @@ def report_analyzer_tab():
 
         with gr.Column():
             processing_status = gr.HTML()
-            output_box = gr.HTML("""<div style="background:#1e1e1e; padding:15px; border-radius:10px;"> <h4 style="color:#ffffff;">📋 Final Explanation Output</h4>""")
+            output_box = gr.HTML("""<div style="background:#1e1e1e; padding:15px; border-radius:10px;">
+                                    <h4 style="color:#ffffff;">📋 Final Explanation Output</h4>""")
             output_explanation = gr.HTML()
             output_close = gr.HTML("</div>")