Update app.py

app.py

@@ -1,26 +1,26 @@
-from flask import Flask, render_template, request, redirect, url_for, session, send_file
+from flask import Flask, render_template, request, send_file, redirect, url_for, session
 import tensorflow as tf
 import numpy as np
 from PIL import Image
+import matplotlib.pyplot as plt
 import pickle
-import io
 import os
-import matplotlib.pyplot as plt
+import io
 from reportlab.lib.pagesizes import A4
 from reportlab.lib import colors
-from reportlab.pdfgen import canvas
-from reportlab.lib.units import inch
+from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
+from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, Image as RLImage
 from datetime import datetime
 import logging
 
 app = Flask(__name__)
-app.secret_key = "e3f6f40bb8b2471b9f07c4025d845be9"  # Replace with secure key if needed
+app.secret_key = "your_secret_key_here"  # Replace with a random string
 
-# Paths
 MODEL_PATH = "skin_lesion_model.h5"
 HISTORY_PATH = "training_history.pkl"
 PLOT_PATH = "/tmp/static/training_plot.png"
-LOGO_PATH = "static/logo.jpg"  # Logo in static folder
+LOGO_PATH = "static/logo.jpg"
+
 IMG_SIZE = (224, 224)
 CONFIDENCE_THRESHOLD = 0.30
 
@@ -35,183 +35,148 @@ label_map = {
     7: "Squamous cell carcinoma"
 }
 
-# Logging setup
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # Load model
-try:
-    logger.info("Loading model from %s", MODEL_PATH)
-    model = tf.keras.models.load_model(MODEL_PATH)
-except Exception as e:
-    logger.error("Failed to load model: %s", str(e))
-    raise
-
-# Load training history and generate plot
-history_dict = {}
+logger.info("Loading model from %s", MODEL_PATH)
+model = tf.keras.models.load_model(MODEL_PATH)
+
+# Load and plot training history
 if os.path.exists(HISTORY_PATH):
-    try:
-        with open(HISTORY_PATH, "rb") as f:
-            history_dict = pickle.load(f)
+    with open(HISTORY_PATH, "rb") as f:
+        history_dict = pickle.load(f)
+
+    if "accuracy" in history_dict:
         os.makedirs("/tmp/static", exist_ok=True)
-        if "accuracy" in history_dict and "val_accuracy" in history_dict:
-            plt.plot(history_dict['accuracy'], label='Train Accuracy')
-            plt.plot(history_dict['val_accuracy'], label='Val Accuracy')
-            plt.xlabel('Epochs')
-            plt.ylabel('Accuracy')
-            plt.title('Training History')
-            plt.legend()
-            plt.grid(True)
-            plt.savefig(PLOT_PATH)
-            plt.close()
-            logger.info("Training plot saved at %s", PLOT_PATH)
-    except Exception as e:
-        logger.error("Failed to process training history: %s", str(e))
+        plt.plot(history_dict['accuracy'], label='Train Accuracy')
+        plt.plot(history_dict.get('val_accuracy', []), label='Val Accuracy')
+        plt.xlabel("Epoch")
+        plt.ylabel("Accuracy")
+        plt.title("Model Training History")
+        plt.legend()
+        plt.grid(True)
+        plt.savefig(PLOT_PATH)
+        plt.close()
+        logger.info("Training plot saved at %s", PLOT_PATH)
 
 def preprocess_image(image_bytes):
-    try:
-        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-        image = image.resize(IMG_SIZE)
-        image_array = tf.keras.utils.img_to_array(image)
-        image_array = np.expand_dims(image_array, axis=0)
-        return image_array / 255.0
-    except Exception as e:
-        logger.error("Image preprocessing failed: %s", str(e))
-        raise
-
-def generate_pdf(report_data, filepath):
-    c = canvas.Canvas(filepath, pagesize=A4)
-    width, height = A4
-
-    # Add logo if exists
-    try:
-        if os.path.exists(LOGO_PATH):
-            c.drawImage(LOGO_PATH, 50, height - 100, width=80, preserveAspectRatio=True, mask='auto')
-    except Exception as e:
-        logger.warning("Could not load logo: %s", str(e))
+    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    image = image.resize(IMG_SIZE)
+    image_array = tf.keras.utils.img_to_array(image)
+    image_array = np.expand_dims(image_array, axis=0)
+    return image_array / 255.0
 
