app.py

# import io
# import os
# import torch
# import torch.nn as nn
# from PIL import Image
# from fastapi import FastAPI, File, UploadFile, Form
# from fastapi.responses import HTMLResponse, FileResponse
# from torchvision import transforms
# from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
# from reportlab.lib.styles import getSampleStyleSheet

# app = FastAPI(title="Multimodal Lung Diagnosis System")

# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# # ---------------- CLASS NAMES ----------------
# CLASS_NAMES = [
#     "COVID",
#     "NORMAL",
#     "PNEUMONIA",
#     "PNEUMOTHORAX"
# ]

# # ---------------- TRANSFORM ----------------
# transform = transforms.Compose([
#     transforms.Resize((224,224)),
#     transforms.Grayscale(1),
#     transforms.ToTensor(),
#     transforms.Normalize([0.5],[0.5])
# ])

# # ---------------- MODEL ----------------
# class LungCNN(nn.Module):
#     def __init__(self):
#         super().__init__()
#         self.features = nn.Sequential(
#             nn.Conv2d(1,32,3,padding=1), nn.ReLU(), nn.MaxPool2d(2),
#             nn.Conv2d(32,64,3,padding=1), nn.ReLU(), nn.MaxPool2d(2),
#             nn.Conv2d(64,128,3,padding=1), nn.ReLU(), nn.MaxPool2d(2),
#             nn.Conv2d(128,256,3,padding=1), nn.ReLU(), nn.MaxPool2d(2)
#         )
#         self.classifier = nn.Sequential(
#             nn.Flatten(),
#             nn.Linear(256*14*14,512),
#             nn.ReLU(),
#             nn.Linear(512,4)
#         )

#     def forward(self,x):
#         return self.classifier(self.features(x))

# # ---------------- LOAD META MODELS ----------------
# models = []
# for i in range(5):
#     model = LungCNN().to(device)
#     model.load_state_dict(
#         torch.load(f"models/meta_model_{i+1}.pth",
#                    map_location=device)
#     )
#     model.eval()
#     models.append(model)

# # ---------------- CLINICAL RISK ENGINE ----------------
# def clinical_analysis(age, spo2, fever, smoking):

#     risk = "Low"

#     if spo2 < 90:
#         risk = "High"

#     if age > 60 and spo2 < 92:
#         risk = "High"

#     if smoking == "Yes":
#         risk = "Moderate"

#     return risk

# # ---------------- RECOMMENDATION ENGINE ----------------
# def generate_recommendations(data, prediction):

#     medical = set()
#     diet = set()
#     avoid = set()
#     lifestyle = set()
#     preventive = set()

#     age = data["age"]
#     spo2 = data["spo2"]
#     fever = data["fever"]
#     smoking = data["smoking"]
#     cough = data["cough"]
#     breathlessness = data["breathlessness"]
#     bmi = data["bmi"]

#     if prediction == "PNEUMOTHORAX":
#         medical.update([
#             "Immediate pulmonologist consultation advised",
#             "Avoid air travel until recovery"
#         ])
#         lifestyle.add("Strict rest required")

#     if spo2 < 90:
#         medical.add("Critical oxygen level detected")
#         preventive.add("Oxygen monitoring required")

#     if smoking == "Yes":
#         medical.add("Smoking cessation strongly advised")
#         avoid.add("Tobacco products")

#     lifestyle.update([
#         "Stay hydrated",
#         "Adequate sleep",
#         "Avoid polluted environments"
#     ])

#     if len(diet) == 0:
#         diet.update([
#             "Balanced nutrition diet",
#             "Fresh fruits & vegetables",
#             "Adequate hydration"
#         ])

#     if len(preventive) == 0:
#         preventive.update([
#             "Regular health checkups",
#             "Maintain active lifestyle"
#         ])

#     return {
#         "medical_advice": list(medical),
#         "diet_plan": list(diet),
#         "foods_to_avoid": list(avoid),
#         "lifestyle_tips": list(lifestyle),
#         "preventive_care": list(preventive)
#     }

# # ---------------- GLOBAL STORAGE ----------------
# latest_report_data = {}

# # ---------------- API ----------------
# @app.post("/predict", response_class=HTMLResponse)
# async def predict(

#     name: str = Form(...),
#     age: int = Form(...),
#     gender: str = Form(...),
#     height: float = Form(...),
#     weight: float = Form(...),
#     smoking: str = Form(...),
#     fever: float = Form(...),
#     cough: str = Form(...),
#     breathlessness: str = Form(...),
#     spo2: int = Form(...),
#     xray: UploadFile = File(...)

# ):

