okeowo1014 commited on
Commit
0b000d9
·
1 Parent(s): 6c9a1c0

Update trainning.py

Browse files
Files changed (1) hide show
  1. trainning.py +2 -1
trainning.py CHANGED
@@ -58,7 +58,8 @@ datagen = ImageDataGenerator(
58
  )
59
 
60
  # Calculate class weights to handle data imbalance
61
- class_weights = compute_class_weight('balanced', np.unique(train_labels), train_labels)
 
62
 
63
  # Model Architecture
64
  model = tf.keras.models.Sequential([
 
58
  )
59
 
60
  # Calculate class weights to handle data imbalance
61
+ # class_weights = compute_class_weight('balanced', np.unique(train_labels), train_labels)
62
+ class_weights = compute_class_weight('balanced', np.unique(train_labels), train_labels.flatten())
63
 
64
  # Model Architecture
65
  model = tf.keras.models.Sequential([