changed pre-trained

app.py

@@ -11,7 +11,7 @@ def main():
     st.write("")
 
     selected_option = st.selectbox(
-        "Number of training samples", [10, 100, 500, 1000, 5000]
+        "Number of training samples", [10, 1000, 5000]
     )
     st.markdown(
         "Move slider to select how many training "
@@ -24,22 +24,23 @@ def main():
 
     preds = load_predictions(selected_option, Path("./images/predictions"))
 
-    zipped_st_images = zip(images, preds["svb"], preds["unet"], labels)
+    zipped_st_images = zip(images, preds["svb"], preds["unet"], preds["unet-ls"], labels)
 
     st.write("")
 
-    titleCol0, titleCol1, titleCol2, titleCol3 = st.columns(4)
+    titleCol0, titleCol1, titleCol2, titleCol3, titleCol4 = st.columns(5)
 
     titleCol0.markdown(f"### MOD09GA [3-2-1] Image Chip")
     titleCol1.markdown(f"### SatVision-B Prediction")
     titleCol2.markdown(f"### UNet (CNN) Prediction")
-    titleCol3.markdown(f"### MCD12Q1 LandCover Target")
+    titleCol3.markdown(f'### UNet (CNN) LS Pretrained Prediction')
+    titleCol4.markdown(f"### MCD12Q1 LandCover Target")
 
     st.write("")
 
-    grid = make_grid(4, 4)
+    grid = make_grid(5, 5)
 
-    for i, (image_data, svb_data, unet_data, label_data) in enumerate(zipped_st_images):
+    for i, (image_data, svb_data, unet_data, unet_ls_data, label_data) in enumerate(zipped_st_images):
         # if i == 0:
 
         #    grid[0][0].markdown(f'## MOD09GA 3-2-1 Image Chip')
@@ -50,7 +51,8 @@ def main():
         grid[i][0].image(image_data[0], image_data[1], use_column_width=True)
         grid[i][1].image(svb_data[0], svb_data[1], use_column_width=True)
         grid[i][2].image(unet_data[0], unet_data[1], use_column_width=True)
-        grid[i][3].image(label_data[0], label_data[1], use_column_width=True)
+        grid[i][3].image(unet_ls_data[0], unet_ls_data[1], use_column_width=True)
+        grid[i][4].image(label_data[0], label_data[1], use_column_width=True)
 
     st.markdown("### Few-Shot Learning with SatVision-Base")
     description = (
@@ -135,6 +137,7 @@ def load_predictions(selected_option: str, pred_dir: Path):
     svb_pred_paths = find_preds(selected_option, pred_dir, "svb")
 
     unet_pred_paths = find_preds(selected_option, pred_dir, "cnn")
+    unet_ls_pred_paths = find_preds(selected_option, pred_dir, "cnn-ls")
 
     svb_preds = [
         (str(path), f"SatVision-B Prediction Example {i}")
@@ -146,7 +149,12 @@ def load_predictions(selected_option: str, pred_dir: Path):
         for i, path in enumerate(unet_pred_paths, 1)
     ]
 
-    prediction_dict = {"svb": svb_preds, "unet": unet_preds}
+    unet_ls_preds = [
+        (str(path), f"Unet LS Pre-trained Prediction Example {i}")
+        for i, path in enumerate(unet_ls_pred_paths, 1)
+    ]
+
+    prediction_dict = {"svb": svb_preds, "unet": unet_preds, "unet-ls": unet_ls_preds}
 
     return prediction_dict
 
@@ -155,8 +163,12 @@ def load_predictions(selected_option: str, pred_dir: Path):
 # find_preds
 # -----------------------------------------------------------------------------
 def find_preds(selected_option: int, pred_dir: Path, model: str):
+
     if model == "cnn":
-        pred_regex = f"ft_cnn_demo_{selected_option}_*_pred.png"
+        pred_regex = f"ft_cnn_demo_{selected_option}_*cnn-plain_pred.png"
+
+    elif model == "cnn-ls":
+        pred_regex = f"ft_cnn_demo_{selected_option}_*cnn-ls_pred.png"
 
     else:
         pred_regex = f"ft_demo_{selected_option}_*_pred.png"
@@ -168,7 +180,7 @@ def find_preds(selected_option: int, pred_dir: Path, model: str):
     preds_matching_regex = sorted(model_specific_dir.glob(pred_regex))
 
     assert (
-        len(preds_matching_regex) == 3
+        len(preds_matching_regex) == 3 
     ), "Should be 3 prediction images matching regex"
 
     assert "1071" in str(preds_matching_regex[0]), "Should be 1071"