adding second model to app

app.py

@@ -91,6 +91,24 @@ def process_image(session, img, colors, mosaic=False):
     return overlay
  
 
+st.title("OpenLander ONNX app")
+st.write("Upload an image to process with the ONNX OpenLander model!")
+st.write("Bear in mind that this model is **much less refined** than the embedded models at the moment.")
+
+models = {    
+        "Embedded model short trained: DeeplabV3+, MobilenetV2, 416px resolution": "end2end.onnx",
+        "Embedded model better trained: DeeplabV3+, MobilenetV2, 416px resolution": "20230608_onnx_416_mbnv2_dl3/end2end.onnx"
+        }
+
+
+
+
+
+# Create a Streamlit radio button to select the desired model
+selected_model = st.radio("Select a model", list(models.keys()))
+
+
+
 # set cuda = true if you have an NVIDIA GPU
 cuda = torch.cuda.is_available()
 
@@ -98,7 +116,10 @@ if cuda:
     print("We have a GPU!")
 providers = ['CUDAExecutionProvider'] if cuda else ['CPUExecutionProvider']
 
-session = ort.InferenceSession('end2end.onnx', providers=providers)
+# Get the selected model's path
+model_path = models[selected_model]
+
+session = ort.InferenceSession(model_path, providers=providers)
 
 
 # Define colors for classes 0, 122 and 244
@@ -113,11 +134,6 @@ def load_image(uploaded_file):
         return None
 
 
-st.title("OpenLander ONNX app")
-st.write("Upload an image to process with the ONNX OpenLander model!")
-st.write("Bear in mind that this model is **much less refined** than the embedded models at the moment.")
-
-
 uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png"])
 if uploaded_file is not None:
     img = load_image(uploaded_file)