Update app.py

app.py

@@ -27,10 +27,19 @@ def main():
             # Show prediction
             if result[0] == 0:
                 result = "Iris-setosa"
-                st.image(image, caption='Iris')
             elif result[0] == 1:
                 result = "Iris-versicolor"
             else:
                 result = "Iris-virginica"
             st.success(f"Your predicted class is {result}")
             
+            if result[0] == 0:
+                st.image('iris.png')
+            elif result[0] == 1:
+                st.image('iris.png')
+            else:
+                st.image('iris.png')
+                
+if __name__ == '__main__':
+    main()
+