Update app.py

app.py

@@ -61,6 +61,7 @@ def plotdata(t, s, i,r,R0, e=None):
         a.legend()
 
     plt.tight_layout()
+    st.pyplot(fig)
 
 #final model
 def SIR(country,t_infective):
@@ -103,8 +104,8 @@ def main():
             # Show SIR
             SIR_param = SIR(country,recovery)
             
-            st.success(all_location[country]['total_cases'].plot())
-            st.success(SIR(country,recovery))
+            st.success(st.pyplot(all_location[country]['total_cases'].plot()))
+            st.success(st.pyplot(SIR(country,recovery)))
             st.success("SIR model parameters for "+str(country)+" is")
             st.success("R0 = "+str(SIR_param[0]))
             st.success("Beta = "+str(SIR_param[3]))