updated README with additional information on model training

README.md

@@ -14,20 +14,37 @@ The segmentation model was trained on a dataset of 800 total images from the Jig
 
 ## Model Description
 
-Keras implementation of Butterfly UNet segmentation model. The model is responsible for taking an input image (256 x 256 x 3) and generating segmentation masks for all classes below that are found in the image.
+Keras implementation of Butterfly UNet segmentation model. The model is responsible for taking an input image (256 x 256 x 3) and generating segmentation masks for all classes below that are found in the image. Vertical flips were applied as data augmentations prior to training the model. Images are converted to grayscale before being fed into the model. The model is trained with categorical cross entropy loss. 
 
 ### Segmentation Classes
-- background
-- generic
-- right_forewing
-- left_forewing
-- right_hindwing
-- left_hindwing
-- ruler
-- white_balance
-- label
-- color_card
-- body
+[`pixel class`] corresponding category
+
+- [0] background
+- [1] generic
+- [2] right_forewing
+- [3] left_forewing
+- [4] right_hindwing
+- [5] left_hindwing
+- [6] ruler
+- [7] white_balance
+- [8] label
+- [9] color_card
+- [19] body
+
+## Class weights 
+The model was fed class weights to help improve performance on wing segmentation categories (pixel classes: 2,3,4,5). The class weights dictionary used is:
+
+{0: 1.0,
+ 1: 1.0,
+ 2: 3.0,
+ 3: 3.0,
+ 4: 3.0,
+ 5: 3.0,
+ 6: 1.0,
+ 7: 1.0,
+ 8: 1.0,
+ 9: 1.0,
+ 10: 1.0}
 
 **Developed by:** Michelle Ramirez