fix water

biomap/inference.py

@@ -13,6 +13,7 @@ preprocess = T.Compose(
         ]
     )
 
+import numpy as np
 def inference(images, model):
     logging.info("Inference on Images")
     x = torch.stack([preprocess(image) for image in images]).cpu()
@@ -25,6 +26,10 @@ def inference(images, model):
             "img": x[i].detach().cpu(),
             "linear_preds": linear_pred[i].detach().cpu(),
         } for i in range(x.shape[0])]
+    
+    # water to natural green
+    for output in outputs:
+        output["linear_preds"] = torch.where(output["linear_preds"] == 5,  3, output["linear_preds"])
     return outputs
 
 
@@ -32,6 +37,7 @@ if __name__ == "__main__":
     import hydra
     from model import LitUnsupervisedSegmenter
     from utils_gee import extract_img, transform_ee_img
+    import os
     latitude = 2.98
     longitude = 48.81
     start_date = '2020-03-20'
@@ -49,7 +55,8 @@ if __name__ == "__main__":
     cfg = hydra.compose(config_name="my_train_config.yml")
 
     # Load the model
-    model_path = "checkpoint/model/model.pt"
+    
+    model_path = os.path.join(os.path.dirname(__file__), "checkpoint/model/model.pt")
     saved_state_dict = torch.load(model_path, map_location=torch.device("cpu"))
 
     nbclasses = cfg.dir_dataset_n_classes

biomap/streamlit_app.py

@@ -64,6 +64,7 @@ def app(model):
     st.markdown("<p style='text-align: center;'>The segmentation model is an association of UNet and DinoV1 trained on the dataset CORINE. Land use is divided into 6 differents classes : Each class is assigned a GBS score from 0 to 1</p>", unsafe_allow_html=True)
     st.markdown("<p style='text-align: center;'>Buildings : 0.1 | Infrastructure : 0.1 | Cultivation : 0.4 | Wetland : 0.9 | Water : 0.9 | Natural green : 1 </p>", unsafe_allow_html=True)
     st.markdown("<p style='text-align: center;'>The score is then averaged on the full image.</p>", unsafe_allow_html=True)
+    
     if st.session_state["submit"]:
             fig = inference_on_location(model, st.session_state["lat"], st.session_state["long"], st.session_state["start_date"], st.session_state["end_date"], st.session_state["segment_interval"])
             st.session_state["infered"] = True
@@ -76,10 +77,6 @@ def app(model):
 
     if st.session_state["infered"]:
             st.plotly_chart(st.session_state["previous_fig"], use_container_width=True)
-
-    
-    
-    
         
     col_1, col_2 = st.columns([0.5, 0.5])
     with col_1: