Update legal_document_analysis.py

legal_document_analysis.py

@@ -7,6 +7,7 @@ from docx import Document
 import matplotlib.pyplot as plt
 import io
 import base64
+import tempfile
 from email.mime.multipart import MIMEMultipart
 from email.mime.text import MIMEText
 from email.mime.application import MIMEApplication
@@ -238,14 +239,11 @@ def plot_risk_assessment_matrix(detected_risks):
     for i in range(len(detected_risks)):
         ax.annotate(detected_risks[i]['phrase'], (likelihood[i], impact[i]))
 
-    buf = io.BytesIO()
-    plt.savefig(buf, format="png", bbox_inches='tight')
-    buf.seek(0)
-    
-    img_str = base64.b64encode(buf.read()).decode('utf-8')
-    buf.close()
-    
-    return img_str
+    # Save to a temporary file
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as tmpfile:
+        plt.savefig(tmpfile.name, format="png", bbox_inches='tight')
+        plt.close()
+        return tmpfile.name  # Return the file path
 
 # Function to plot risk level distribution pie chart
 def plot_risk_level_distribution(detected_risks):
@@ -258,14 +256,11 @@ def plot_risk_level_distribution(detected_risks):
 
     plt.title("Risk Level Distribution", fontsize=10)
 
-    buf = io.BytesIO()
-    plt.savefig(buf, format="png", bbox_inches='tight')
-    buf.seek(0)
-    
-    img_str = base64.b64encode(buf.read()).decode('utf-8')
-    buf.close()
-    
-    return img_str
+    # Save to a temporary file
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as tmpfile:
+        plt.savefig(tmpfile.name, format="png", bbox_inches='tight')
+        plt.close()
+        return tmpfile.name  # Return the file path
 
 # Function to plot risks by type bar chart
 def plot_risks_by_type(detected_risks):
@@ -278,14 +273,11 @@ def plot_risks_by_type(detected_risks):
     ax.set_title("Risks by Type", fontsize=10)
     ax.set_ylabel("Count")
 
-    buf = io.BytesIO()
-    plt.savefig(buf, format="png", bbox_inches='tight')
-    buf.seek(0)
-
-    img_str = base64.b64encode(buf.read()).decode('utf-8')
-    buf.close()
-
-    return img_str
+    # Save to a temporary file
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as tmpfile:
+        plt.savefig(tmpfile.name, format="png", bbox_inches='tight')
+        plt.close()
+        return tmpfile.name  # Return the file path
 
 # Function to plot stacked bar chart of risks by level
 def plot_stacked_bar_chart(detected_risks):
@@ -300,14 +292,11 @@ def plot_stacked_bar_chart(detected_risks):
     ax.set_title("Stacked Bar Chart of Risks by Level", fontsize=10)
     ax.set_ylabel("Count")
 
-    buf = io.BytesIO()
-    plt.savefig(buf, format="png", bbox_inches='tight')
-    buf.seek(0)
-
-    img_str = base64.b64encode(buf.read()).decode('utf-8')
-    buf.close()
-
-    return img_str
+    # Save to a temporary file
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as tmpfile:
+        plt.savefig(tmpfile.name, format="png", bbox_inches='tight')
+        plt.close()
+        return tmpfile.name  # Return the file path
 
 # Function to plot risk heatmap
 def plot_risk_heatmap(detected_risks):
@@ -324,21 +313,14 @@ def plot_risk_heatmap(detected_risks):
     sns.heatmap(heatmap_data.pivot_table(index='Risk Level', values='Count'), annot=True, cmap='YlGnBu', ax=ax)
     ax.set_title("Risk Heatmap")
 
-    buf = io.BytesIO()
-    plt.savefig(buf, format="png", bbox_inches='tight')
-    buf.seek(0)
-
-    img_str = base64.b64encode(buf.read()).decode('utf-8')
-    buf.close()
-
-    return img_str
-
-# Function to convert base64 to image
-def base64_to_image(data):
-    return io.BytesIO(base64.b64decode(data))
+    # Save to a temporary file
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as tmpfile:
+        plt.savefig(tmpfile.name, format="png", bbox_inches='tight')
+        plt.close()
+        return tmpfile.name  # Return the file path
 
 # Function to generate PDF document with improved aesthetics
