Update app.py

app.py

@@ -30,18 +30,15 @@ def main():
             # Show prediction
             if result[0] == 0:
                 result = "Iris-setosa"
+                st.image(image_setosa, caption='Iris Setosa', width = 350)
             elif result[0] == 1:
                 result = "Iris-versicolor"
+                st.image(image_versicolor, caption='Iris Versicolor', width = 350)
             else:
                 result = "Iris-virginica"
+                st.image(image_virginica, caption='Iris Virginica', width = 350)
             st.success(f"Your predicted class is {result}")
             
-            if result[0] == 0:
-                st.image(image_setosa, caption='Iris Setosa', width = 350)
-            elif result[0] == 1:
-                st.image(image_versicolor, caption='Iris Versicolor', width = 350)
-            else:
-                st.image(image_virginica, caption='Iris Virginica', width = 350)
                 
 if __name__ == '__main__':
     main()