app.py

import gradio as gr
import tensorflow as tf
import pdfplumber
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import timm
import torch
import pandas as pd

# Load pre-trained zero-shot model for text classification (using PyTorch for compatibility)
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli", framework="pt")

# Pre-trained ResNet50 model for X-ray or image analysis using Timm
image_model = timm.create_model('resnet50', pretrained=True)
image_model.eval()

from tensorflow import keras
from tensorflow.keras.layers import TFSMLayer

# Load the model as a layer (in the SavedModel format)
#eye_model = TFSMLayer('model.h5')

# Patient database
patients_db = []

# Disease details for medical report analyzer
disease_details = {
    "anemia": {"medication": "Iron supplements", "precaution": "Eat iron-rich foods", "doctor": "Hematologist"},
    "viral infection": {"medication": "Antiviral drugs", "precaution": "Stay hydrated", "doctor": "Infectious Disease Specialist"},
    "liver disease": {"medication": "Hepatoprotective drugs", "precaution": "Avoid alcohol", "doctor": "Hepatologist"},
    "diabetes": {"medication": "Metformin or insulin", "precaution": "Monitor sugar levels", "doctor": "Endocrinologist"},
}

# Passwords
doctor_password = "doctor123"

# Loading the custom model for consultation with the doctor
try:
    # Force using the slow tokenizer for compatibility
    tokenizer = AutoTokenizer.from_pretrained("ahmed-7124/NeuraMedAW", use_fast=False)
except Exception as e:
    print(f"Tokenizer error: {e}")
    tokenizer = AutoTokenizer.from_pretrained("ahmed-7124/NeuraMedAW", use_fast=False)

model = AutoModelForCausalLM.from_pretrained("ahmed-7124/NeuraMedAW")

def consult_doctor(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=100)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Functions for the app

def register_patient(name, age, gender, password):
    patient_id = len(patients_db) + 1
    patients_db.append({
        "ID": patient_id,
        "Name": name,
        "Age": age,
        "Gender": gender,
        "Password": password,
        "Diagnosis": "",
        "Medications": "",
        "Precautions": "",
        "Doctor": ""
    })
    return f"✅ Patient {name} registered successfully. Patient ID: {patient_id}"

def analyze_or_extract_report(patient_id, pdf=None, report_text=None):
    if pdf:
        # Extract text from PDF
        with pdfplumber.open(pdf.name) as pdf_file:
            report_text = "".join([page.extract_text() for page in pdf_file.pages])
    
    if not report_text:
        return "❌ Please provide a report text or upload a PDF."

    # Analyze the report
    candidate_labels = list(disease_details.keys())
    result = classifier(report_text, candidate_labels)
    diagnosis = result['labels'][0]

    # Update patient's record
    medication = disease_details[diagnosis]['medication']
    precaution = disease_details[diagnosis]['precaution']
    doctor = disease_details[diagnosis]['doctor']
    for patient in patients_db:
        if patient['ID'] == patient_id:
            patient.update(Diagnosis=diagnosis, Medications=medication, Precautions=precaution, Doctor=doctor)
    return f"🔍 Diagnosis: {diagnosis}\n💊 Medication: {medication}\n⚠ Precaution: {precaution}\n👩‍⚕ Recommended Doctor: {doctor}"


# def extract_pdf_report(pdf):
#     text = ""
#     with pdfplumber.open(pdf.name) as pdf_file:
#         for page in pdf_file.pages:
#             text += page.extract_text()
#     return text
#
# def predict_eye_disease(input_image):
#     input_image = tf.image.resize(input_image, [224, 224]) / 255.0
#     input_image = tf.expand_dims(input_image, 0)
#     predictions = eye_model(input_image)
#     labels = ['Cataract', 'Conjunctivitis', 'Glaucoma', 'Normal']
#     confidence_scores = {labels[i]: round(predictions[i] * 100, 2) for i in range(len(labels))}
#     if confidence_scores['Normal'] > 50:
#         return f"Congrats! No disease detected. Confidence: {confidence_scores['Normal']}%"
#     return "\n".join([f"{label}: {confidence}%" for label, confidence in confidence_scores.items()])

def doctor_space(patient_id):
    for patient in patients_db:
        if patient["ID"] == patient_id:
            return f"⚠ Precautions: {patient['Precautions']}\n👩‍⚕ Recommended Doctor: {patient['Doctor']}"
    return "❌ Patient not found. Please check the ID."

