Update app.py

app.py

@@ -4,7 +4,7 @@ import torch
 from transformers import pipeline
 import time, logging, re
 import matplotlib
-matplotlib.use('Agg')  # Use non-interactive backend for headless environments
+matplotlib.use('Agg')
 import matplotlib.pyplot as plt
 import io
 from PIL import Image
@@ -34,7 +34,30 @@ def visualize_chunk_status(chunk_data):
     buf = io.BytesIO()
     plt.savefig(buf, format='png')
     buf.seek(0)
-    plt.close(fig)  # Release memory
+    plt.close(fig)
+    return Image.open(buf)
+
+def create_summary_flowchart(summaries):
+    fig, ax = plt.subplots(figsize=(6, len(summaries) * 0.8 + 1))
+    ax.axis('off')
+    
+    ypos = list(range(len(summaries) * 2, 0, -2))
+    boxprops = dict(boxstyle="round,pad=0.5", facecolor="lightblue", edgecolor="black")
+
+    for i, (y, summary) in enumerate(zip(ypos, summaries)):
+        summary_text = summary.split("**Chunk")[1] if summary.startswith("**Chunk") else summary
+        summary_text = summary_text.strip().replace("**:", ":")[:120] + ("..." if len(summary) > 120 else "")
+        ax.text(0.5, y, summary_text, ha='center', va='center', bbox=boxprops, fontsize=9, wrap=True)
+
+        if i < len(summaries) - 1:
+            ax.annotate('', xy=(0.5, y - 1), xytext=(0.5, y - 0.2),
+                        arrowprops=dict(arrowstyle="->", lw=1.5))
+
+    buf = io.BytesIO()
+    plt.tight_layout()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    plt.close(fig)
     return Image.open(buf)
 
 def summarize_file(file_bytes):
@@ -49,10 +72,10 @@ def summarize_file(file_bytes):
         text = re.sub(r"\s+", " ", text).strip()
         text = "".join(c for c in text if ord(c) < 128)
     except Exception as e:
-        return f"❌ Text extraction failed: {str(e)}", None
+        return f"❌ Text extraction failed: {str(e)}", None, None
 
     if not text.strip():
-        return "❌ No text found", None
+        return "❌ No text found", None, None
 
     text = text[:300000]
     chunks = [text[i:i+2000] for i in range(0, len(text), 2000)]
@@ -82,18 +105,20 @@ def summarize_file(file_bytes):
         chunk_info.append(chunk_result)
 
     final_summary = f"**Chars**: {len(text)}\n**Time**: {time.time()-start:.2f}s\n\n" + "\n\n".join(summaries)
-    image = visualize_chunk_status(chunk_info)
-    return final_summary, image
+    process_img = visualize_chunk_status(chunk_info)
+    flow_img = create_summary_flowchart(summaries)
+    return final_summary, process_img, flow_img
 
 demo = gr.Interface(
     fn=summarize_file,
     inputs=gr.File(label="📄 Upload PDF", type="binary"),
     outputs=[
         gr.Textbox(label="📝 Summarized Output"),
-        gr.Image(label="📊 Visual Process Flow", type="pil")
+        gr.Image(label="📊 Chunk Status", type="pil"),
+        gr.Image(label="🔁 Flowchart Summary", type="pil")
     ],
     title="AI-Powered PDF Summarizer",
-    description="Summarizes long PDFs (up to 300,000 characters) and visualizes chunk-level automation status."
+    description="Summarizes long PDFs (up to 300,000 characters) and visualizes chunk processing + flow of content."
 )
 
 if __name__ == "__main__":
@@ -102,3 +127,4 @@ if __name__ == "__main__":
     except Exception as e:
         print(f"❌ Gradio launch failed: {str(e)}")
 
+