-    # Title
-    c.setFillColor(colors.HexColor("#007ACC"))
-    c.setFont("Helvetica-Bold", 20)
-    c.drawCentredString(width / 2, height - 80, "Skin Lesion Diagnosis Report")
-    c.setStrokeColor(colors.HexColor("#007ACC"))
-    c.setLineWidth(2)
-    c.line(60, height - 90, width - 60, height - 90)
-
-    # Info box background
-    c.setFillColor(colors.lightgrey)
-    c.rect(50, height - 250, width - 100, 140, fill=1, stroke=0)
-
-    # Patient Info
-    c.setFillColor(colors.black)
-    c.setFont("Helvetica-Bold", 12)
-    y = height - 120
-    spacing = 20
-
-    def draw_field(label, value):
-        nonlocal y
-        c.setFont("Helvetica-Bold", 12)
-        c.drawString(70, y, f"{label}:")
-        c.setFont("Helvetica", 12)
-        c.drawString(180, y, value)
-        y -= spacing
-
-    draw_field("Full Name", report_data.get("name", "N/A"))
-    draw_field("Email", report_data.get("email", "N/A"))
-    draw_field("Gender", report_data.get("gender", "N/A"))
-    draw_field("Age", str(report_data.get("age", "N/A")))
-
-    # Prediction
-    y -= 20
-    c.setFont("Helvetica-Bold", 14)
-    c.setFillColor(colors.HexColor("#007ACC"))
-    c.drawString(50, y, "AI Diagnosis Result")
-    c.setFillColor(colors.black)
-    y -= spacing
-    draw_field("Prediction", report_data.get("prediction", "N/A"))
-    draw_field("Confidence", report_data.get("confidence", "N/A"))
-
-    # Optional message
-    message = report_data.get("message", "")
-    if message:
-        y -= 10
-        c.setFont("Helvetica-Oblique", 11)
-        c.setFillColor(colors.red)
-        c.drawString(70, y, message)
-
-    # Timestamp
-    y -= 40
-    c.setFont("Helvetica", 10)
-    c.setFillColor(colors.grey)
-    c.drawString(50, y, f"Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
-
-    c.save()
-
-@app.route("/form", methods=["GET"])
+@app.route("/form")
 def form():
-    return render_template("form.html", history_plot="/training_plot.png")
+    result = session.pop("result", None)
+    patient = session.pop("patient", None)
+    return render_template("form.html", history_plot="/training_plot.png", result=result, patient=patient)
 
 @app.route("/training_plot.png")
 def training_plot():
-    return send_file(PLOT_PATH, mimetype="image/png")
+    return send_file(PLOT_PATH, mimetype='image/png')
 
 @app.route("/predict", methods=["POST"])
 def predict():
-    try:
-        if "image" not in request.files:
-            raise ValueError("⚠ No image uploaded.")
-        
-        image = request.files["image"].read()
-        img_array = preprocess_image(image)
-        prediction = model.predict(img_array)[0]
-        predicted_index = int(np.argmax(prediction))
-        confidence = float(prediction[predicted_index])
-
-        name = request.form.get("name")
-        email = request.form.get("email")
-        gender = request.form.get("gender")
-        age = request.form.get("age")
-
-        if confidence < CONFIDENCE_THRESHOLD:
-            pred_label = "Low confidence"
-            msg = "⚠ This image is not confidently recognized. Please upload a clearer image."
-        else:
-            pred_label = label_map.get(predicted_index, "Unknown")
-            msg = ""
-
-        session["report"] = {
-            "name": name,
-            "email": email,
-            "gender": gender,
-            "age": age,
-            "prediction": pred_label,
+    if "image" not in request.files:
+        return redirect(url_for("form"))
+
+    name = request.form.get("name")
+    email = request.form.get("email")
+    gender = request.form.get("gender")
+    age = request.form.get("age")
+    image = request.files["image"].read()
+
+    img_array = preprocess_image(image)
+    prediction = model.predict(img_array)[0]
+    predicted_index = int(np.argmax(prediction))
+    confidence = float(prediction[predicted_index])
+
+    if confidence < CONFIDENCE_THRESHOLD:
+        result = {
+            "prediction": "Low confidence",
             "confidence": f"{confidence * 100:.2f}%",
-            "message": msg
+            "message": "⚠️ Image not confidently classified. Try uploading a clearer image."
+        }
+    else:
+        result = {
+            "prediction": label_map.get(predicted_index, "Unknown"),
+            "confidence": f"{confidence * 100:.2f}%"
         }
 
-        return redirect(url_for("result"))
-    except Exception as e:
-        logger.error("Prediction error: %s", str(e))
-        return render_template("form.html", history_plot="/training_plot.png", result={
-            "prediction": "Error",
-            "confidence": "N/A",
-            "message": f"An error occurred: {str(e)}"
-        })
-
-@app.route("/result")
-def result():
-    report = session.get("report", {})
-    return render_template("result.html", **report)
-
-@app.route("/download-report")
-def download_report():
-    report = session.get("report", {})
-    if not report:
+    session["result"] = result
+    session["patient"] = {"name": name, "email": email, "gender": gender, "age": age}
+    return redirect(url_for("form"))
+
+@app.route("/download-pdf")
+def download_pdf():
+    patient = session.get("patient")
+    result = session.get("result")
+    if not patient or not result:
         return redirect(url_for("form"))
-    os.makedirs("/tmp/reports", exist_ok=True)
-    filepath = "/tmp/reports/report.pdf"
-    generate_pdf(report, filepath)
-    return send_file(filepath, as_attachment=True)
+
+    buffer = io.BytesIO()
+    doc = SimpleDocTemplate(buffer, pagesize=A4)
+
+    elements = []
+    styles = getSampleStyleSheet()
+    styles.add(ParagraphStyle(name="Title", fontSize=20, textColor=colors.HexColor("#007acc"), spaceAfter=16, alignment=1))
+    styles.add(ParagraphStyle(name="SectionHeader", fontSize=14, textColor=colors.HexColor("#007acc"), spaceBefore=10, spaceAfter=10))
+    styles.add(ParagraphStyle(name="NormalBold", fontSize=12, leading=14, spaceAfter=6, fontName='Helvetica-Bold'))
+
+    # Logo
+    if os.path.exists(LOGO_PATH):
+        elements.append(RLImage(LOGO_PATH, width=100, height=50))
+        elements.append(Spacer(1, 12))
+
+    # Title
+    elements.append(Paragraph("Skin Lesion Diagnosis Report", styles["Title"]))
+    elements.append(Spacer(1, 6))
+
+    # Patient Info Table
+    patient_data = [
+        ["Full Name:", patient["name"]],
+        ["Email:", patient["email"]],
+        ["Gender:", patient["gender"]],
+        ["Age:", patient["age"]],
+    ]
+    patient_table = Table(patient_data, colWidths=[100, 300])
+    patient_table.setStyle(TableStyle([
+        ("BACKGROUND", (0, 0), (-1, -1), colors.whitesmoke),
+        ("FONTNAME", (0, 0), (-1, -1), "Helvetica"),
+        ("FONTSIZE", (0, 0), (-1, -1), 11),
+        ("BOTTOMPADDING", (0, 0), (-1, -1), 8),
+        ("ROWBACKGROUNDS", (0, 0), (-1, -1), [colors.lightgrey, colors.whitesmoke]),
+    ]))
+    elements.append(Paragraph("Patient Information", styles["SectionHeader"]))
+    elements.append(patient_table)
+
+    # Prediction Info Table
+    result_data = [
+        ["Prediction:", result["prediction"]],
+        ["Confidence:", result["confidence"]],
+    ]
+    result_table = Table(result_data, colWidths=[100, 300])
+    result_table.setStyle(TableStyle([
+        ("BACKGROUND", (0, 0), (-1, -1), colors.whitesmoke),
+        ("FONTNAME", (0, 0), (-1, -1), "Helvetica"),
+        ("FONTSIZE", (0, 0), (-1, -1), 11),
+        ("BOTTOMPADDING", (0, 0), (-1, -1), 8),
+        ("ROWBACKGROUNDS", (0, 0), (-1, -1), [colors.lightgrey, colors.whitesmoke]),
+    ]))
+    elements.append(Spacer(1, 16))
+    elements.append(Paragraph("AI Diagnosis Result", styles["SectionHeader"]))
+    elements.append(result_table)
+
+    # Footer
+    elements.append(Spacer(1, 30))
+    date_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    elements.append(Paragraph(f"<font size='10'>Generated on: {date_str}</font>", styles["Normal"]))
+
+    doc.build(elements)
+    buffer.seek(0)
+    return send_file(buffer, as_attachment=True, download_name="diagnosis_report.pdf", mimetype="application/pdf")
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)