Update app.py

app.py

@@ -72,23 +72,25 @@ elif mode == "Train Mode":
             # Label input
             user_label = st.selectbox("Select the correct label", labels)
 
+            # Train button
+           # Inside the "Train Mode" section
+
+
             # Train button
             if st.button('Train Model'):
                 # Update the model with the user-provided image and label
-                target = np.zeros(len(labels))
-                target[labels.index(user_label)] = 1
-
+                target = np.zeros((1, len(labels)))  # Adjust dimensions for one-hot encoding
+                target[0, labels.index(user_label)] = 1
+        
                 # Reshape the image to match the model's input shape
                 image = tf.image.resize(image, (224, 224))  # adjust the size as needed
                 image = tf.expand_dims(image, axis=0)
-
+        
                 # Create a dataset with the new image and label
                 dataset = tf.data.Dataset.from_tensor_slices((image, target))
                 dataset = dataset.batch(1)  # Batch size 1 as we are using a single image
-
+        
                 # Train the model
                 model.fit(dataset, epochs=1)  # You might want to adjust the number of epochs
-
+        
                 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.")