Update app.py

app.py

@@ -2,17 +2,28 @@ import gradio as gr
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib_venn import venn3
+from io import BytesIO
 
 def calculate_probabilities(A, B, C, AB, AC, BC, ABC):
     total = A + B + C - AB - AC - BC + ABC
+    if total == 0:
+        return {
+            "P(A)": 0,
+            "P(B)": 0,
+            "P(C)": 0,
+            "P(A ∩ B)": 0,
+            "P(A ∩ C)": 0,
+            "P(B ∩ C)": 0,
+            "P(A ∩ B ∩ C)": 0,
+        }
 
-    P_A = A / total
-    P_B = B / total
-    P_C = C / total
-    P_AB = AB / total
-    P_AC = AC / total
-    P_BC = BC / total
-    P_ABC = ABC / total
+    P_A = A / total if total > 0 else 0
+    P_B = B / total if total > 0 else 0
+    P_C = C / total if total > 0 else 0
+    P_AB = AB / total if total > 0 else 0
+    P_AC = AC / total if total > 0 else 0
+    P_BC = BC / total if total > 0 else 0
+    P_ABC = ABC / total if total > 0 else 0
     
     return {
         "P(A)": P_A,
@@ -28,12 +39,17 @@ def plot_venn(A, B, C, AB, AC, BC, ABC):
     plt.figure(figsize=(8, 8))
     venn = venn3(subsets=(A, B, AB, C, AC, BC, ABC), set_labels=('A', 'B', 'C'))
     plt.title("Venn Diagram")
-    plt.show()
+    
+    buf = BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    
+    return buf
 
 def main(A, B, C, AB, AC, BC, ABC):
-    plot_venn(A, B, C, AB, AC, BC, ABC)
+    venn_diagram = plot_venn(A, B, C, AB, AC, BC, ABC)
     probabilities = calculate_probabilities(A, B, C, AB, AC, BC, ABC)
-    return probabilities
+    return venn_diagram, probabilities
 
 iface = gr.Interface(
     fn=main,
@@ -47,7 +63,7 @@ iface = gr.Interface(
         gr.Number(label="Cantidad en A ∩ B ∩ C")
     ],
     outputs=[
-        gr.Plot(label="Diagrama de Venn"),
+        gr.Image(type="plot", label="Diagrama de Venn"),
         gr.JSON(label="Probabilidades Calculadas")
     ]
 )