#     bmi = weight / ((height/100) ** 2)

#     image_bytes = await xray.read()
#     image = Image.open(io.BytesIO(image_bytes)).convert("L")
#     image = transform(image).unsqueeze(0).to(device)

#     probs_total = 0
#     for m in models:
#         with torch.no_grad():
#             out = m(image)
#             probs_total += torch.softmax(out, dim=1)

#     avg_probs = probs_total / len(models)
#     confidence, pred_index = torch.max(avg_probs,1)
#     prediction = CLASS_NAMES[pred_index.item()]
#     confidence = round(confidence.item()*100,2)

#     risk = clinical_analysis(age, spo2, fever, smoking)

#     data = {
#         "age": age,
#         "spo2": spo2,
#         "fever": fever,
#         "smoking": smoking,
#         "cough": cough,
#         "breathlessness": breathlessness,
#         "bmi": bmi
#     }

#     recommendations = generate_recommendations(data, prediction)

#     # Store for PDF
#     global latest_report_data
#     latest_report_data = {
#         "name": name,
#         "age": age,
#         "gender": gender,
#         "BMI": round(bmi,2),
#         "smoking": smoking,
#         "prediction": prediction,
#         "confidence": confidence,
#         "risk": risk,
#         "recommendations": recommendations
#     }

#     # HTML Report
#     report_html = f"""
#     <h1>Lung Health Assessment Report</h1>

#     <h2>Patient Information</h2>
#     <p>Name: {name}</p>
#     <p>Age: {age}</p>
#     <p>Gender: {gender}</p>
#     <p>BMI: {round(bmi,2)}</p>
#     <p>Smoking: {smoking}</p>

#     <h2>Radiology Prediction</h2>
#     <p>Disease: <b>{prediction}</b></p>
#     <p>Confidence: {confidence}%</p>
#     <p>Risk Level: {risk}</p>

#     <h2>Medical Advice</h2>
#     <ul>
#     {''.join([f"<li>{i}</li>" for i in recommendations['medical_advice']])}
#     </ul>

#     <h2>Diet Plan</h2>
#     <ul>
#     {''.join([f"<li>{i}</li>" for i in recommendations['diet_plan']])}
#     </ul>

#     <h2>Lifestyle Tips</h2>
#     <ul>
#     {''.join([f"<li>{i}</li>" for i in recommendations['lifestyle_tips']])}
#     </ul>

#     <h2>Preventive Care</h2>
#     <ul>
#     {''.join([f"<li>{i}</li>" for i in recommendations['preventive_care']])}
#     </ul>

#     <br><br>
#     <a href="/download-report">
#     <button style="padding:10px 20px;font-size:16px;">
#     Download Your Health Report (PDF)
#     </button>
#     </a>
#     """

#     return HTMLResponse(content=report_html)

# # ---------------- PDF DOWNLOAD ----------------
# @app.get("/download-report")
# def download_report():

#     global latest_report_data

#     file_path = "Health_Report.pdf"
#     doc = SimpleDocTemplate(file_path)
#     styles = getSampleStyleSheet()
#     content = []

#     content.append(Paragraph("Lung Health Assessment Report", styles['Title']))
#     content.append(Spacer(1,12))

#     for key, value in latest_report_data.items():
#         if key != "recommendations":
#             content.append(Paragraph(f"{key}: {value}", styles['BodyText']))
#             content.append(Spacer(1,8))

#     content.append(Spacer(1,12))
#     content.append(Paragraph("Recommendations:", styles['Heading2']))
#     content.append(Spacer(1,10))

#     for section, items in latest_report_data["recommendations"].items():
#         content.append(Paragraph(section.replace("_"," ").title(), styles['Heading3']))
#         content.append(Spacer(1,6))
#         for item in items:
#             content.append(Paragraph(f"- {item}", styles['BodyText']))
#         content.append(Spacer(1,8))

#     doc.build(content)

#     return FileResponse(file_path, media_type='application/pdf', filename="Health_Report.pdf")

import io
import os
import torch
import torch.nn as nn
from PIL import Image
from fastapi import FastAPI, File, UploadFile, Form, Request
from fastapi.responses import HTMLResponse, FileResponse
from fastapi.templating import Jinja2Templates
from fastapi.staticfiles import StaticFiles
from torchvision import transforms
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
from reportlab.lib.styles import getSampleStyleSheet

# ---------------- APP INIT ----------------
app = FastAPI(title="Multimodal Lung Diagnosis System")

templates = Jinja2Templates(directory="templates")
app.mount("/static", StaticFiles(directory="static"), name="static")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ---------------- CLASS NAMES ----------------
CLASS_NAMES = [
    "COVID",
    "NORMAL",
    "PNEUMONIA",
    "PNEUMOTHORAX"
]

