Update app.py

app.py

@@ -61,7 +61,8 @@ def plotdata(t, s, i,r,R0, e=None):
         a.legend()
 
     plt.tight_layout()
-    st.pyplot(fig)
+    
+    return fig
 
 #final model
 def SIR(country,t_infective):
@@ -83,9 +84,8 @@ def SIR(country,t_infective):
   s, i, r = soln.T
   e = None
 
-  plotdata(t, s, i,r,R0)
 
-  return R0,t_infective,gamma,beta
+  return plotdata(t, s, i,r,R0),R0,t_infective,gamma,beta
 
 def main():
     st.title("SIR Model for Monkeypox")
@@ -104,12 +104,12 @@ def main():
             # Show SIR
             SIR_param = SIR(country,recovery)
             
-            st.success(st.pyplot(all_location[country]['total_cases'].plot()))
-            st.success(st.pyplot(SIR(country,recovery)))
+            st.success(st.pyplot(plt.plot(all_location[country]['total_cases'])))
+            st.success(st.pyplot(SIR(country,recovery)[0]))
             st.success("SIR model parameters for "+str(country)+" is")
-            st.success("R0 = "+str(SIR_param[0]))
-            st.success("Beta = "+str(SIR_param[3]))
-            st.success("Gamma = "+str(SIR_param[2]))
+            st.success("R0 = "+str(SIR_param[1]))
+            st.success("Beta = "+str(SIR_param[4]))
+            st.success("Gamma = "+str(SIR_param[3]))
 
 # Run main()
 if __name__ == "__main__":