def pharmacist_space(patient_id):
    for patient in patients_db:
        if patient["ID"] == patient_id:
            return f"💊 Medications: {patient['Medications']}"
    return "❌ Patient not found. Please check the ID."

def patient_dashboard(patient_id, password):
    for patient in patients_db:
        if patient["ID"] == patient_id and patient["Password"] == password:
            return (f"🩺 Name: {patient['Name']}\n"
                    f"📋 Diagnosis: {patient['Diagnosis']}\n"
                    f"💊 Medications: {patient['Medications']}\n"
                    f"⚠ Precautions: {patient['Precautions']}\n"
                    f"👩‍⚕ Recommended Doctor: {patient['Doctor']}")
    return "❌ Access Denied: Invalid ID or Password."

def doctor_dashboard(password):
    if password != doctor_password:
        return "❌ Access Denied: Incorrect Password"
    if not patients_db:
        return "No patient records available."
    details = []
    for patient in patients_db:
        details.append(f"🩺 Name: {patient['Name']}\n"
                       f"📋 Diagnosis: {patient['Diagnosis']}\n"
                       f"💊 Medications: {patient['Medications']}\n"
                       f"⚠ Precautions: {patient['Precautions']}\n"
                       f"👩‍⚕ Recommended Doctor: {patient['Doctor']}")
    return "\n\n".join(details)

# Gradio Interfaces
registration_interface = gr.Interface(
    fn=register_patient,
    inputs=[ 
        gr.Textbox(label="Patient Name"),
        gr.Number(label="Age"),
        gr.Radio(label="Gender", choices=["Male", "Female", "Other"]),
        gr.Textbox(label="Set Password", type="password"),
    ],
    outputs="text",
)

#pdf_extraction_interface = gr.Interface(
 #   fn=extract_pdf_report,
  #  inputs=gr.File(label="Upload PDF Report"),
   # outputs="text",
#)

# report_analysis_interface = gr.Interface(
#     fn=analyze_report,
#     inputs=[ 
#         gr.Number(label="Patient ID"),
#         gr.Textbox(label="Report Text"),
#     ],
#     outputs="text",
# )

# Unified Gradio Interface
analyze_report_interface = gr.Interface(
    fn=analyze_or_extract_report,
    inputs=[
        gr.Number(label="Patient ID"),
        gr.File(label="Upload PDF Report"),  # Removed optional=True
        gr.Textbox(label="Report Text (Optional)"),
    ],
    outputs="text",
)

# eye_disease_interface = gr.Interface(
#     fn=predict_eye_disease,
#     inputs=gr.Image(label="Upload an Eye Image", type="numpy"),
#     outputs="text",
# )

doctor_space_interface = gr.Interface(
    fn=doctor_space,
    inputs=gr.Number(label="Patient ID"),
    outputs="text",
)

pharmacist_space_interface = gr.Interface(
    fn=pharmacist_space,
    inputs=gr.Number(label="Patient ID"),
    outputs="text",
)

patient_dashboard_interface = gr.Interface(
    fn=patient_dashboard,
    inputs=[ 
        gr.Number(label="Patient ID"),
        gr.Textbox(label="Password", type="password"),
    ],
    outputs="text",
)

doctor_dashboard_interface = gr.Interface(
    fn=doctor_dashboard,
    inputs=gr.Textbox(label="Doctor Password", type="password"),
    outputs="text",
)

consult_doctor_interface = gr.Interface(
    fn=consult_doctor,
    inputs=gr.Textbox(label="Enter Your Query for the Doctor"),
    outputs="text",
)

# Gradio App Layout
with gr.Blocks() as app:
    gr.Markdown("# Medico GPT")
    
    with gr.Tab("Patient Registration"):
        registration_interface.render()
    
    # with gr.Tab("Analyze Medical Report"):
    #     report_analysis_interface.render()
    with gr.Tab("Analyze Medical Report"):
        analyze_report_interface.render()
    
    with gr.Tab("Extract PDF Report"):
        pdf_extraction_interface.render()
    
    # with gr.Tab("Ophthalmologist Space"):
    #     eye_disease_interface.render()
    
    with gr.Tab("Doctor Space"):
        doctor_space_interface.render()
    
    with gr.Tab("Pharmacist Space"):
        pharmacist_space_interface.render()
    
    with gr.Tab("Patient Dashboard"):
        patient_dashboard_interface.render()
    
    with gr.Tab("Doctor Dashboard"):
        doctor_dashboard_interface.render()

    with gr.Tab("Doctor Consult"):
        consult_doctor_interface.render()

# Launch the app
app.launch(share=True)