Update app.py


app.py

@@ -1,14 +1,14 @@
 # Install necessary libraries
-# !pip install gradio transformers pandas PyPDF2 pdfplumber torch torchvision timm sentencepiece
+#!pip install gradio transformers pandas PyPDF2 pdfplumber torch torchvision timm sentencepiece
 
 import gradio as gr
 from transformers import pipeline
 import pandas as pd
+import PyPDF2
 import pdfplumber
 import torch
-from torchvision import transforms
-from PIL import Image
 import timm
+from PIL import Image
 
 # Load pre-trained model for zero-shot classification
 classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
@@ -20,11 +20,44 @@ image_model.eval()
 # Initialize patient database
 patients_db = []
 
+# Disease and medication mapping
+disease_details = {
+    "anemia": {
+        "medication": "Iron supplements (e.g., Ferrous sulfate)",
+        "precaution": "Increase intake of iron-rich foods like spinach and red meat."
+    },
+    "viral infection": {
+        "medication": "Antiviral drugs (e.g., Oseltamivir for flu)",
+        "precaution": "Rest, stay hydrated, and avoid close contact with others."
+    },
+    "liver disease": {
+        "medication": "Hepatoprotective drugs (e.g., Ursodeoxycholic acid)",
+        "precaution": "Avoid alcohol and maintain a balanced diet."
+    },
+    "kidney disease": {
+        "medication": "Angiotensin-converting enzyme inhibitors (e.g., Lisinopril)",
+        "precaution": "Monitor salt intake and stay hydrated."
+    },
+    "diabetes": {
+        "medication": "Metformin or insulin therapy",
+        "precaution": "Follow a low-sugar diet and exercise regularly."
+    },
+    "hypertension": {
+        "medication": "Antihypertensive drugs (e.g., Amlodipine)",
+        "precaution": "Reduce salt intake and manage stress."
+    },
+    "COVID-19": {
+        "medication": "Supportive care, antiviral drugs (e.g., Remdesivir in severe cases)",
+        "precaution": "Follow isolation protocols, wear a mask, and stay hydrated."
+    },
+    "pneumonia": {
+        "medication": "Antibiotics (e.g., Amoxicillin) if bacterial",
+        "precaution": "Rest, avoid smoking, and stay hydrated."
+    }
+}
+
 # Function to register patients
 def register_patient(name, age, gender):
-    if not name or age <= 0 or gender not in ["Male", "Female", "Other"]:
-        return "❌ Invalid input. Please check the details and try again."
-
     patient_id = len(patients_db) + 1
     patients_db.append({
         "ID": patient_id, 
@@ -35,74 +68,68 @@ def register_patient(name, age, gender):
         "Diagnosis": "", 
         "Action Plan": "", 
         "Medications": "", 
+        "Precautions": "",
         "Tests": ""
     })
     return f"✅ Patient {name} registered successfully. Patient ID: {patient_id}"
 
 # Function to analyze text reports
 def analyze_report(patient_id, report_text):
-    if not report_text:
-        return "❌ Report text cannot be empty."
-
-    candidate_labels = ["anemia", "viral infection", "liver disease", "kidney disease", "diabetes"]
+    candidate_labels = list(disease_details.keys())
     result = classifier(report_text, candidate_labels)
     diagnosis = result['labels'][0]
-    action_plan = f"Based on your report, you might have {diagnosis}. Please consult a doctor for confirmation."
+
+    # Fetch medication and precaution
+    medication = disease_details[diagnosis]["medication"]
+    precaution = disease_details[diagnosis]["precaution"]
+    action_plan = f"You might have {diagnosis}. Please consult a doctor for confirmation."
 
     # Store diagnosis in the database
     for patient in patients_db:
         if patient["ID"] == patient_id:
             patient["Diagnosis"] = diagnosis
             patient["Action Plan"] = action_plan
+            patient["Medications"] = medication
+            patient["Precautions"] = precaution
             break
-    else:
-        return "❌ Patient ID not found."
-
-    return f"🔍 Diagnosis: {diagnosis}. {action_plan}"
+    
+    return (f"🔍 Diagnosis: {diagnosis}\n"
+            f"🩺 Medications: {medication}\n"
+            f"⚠️ Precautions: {precaution}\n"
+            f"💡 {action_plan}")
 
 # Function to extract text from PDF reports
 def extract_pdf_report(pdf):
     text = ""
-    try:
-        with pdfplumber.open(pdf.name) as pdf_file:
-            for page in pdf_file.pages:
-                text += page.extract_text()
-    except Exception as e:
-        return f"❌ Error extracting text from PDF: {str(e)}"
-
+    with pdfplumber.open(pdf.name) as pdf_file:
+        for page in pdf_file.pages:
+            text += page.extract_text()
     return text
 
 # Function to analyze uploaded images (X-ray/CT-scan)
 def analyze_image(patient_id, img):
-    try:
-        image = Image.open(img).convert('RGB')
-        transform = transforms.Compose([
-            transforms.Resize((224, 224)),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-        ])
-        image_tensor = transform(image).unsqueeze(0)
-
-        # Run the image through the model (for simplicity, assuming ResNet50 output)
-        output = image_model(image_tensor)
-        _, predicted = torch.max(output, 1)
-
-        # Map prediction to a label (placeholder example)
-        labels = {0: "Normal", 1: "Pneumonia", 2: "Liver Disorder", 3: "COVID-19"}
-        diagnosis = labels.get(predicted.item(), "Unknown")
-
-        # Store diagnosis in the database
-        for patient in patients_db:
-            if patient["ID"] == patient_id:
-                patient["Diagnosis"] = diagnosis
-                break
-        else:
-            return "❌ Patient ID not found."
-
-        return f"🔍 Diagnosis from image: {diagnosis}"
-
-    except Exception as e:
-        return f"❌ Error analyzing image: {str(e)}"
+    image = Image.open(img).convert('RGB')
+    transform = torch.nn.Sequential(
+        torch.nn.Upsample(size=(224, 224)),
+        torch.nn.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    )
+    image_tensor = transform(torch.unsqueeze(torch.tensor(image), 0))
+
+    # Run the image through the model (for simplicity, assuming ResNet50 output)
+    output = image_model(image_tensor)
+    _, predicted = torch.max(output, 1)
+
+    # Map prediction to a label
+    labels = {0: "Normal", 1: "Pneumonia", 2: "Liver Disorder", 3: "COVID-19"}
+    diagnosis = labels.get(predicted.item(), "Unknown")
+
+    # Store diagnosis in the database
+    for patient in patients_db:
+        if patient["ID"] == patient_id:
+            patient["Diagnosis"] = diagnosis
+            break
+    
+    return f"🔍 Diagnosis from image: {diagnosis}"
 
 # Function to display the dashboard
 def show_dashboard():
@@ -170,9 +197,9 @@ with gr.Blocks() as demo:
     with gr.TabItem("Analyze Report (PDF)"):
         pdf_report_interface.render()
     with gr.TabItem("Analyze Image (X-ray/CT)"):
-
         image_interface.render()
     with gr.TabItem("Dashboard"):
         dashboard_interface.render()
 
 demo.launch(share=True)
+