Update app.py

app.py

@@ -115,6 +115,7 @@ if mode == "Predict Mode":
     except Exception as e:
         st.error(f"An error occurred: {e}. Please contact us EcoClim Solutions at EcoClimSolutions.wordpress.com.")
 
+
 elif mode == "Train Mode":
     # Train the model with a new image and label
     st.header("Train the model with a new image and label")
@@ -132,11 +133,10 @@ elif mode == "Train Mode":
             # Train button
             if st.button('Train Model'):
                 # Update the model with the user-provided image and label
-                
-                # You might want to adjust the number of epochs
-                
-                for i in range(2):
-                    model.fit(image, user_label, epochs=1, batch_size=32)
+                image = image[np.newaxis, ...]  # Add a batch dimension
+                label_index = labels.index(user_label)
+                label_one_hot = tf.one_hot(label_index, len(labels))
+                model.fit(image, label_one_hot, epochs=2, batch_size=1)
                 st.success(f'Model has been trained with the new image and label.')
         except Exception as e:
             st.error(f"An error occurred: {e}. Please contact us EcoClim Solutions at EcoClimSolutions.wordpress.com.")