Update app.py

app.py

@@ -1,4 +1,4 @@
-import gradio as gr 
+import gradio as gr
 import pdfplumber
 import matplotlib.pyplot as plt
 import numpy as np
@@ -47,7 +47,7 @@ def find_penalty_values(text: str) -> List[float]:
         for match in matches:
             penalty_text = match.group()
             try:
-                if any(word in penalty_text.lower() for word in ['one', 'two', 'three', 'four', 'five', 'six', 'seven', 'eight', 'nine', 'ten', 'hundred', 'thousand', 'million']): 
+                if any(word in penalty_text.lower() for word in ['one', 'two', 'three', 'four', 'five', 'six', 'seven', 'eight', 'nine', 'ten', 'hundred', 'thousand', 'million']):
                     penalty_value = w2n.word_to_num(penalty_text.split('dollars')[0].strip())
                 else:
                     penalty_value = float(re.sub(r'[^\d.]', '', penalty_text))
@@ -87,7 +87,7 @@ def generate_heatmap(risk_level: str) -> plt.Figure:
     """Generate a horizontal heatmap based on detected risk levels"""
     fig, ax = plt.subplots(figsize=(8, 2))
 
-    # Set default values for the risks
+    # Initialize all risks to 0
     risk_values = {
         "Low": 0,
         "Medium": 0,
@@ -102,24 +102,26 @@ def generate_heatmap(risk_level: str) -> plt.Figure:
     elif risk_level == "High":
         risk_values["High"] = 100
 
-    # Create horizontal heatmap using the risk values
-    gradient = np.array([risk_values["Low"], risk_values["Medium"], risk_values["High"]]).reshape(1, -1)
-    ax.imshow(gradient, aspect='auto', cmap='RdYlGn')  # RdYlGn for Green to Red
+    # Define the colors for each risk level
+    colors = ['#28a745', '#ffc107', '#dc3545']  # Green, Orange, Red
 
-    # Set labels for the heatmap
+    # Plot the heatmap with each color representing a risk
+    gradient = np.array([risk_values["Low"], risk_values["Medium"], risk_values["High"]]).reshape(1, -1)
+    ax.imshow(gradient, aspect='auto', cmap='RdYlGn')  # Color scale: Green to Red
+    
+    # Set the labels for the heatmap
     ax.set_yticks([])
     ax.set_xticks([0, 1, 2])
     ax.set_xticklabels(['Low Risk', 'Medium Risk', 'High Risk'])
     
-    # Add text for the risk level on the heatmap
-    ax.text(0, 0.5, f"Low Risk: {risk_values['Low']}%", color='black', ha='center', va='center', fontsize=14, fontweight='bold')
+    # Add the corresponding risk labels
+    ax.text(0, 0.5, f"Low Risk: {risk_values['Low']}%", color='white', ha='center', va='center', fontsize=14, fontweight='bold')
     ax.text(1, 0.5, f"Medium Risk: {risk_values['Medium']}%", color='black', ha='center', va='center', fontsize=14, fontweight='bold')
-    ax.text(2, 0.5, f"High Risk: {risk_values['High']}%", color='black', ha='center', va='center', fontsize=14, fontweight='bold')
+    ax.text(2, 0.5, f"High Risk: {risk_values['High']}%", color='white', ha='center', va='center', fontsize=14, fontweight='bold')
     
     ax.set_axis_off()
     plt.tight_layout()
 
-    # Return as a file object for Gradio to handle
     return fig
 
 def analyze_pdf(file_obj) -> List: