Update app.py

app.py

@@ -115,13 +115,16 @@ def predict(input_image, model_choice):
 # Initialize model (do this outside the inference function for better performance)
 w_noise_model_path = "./models/best_model_w_noise.pth"
 wo_noise_model_path = "./models/best_model_wo_noise.pth"
+w_noise_model_v2_path = "./models/best_model_w_noise_v2.pth"
 
 w_noise_model, device = initialize_model(w_noise_model_path)
 wo_noise_model, device = initialize_model(wo_noise_model_path)
+w_noise_model_v2, device = initialize_model(w_noise_model_v2_path)
 
 models = {
     "Without Noise": wo_noise_model,
     "With Noise": w_noise_model
+    "With Noise V2": w_noise_model_v2
 }
 
 # Create Gradio interface
@@ -129,7 +132,7 @@ iface = gr.Interface(
     fn=predict,
     inputs=[
         gr.Image(type="pil"),
-        gr.Dropdown(choices=["Without Noise", "With Noise"], value="With Noise"),
+        gr.Dropdown(choices=["Without Noise", "With Noise", "With Noise V2"], value="With Noise V2"),
     ],
     outputs=[
         gr.Image(type="pil", label="Segmentation Mask"),