Update app.py

app.py

@@ -184,22 +184,20 @@ def segment_image(image: Image.Image, model: nn.Module) -> dict:
 def generate_pdf_report(original_img: Image.Image, mask: np.ndarray, overlay: Image.Image, 
                        diagnosis: str, grade: str, bright_area: float, red_area: float):
     try:
-        # Create PDF in memory
         pdf = FPDF()
         pdf.add_page()
         
-        # Use built-in font that supports basic characters
+        # Header and patient info
         pdf.set_font("helvetica", "B", 16)
         pdf.cell(text="Diabetic Retinopathy Diagnosis Report", new_x="LMARGIN", new_y="NEXT", align='C')
         pdf.ln(10)
         
-        # Patient Info
         pdf.set_font("helvetica", "", 12)
         pdf.cell(text="Patient: ___________________________", new_x="LMARGIN", new_y="NEXT")
         pdf.cell(text="Date: _____________________________", new_x="LMARGIN", new_y="NEXT")
         pdf.ln(10)
         
-        # Diagnosis
+        # Diagnosis section
         pdf.set_font("helvetica", "B", 14)
         pdf.cell(text="Diagnosis:", new_x="LMARGIN", new_y="NEXT")
         pdf.set_font("helvetica", "", 12)
@@ -207,7 +205,7 @@ def generate_pdf_report(original_img: Image.Image, mask: np.ndarray, overlay: Im
         pdf.cell(text=f"UK Grading: {grade}", new_x="LMARGIN", new_y="NEXT")
         pdf.ln(10)
         
-        # Lesion Analysis
+        # Lesion analysis
         pdf.set_font("helvetica", "B", 14)
         pdf.cell(text="Lesion Analysis:", new_x="LMARGIN", new_y="NEXT")
         pdf.set_font("helvetica", "", 12)
@@ -216,30 +214,38 @@ def generate_pdf_report(original_img: Image.Image, mask: np.ndarray, overlay: Im
         pdf.cell(text=f"Total Affected Area: {bright_area + red_area:.2f}%", new_x="LMARGIN", new_y="NEXT")
         pdf.ln(15)
         
-        # Helper function to add images
-        def add_pdf_image(pdf, image, title):
-            img_byte_arr = io.BytesIO()
-            resized_img = image.resize((150, 150))  
-            resized_img.save(img_byte_arr, format='PNG')
-            img_bytes = img_byte_arr.getvalue()
-            pdf.set_font("helvetica", "B", 12)
-            pdf.cell(text=title, new_x="LMARGIN", new_y="NEXT")
-            pdf.image(io.BytesIO(img_bytes), x=10, w=100)  
-            pdf.ln(5)
+        # Original image on first page
+        pdf.set_font("helvetica", "B", 12)
+        pdf.cell(text="Original Retinal Image:", new_x="LMARGIN", new_y="NEXT")
+        img_byte_arr = io.BytesIO()
+        original_img.save(img_byte_arr, format='PNG')
+        pdf.image(io.BytesIO(img_byte_arr.getvalue()), x=10, w=100)
+        pdf.ln(10)
+        
+        # Add new page for segmentation results
+        pdf.add_page()
+        
+        # Segmentation mask
+        pdf.set_font("helvetica", "B", 12)
+        pdf.cell(text="Lesion Segmentation Mask:", new_x="LMARGIN", new_y="NEXT")
+        img_byte_arr = io.BytesIO()
+        Image.fromarray((mask * 85).astype(np.uint8)).save(img_byte_arr, format='PNG')
+        pdf.image(io.BytesIO(img_byte_arr.getvalue()), x=10, w=100)
+        pdf.ln(10)
         
-        # Add images
-        add_pdf_image(pdf, original_img, "Original Retinal Image:")
-        add_pdf_image(pdf, Image.fromarray((mask * 85).astype(np.uint8)), "Lesion Segmentation Mask:")
-        add_pdf_image(pdf, overlay, "Lesion Overlay:")
+        # Lesion overlay
+        pdf.set_font("helvetica", "B", 12)
+        pdf.cell(text="Lesion Overlay:", new_x="LMARGIN", new_y="NEXT")
+        img_byte_arr = io.BytesIO()
+        overlay.save(img_byte_arr, format='PNG')
+        pdf.image(io.BytesIO(img_byte_arr.getvalue()), x=10, w=100)
         
-        # Footer
+        # Footer on last page
         pdf.ln(10)
         pdf.set_font("helvetica", "I", 10)
         pdf.cell(text="This report was generated by DR Analysis System", new_x="LMARGIN", new_y="NEXT", align='C')
         
-        # Get PDF as bytes - updated for fpdf2>=2.7.6
-        pdf_bytes = bytes(pdf.output()) # Returns bytes directly in newer versions
-        return pdf_bytes
+        return bytes(pdf.output())
     
     except Exception as e:
         st.error(f"PDF generation failed: {str(e)}")