# ---------------- TRANSFORM ----------------
transform = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.Grayscale(1),
    transforms.ToTensor(),
    transforms.Normalize([0.5],[0.5])
])

# ---------------- MODEL ----------------
class LungCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1,32,3,padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(32,64,3,padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(64,128,3,padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(128,256,3,padding=1), nn.ReLU(), nn.MaxPool2d(2)
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(256*14*14,512),
            nn.ReLU(),
            nn.Linear(512,4)
        )

    def forward(self,x):
        return self.classifier(self.features(x))

# ---------------- LOAD META MODELS ----------------
models = []
for i in range(5):
    model = LungCNN().to(device)
    model.load_state_dict(
        torch.load(f"models/meta_model_{i+1}.pth",
                   map_location=device)
    )
    model.eval()
    models.append(model)

# ---------------- CLINICAL RISK ENGINE ----------------
def clinical_analysis(age, spo2, fever, smoking):

    risk = "Low"

    if spo2 < 90:
        risk = "High"

    if age > 60 and spo2 < 92:
        risk = "High"

    if smoking == "Yes":
        risk = "Moderate"

    return risk

# ---------------- RECOMMENDATION ENGINE ----------------
def generate_recommendations(data, prediction):

    medical = set()
    diet = set()
    avoid = set()
    lifestyle = set()
    preventive = set()

    age = data["age"]
    spo2 = data["spo2"]
    fever = data["fever"]
    smoking = data["smoking"]
    cough = data["cough"]
    breathlessness = data["breathlessness"]
    bmi = data["bmi"]

    if prediction == "PNEUMOTHORAX":
        medical.update([
            "Immediate pulmonologist consultation advised",
            "Avoid air travel until recovery"
        ])
        lifestyle.add("Strict rest required")

    if spo2 < 90:
        medical.add("Critical oxygen level detected")
        preventive.add("Oxygen monitoring required")

    if smoking == "Yes":
        medical.add("Smoking cessation strongly advised")
        avoid.add("Tobacco products")

    lifestyle.update([
        "Stay hydrated",
        "Adequate sleep",
        "Avoid polluted environments"
    ])

    if len(diet) == 0:
        diet.update([
            "Balanced nutrition diet",
            "Fresh fruits & vegetables",
            "Adequate hydration"
        ])

    if len(preventive) == 0:
        preventive.update([
            "Regular health checkups",
            "Maintain active lifestyle"
        ])

    return {
        "medical_advice": list(medical),
        "diet_plan": list(diet),
        "foods_to_avoid": list(avoid),
        "lifestyle_tips": list(lifestyle),
        "preventive_care": list(preventive)
    }

# ---------------- GLOBAL STORAGE ----------------
latest_report_data = {}

# ---------------- FORM PAGE ----------------
@app.get("/", response_class=HTMLResponse)
def form_page(request: Request):
    return templates.TemplateResponse(
        request=request,
        name="form.html",
    )

# ---------------- PREDICTION API ----------------
# @app.post("/predict", response_class=HTMLResponse)
# async def predict(
#     request: Request,
#     name: str = Form(...),
#     age: int = Form(...),
#     gender: str = Form(...),
#     height: float = Form(...),
#     weight: float = Form(...),
#     smoking: str = Form(...),
#     fever: float = Form(...),
#     cough: str = Form(...),
#     breathlessness: str = Form(...),
#     spo2: int = Form(...),
#     xray: UploadFile = File(...)
# ):

#     bmi = weight / ((height/100) ** 2)

#     image_bytes = await xray.read()
#     image = Image.open(io.BytesIO(image_bytes)).convert("L")
#     image = transform(image).unsqueeze(0).to(device)

#     probs_total = 0
#     for m in models:
#         with torch.no_grad():
#             out = m(image)
#             probs_total += torch.softmax(out, dim=1)

#     avg_probs = probs_total / len(models)
#     # confidence, pred_index = torch.max(avg_probs,1)
#     # prediction = CLASS_NAMES[pred_index.item()]
#     # confidence = round(confidence.item()*100,2)

# # -------- Prediction --------
# confidence, pred_index = torch.max(avg_probs,1)
# prediction = CLASS_NAMES[pred_index.item()]
# confidence = round(confidence.item()*100,2)

# # -------- All Class Probabilities --------
# probabilities = {}

# for i, cls in enumerate(CLASS_NAMES):
#     probabilities[cls] = round(
#         avg_probs[0][i].item() * 100, 2
#     )
#     risk = clinical_analysis(age, spo2, fever, smoking)

#     data = {
#         "age": age,
#         "spo2": spo2,
#         "fever": fever,
#         "smoking": smoking,
#         "cough": cough,
#         "breathlessness": breathlessness,
#         "bmi": bmi
#     }

#     recommendations = generate_recommendations(data, prediction)

#     global latest_report_data
#     latest_report_data = {
#         "name": name,
#         "age": age,
#         "gender": gender,
#         "BMI": round(bmi,2),
#         "prediction": prediction,
#         "confidence": confidence,
#         "risk": risk,
#         "recommendations": recommendations
#     }

#     return templates.TemplateResponse(
#         "report.html",
#         {
#             "request": request,
#             "name": name,
#             "age": age,
#             "gender": gender,
#             "bmi": round(bmi,2),
#             "prediction": prediction,
#             "confidence": confidence,
#             "risk": risk,
#             "recommendations": recommendations
#         }
#     )

@app.post("/predict", response_class=HTMLResponse)
async def predict(
    request: Request,
    name: str = Form(...),
    age: int = Form(...),
    gender: str = Form(...),
    height: float = Form(...),
    weight: float = Form(...),
    smoking: str = Form(...),
    fever: float = Form(...),
    cough: str = Form(...),
    breathlessness: str = Form(...),
    spo2: int = Form(...),
    xray: UploadFile = File(...)
):

    global latest_report_data

    # ---------------- BMI ----------------
    bmi = weight / ((height / 100) ** 2)

    # ---------------- Image Processing ----------------
    image_bytes = await xray.read()
    image = Image.open(io.BytesIO(image_bytes)).convert("L")
    image = transform(image).unsqueeze(0).to(device)

    probs_total = 0
    for m in models:
        with torch.no_grad():
            out = m(image)
            probs_total += torch.softmax(out, dim=1)

    avg_probs = probs_total / len(models)

    # ---------------- Final Prediction ----------------
    confidence, pred_index = torch.max(avg_probs, 1)
    prediction = CLASS_NAMES[pred_index.item()]
    confidence = round(confidence.item() * 100, 2)

    # ---------------- All Class Probabilities ----------------
    probabilities = {}
    for i, cls in enumerate(CLASS_NAMES):
        probabilities[cls] = round(
            avg_probs[0][i].item() * 100, 2
        )

    # ---------------- Risk ----------------
    risk = clinical_analysis(age, spo2, fever, smoking)

    # ---------------- Recommendation ----------------
    data = {
        "age": age,
        "spo2": spo2,
        "fever": fever,
        "smoking": smoking,
        "cough": cough,
        "breathlessness": breathlessness,
        "bmi": bmi
    }

    recommendations = generate_recommendations(data, prediction)

    # ---------------- Store for PDF ----------------
    latest_report_data = {
        "name": name,
        "age": age,
        "gender": gender,
        "BMI": round(bmi, 2),
        "prediction": prediction,
        "confidence": confidence,
        "risk": risk,
        "recommendations": recommendations,
        "probabilities": probabilities
    }

    # ---------------- Render Report ----------------
    return templates.TemplateResponse(
        request=request,
        name="report.html",
        context={
            "name": name,
            "age": age,
            "gender": gender,
            "bmi": round(bmi, 2),
            "prediction": prediction,
            "confidence": confidence,
            "risk": risk,
            "recommendations": recommendations,
            "probabilities": probabilities,
        },
    )

# ---------------- PDF DOWNLOAD ----------------
@app.get("/download-report")
def download_report():

    global latest_report_data

    file_path = "Health_Report.pdf"
    doc = SimpleDocTemplate(file_path)
    styles = getSampleStyleSheet()
    content = []

    content.append(Paragraph("Lung Health Assessment Report", styles['Title']))
    content.append(Spacer(1,12))

    for key, value in latest_report_data.items():
        if key != "recommendations":
            content.append(Paragraph(f"{key}: {value}", styles['BodyText']))
            content.append(Spacer(1,8))

    content.append(Spacer(1,12))
    content.append(Paragraph("Recommendations:", styles['Heading2']))
    content.append(Spacer(1,10))

    for section, items in latest_report_data["recommendations"].items():
        content.append(Paragraph(section.replace("_"," ").title(), styles['Heading3']))
        content.append(Spacer(1,6))
        for item in items:
            content.append(Paragraph(f"- {item}", styles['BodyText']))
        content.append(Spacer(1,8))

    doc.build(content)

    return FileResponse(
        file_path,
        media_type='application/pdf',
        filename="Health_Report.pdf"
    )