diff --git a/classifier_base.py b/classifier_base.py
index f76919d..1a90da6 100644
--- a/classifier_base.py
+++ b/classifier_base.py
@@ -93,7 +93,7 @@ def compile_mode(self, force=False):
                 self.optimizer = Nadam(self.lr)
             self.model.compile(loss='categorical_crossentropy', optimizer=self.optimizer, metrics=['accuracy'])
 
-    def train(self):
+    def train(self, class_weight=None):
         # calculate files number
         file_num = utils.calculate_file_num(PATH_TRAIN_IMAGES)
         steps_train = file_num // self.batch_size
@@ -127,6 +127,7 @@ def train(self):
                 validation_data=val_generator,
                 validation_steps=steps_val,
                 verbose=1,
+                class_weight=class_weight,
             )
         except KeyboardInterrupt:
             print('\nStop by keyboardInterrupt, try saving weights.')
diff --git a/classifier_xception.py b/classifier_xception.py
index 65193dd..7e43b14 100644
--- a/classifier_xception.py
+++ b/classifier_xception.py
@@ -1,3 +1,4 @@
+import utils
 from classifier_base import BaseClassifier
 from keras.applications import *
 from keras.optimizers import *
@@ -38,4 +39,4 @@ def data_generator(self, path_image, train=True):
     # classifier = XceptionClassifier(lr=2e-4)
     # classifier = XceptionClassifier(lr=2e-5)
     classifier = XceptionClassifier('xception_resize', optimizer=Adam(1e-4))
-    classifier.train()
+    classifier.train(class_weight=utils.calculate_class_weight())
diff --git a/test_augment.py b/test_augment.py
new file mode 100644
index 0000000..e7f3247
--- /dev/null
+++ b/test_augment.py
@@ -0,0 +1,53 @@
+from imgaug import augmenters as iaa
+from PIL import Image
+import imgaug as ia
+import numpy as np
+import os
+
+ia.seed(1)
+
+# Example batch of images.
+# The array has shape (32, 64, 64, 3) and dtype uint8.
+images = np.array(
+    [ia.quokka(size=(64, 64)) for _ in range(32)],
+    dtype=np.uint8
+)
+
+seq = iaa.Sequential([
+    iaa.Fliplr(0.5),  # horizontal flips
+    iaa.Crop(percent=(0, 0.1)),  # random crops
+    # Small gaussian blur with random sigma between 0 and 0.5.
+    # But we only blur about 50% of all images.
+    iaa.Sometimes(0.5,
+                  iaa.GaussianBlur(sigma=(0, 0.5))
+                  ),
+    # Strengthen or weaken the contrast in each image.
+    iaa.ContrastNormalization((0.75, 1.5)),
+    # Add gaussian noise.
+    # For 50% of all images, we sample the noise once per pixel.
+    # For the other 50% of all images, we sample the noise per pixel AND
+    # channel. This can change the color (not only brightness) of the
+    # pixels.
+    iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.05 * 255), per_channel=0.5),
+    # Make some images brighter and some darker.
+    # In 20% of all cases, we sample the multiplier once per channel,
+    # which can end up changing the color of the images.
+    iaa.Multiply((0.8, 1.2), per_channel=0.2),
+    # Apply affine transformations to each image.
+    # Scale/zoom them, translate/move them, rotate them and shear them.
+    iaa.Affine(
+        scale={"x": (0.8, 1.2), "y": (0.8, 1.2)},
+        translate_percent={"x": (-0.2, 0.2), "y": (-0.2, 0.2)},
+        rotate=(-25, 25),
+        shear=(-8, 8)
+    )
+], random_order=True)  # apply augmenters in random order
+
+path = '/Users/zijiao/Desktop/1'
+images_aug = seq.augment_images(images)
+for i, im in enumerate(images_aug):
+    im = Image.fromarray(im)
+    im.show()
+    # with open(os.path.join(path, '%d.jpg' % i), 'wb') as f:
+    #     im.save(f)
+print('Done.')
diff --git a/utils.py b/utils.py
index 3f58228..f61f52e 100644
--- a/utils.py
+++ b/utils.py
@@ -3,6 +3,7 @@
 from keras.applications.xception import preprocess_input
 import keras.backend as K
 import tensorflow as tf
+import config
 import os
 
 
@@ -36,6 +37,15 @@ def ensure_dir(dir):
         os.makedirs(dir)
 
 
+def calculate_class_weight(train_path=config.PATH_TRAIN_IMAGES):
+    if not os.path.isdir(train_path):
+        raise Exception('Dir "%s" not exists.' % train_path)
+    n_classes = [len(os.listdir(os.path.join(train_path, subdir))) for subdir in os.listdir(train_path)]
+    print n_classes
+    n_all = sum(n_classes)
+    return [num / float(n_all) for num in n_classes]
+
+
 def get_best_weights(path_weights, mode='acc', postfix='.h5'):
     if not os.path.isdir(path_weights):
         return None