-def generate_pdf_analysis(document_text, summary, detected_clauses, hidden_obligations, detected_risks, risk_assessment_matrix, risk_level_distribution, risks_by_type, stacked_bar_chart, risk_heatmap):
+def generate_pdf_analysis(document_text, summary, detected_clauses, hidden_obligations, detected_risks, risk_assessment_matrix_path, risk_level_distribution_path, risks_by_type_path, stacked_bar_chart_path, risk_heatmap_path):
     pdf = FPDF()
     pdf.add_page()
 
@@ -371,15 +353,15 @@ def generate_pdf_analysis(document_text, summary, detected_clauses, hidden_oblig
     pdf.ln(10)
 
     # Add visualizations for risks
-    pdf.image(base64_to_image(risk_assessment_matrix), x=10, y=pdf.get_y(), w=90)
-    pdf.image(base64_to_image(risk_level_distribution), x=110, y=pdf.get_y()-50, w=90)  # Position next to the first image
+    pdf.image(risk_assessment_matrix_path, x=10, y=pdf.get_y(), w=90)
+    pdf.image(risk_level_distribution_path, x=110, y=pdf.get_y()-50, w=90)  # Position next to the first image
     pdf.ln(60)
 
-    pdf.image(base64_to_image(risks_by_type), x=10, y=pdf.get_y(), w=90)
-    pdf.image(base64_to_image(stacked_bar_chart), x=110, y=pdf.get_y()-50, w=90)  # Position next to the previous image
+    pdf.image(risks_by_type_path, x=10, y=pdf.get_y(), w=90)
+    pdf.image(stacked_bar_chart_path, x=110, y=pdf.get_y()-50, w=90)  # Position next to the previous image
     pdf.ln(60)
 
-    pdf.image(base64_to_image(risk_heatmap), x=10, y=pdf.get_y(), w=190)  # Fit image to width
+    pdf.image(risk_heatmap_path, x=10, y=pdf.get_y(), w=190)  # Fit image to width
     pdf.ln(10)
 
     # Footer
@@ -559,18 +541,18 @@ def display_legal_analysis_page():
                 st.write("No risks detected.")
 
             # Generate all visualizations
-            risk_assessment_matrix = plot_risk_assessment_matrix(detected_risks)
-            risk_level_distribution = plot_risk_level_distribution(detected_risks)
-            risks_by_type = plot_risks_by_type(detected_risks)
-            stacked_bar_chart = plot_stacked_bar_chart(detected_risks)
-            risk_heatmap = plot_risk_heatmap(detected_risks)
+            risk_assessment_matrix_path = plot_risk_assessment_matrix(detected_risks)
+            risk_level_distribution_path = plot_risk_level_distribution(detected_risks)
+            risks_by_type_path = plot_risks_by_type(detected_risks)
+            stacked_bar_chart_path = plot_stacked_bar_chart(detected_risks)
+            risk_heatmap_path = plot_risk_heatmap(detected_risks)
 
             # Display the charts
-            st.image(f"data:image/png;base64,{risk_assessment_matrix}", caption="Risk Assessment Matrix")
-            st.image(f"data:image/png;base64,{risk_level_distribution}", caption="Risk Level Distribution")
-            st.image(f"data:image/png;base64,{risks_by_type}", caption="Risks by Type")
-            st.image(f"data:image/png;base64,{stacked_bar_chart}", caption="Stacked Bar Chart of Risks by Level")
-            st.image(f"data:image/png;base64,{risk_heatmap}", caption="Risk Heatmap")
+            st.image(risk_assessment_matrix_path, caption="Risk Assessment Matrix")
+            st.image(risk_level_distribution_path, caption="Risk Level Distribution")
+            st.image(risks_by_type_path, caption="Risks by Type")
+            st.image(stacked_bar_chart_path, caption="Stacked Bar Chart of Risks by Level")
+            st.image(risk_heatmap_path, caption="Risk Heatmap")
 
         with tabs[5]:
             st.subheader("Suggestions for Improvement")
@@ -591,7 +573,7 @@ def display_legal_analysis_page():
             # Download PDF Analysis Button
             st.subheader("Download Analysis as PDF")
             pdf_buffer = io.BytesIO()
-            pdf = generate_pdf_analysis(document_text, summary, detected_clauses, hidden_obligations, detected_risks, risk_assessment_matrix, risk_level_distribution, risks_by_type, stacked_bar_chart, risk_heatmap)
+            pdf = generate_pdf_analysis(document_text, summary, detected_clauses, hidden_obligations, detected_risks, risk_assessment_matrix_path, risk_level_distribution_path, risks_by_type_path, stacked_bar_chart_path, risk_heatmap_path)
             pdf.output(pdf_buffer, 'F')
             pdf_buffer.seek(0)