Update README.md

README.md

@@ -5,62 +5,85 @@ library_name: keras
 ## Model description
 
 This model identifies contrails in satellite images. It takes pre-processed .npy files (images) as its inputs, and returns a "mask" image showing only the contrails overlayed on the same area. 
-We used a TransUNet model architecture ...
+We used a UNet model architecture ...
 
-## Intended uses & limitations
 
-Note, this is in progress - 
 
-We hope that data scientists and researchers focused on reducing contrails (towards the goal of reducing global warming) will use this model to improve their work. 
-There are current efforts underway to develop software that re-routes planes to avoid contrails. Researchers are building models to predict contrails based on atmospheric conditions and other factors, but they need a way to validate those predictions.
-That's where we come in. Let's say you have a model that suggests there should be contrails in an image (based on the time/location the picture was taken).
-Our model can find the contrails (or lack thereof) in your image without a human labeler, allowing you to validate whether your predictions were correct. 
-This becomes valuable at scale, when you need to validate your model on many images - contrail detection is a tough task for humans and machines alike!
+## Intended uses & limitations
 
+We hope that data scientists and researchers focused on reducing contrails (towards the goal of reducing global warming) will use this model to improve their work. 
+There are current efforts underway to develop software that re-routes planes to avoid contrails. 
+Researchers are building models to predict contrails based on atmospheric conditions and other factors, but they need a way to validate those predictions.
 
 ## How to Get Started with the Model
 
 Use the code below to get started with the model.
 
 ```
+#Required imports and Huggingface authentication
+import os
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
+os.environ["SM_FRAMEWORK"] = "tf.keras"
+import segmentation_models as sm
+import tensorflow as tf
+from huggingface_hub import from_pretrained_keras
 from huggingface_hub import notebook_login
-notebook_login()
+from PIL import Image
+import numpy as np
+import matplotlib.pyplot as plt
 
-from huggingface_hub import from_pretrained_keras
 
+weights = [0.5,0.5] # hyper parameter
+
+dice_loss = sm.losses.DiceLoss(class_weights = weights)
+focal_loss = sm.losses.CategoricalFocalLoss()
+TOTAL_LOSS_FACTOR = 5
+total_loss = dice_loss + (TOTAL_LOSS_FACTOR * focal_loss)
+
+def jaccard_coef(y_true, y_pred):
+    """
+    Defines custom jaccard coefficient metric
+    """
+    
+    y_true_flatten = K.flatten(y_true)
+    y_pred_flatten = K.flatten(y_pred)
+    intersection = K.sum(y_true_flatten * y_pred_flatten)
+    final_coef_value = (intersection + 1.0) / (K.sum(y_true_flatten) + K.sum(y_pred_flatten) - intersection + 1.0)
+    return final_coef_value
+
+metrics = [tf.keras.metrics.MeanIoU(num_classes=2, sparse_y_true= False, sparse_y_pred=False, name="Mean IOU")]
+notebook_login() 
+
+# Load model from Huggingface Hub
 model = from_pretrained_keras("MIDSCapstoneTeam/ContrailSentinel", custom_objects={'dice_loss_plus_5focal_loss': total_loss, 'jaccard_coef': jaccard_coef, 'IOU score' : sm.metrics.IOUScore(threshold=0.9, name="IOU score"), 'Dice Coeficient' : sm.metrics.FScore(threshold=0.6, name="Dice Coeficient")}, compile=False)
-```
+model.compile(metrics=metrics)
+
+# Inference -- User needs to specify the image path where label and ash images are stored
+label = np.load({Image path} + 'human_pixel_masks.npy')
+ash_image = np.load({Image path} + 'ash_image.npy')[...,4] 
+y_pred = model.predict(ash_image.reshape(1,256, 256, 3))
+prediction = np.argmax(y_pred[0], axis=2).reshape(256,256,1)
 
+fig, ax = plt.subplots(1, 2, figsize=(9, 5))
+fig.tight_layout(pad=5.0)
+ax[1].set_title("Contrail prediction")
+ax[1].imshow(ash_image)
+ax[1].imshow(prediction)
+ax[1].axis('off')
+
+ax[0].set_title("False colored satellite image")
+ax[0].imshow(ash_image)
+ax[0].axis('off')
+
+```
 
 ## Training and evaluation data
 
 (will add more info here) 
 OpenContrails dataset [here](https://arxiv.org/abs/2304.02122)
+Can get images 
 
 ## Training procedure
 
 
 
-### Training hyperparameters
-
-The following hyperparameters were used during training:
-
-| Hyperparameters | Value |
-| :-- | :-- |
-| name | RMSprop |
-| weight_decay | None |
-| clipnorm | None |
-| global_clipnorm | None |
-| clipvalue | None |
-| use_ema | False |
-| ema_momentum | 0.99 |
-| ema_overwrite_frequency | 100 |
-| jit_compile | True |
-| is_legacy_optimizer | False |
-| learning_rate | 0.0010000000474974513 |
-| rho | 0.9 |
-| momentum | 0.0 |
-| epsilon | 1e-07 |
-| centered | False |
-| training_precision | float32 |
-