Upload app.py

app.py

@@ -24,12 +24,18 @@ def load_model(model_name):
             model = models.segmentation.deeplabv3_resnet101(weights=weights)
         except:
             model = models.segmentation.deeplabv3_resnet101(pretrained=True)
-    elif model_name == "FCN (ResNet50)":
+    elif model_name == "DeepLabV3+ (MobileNetV3)":
         try:
-            weights = models.segmentation.FCN_ResNet50_Weights.DEFAULT
-            model = models.segmentation.fcn_resnet50(weights=weights)
+            weights = models.segmentation.DeepLabV3_MobileNet_V3_Large_Weights.DEFAULT
+            model = models.segmentation.deeplabv3_mobilenet_v3_large(weights=weights)
         except:
-            model = models.segmentation.fcn_resnet50(pretrained=True)
+            model = models.segmentation.deeplabv3_mobilenet_v3_large(pretrained=True)
+    elif model_name == "LRASPP (MobileNetV3)":
+        try:
+            weights = models.segmentation.LRASPP_MobileNet_V3_Large_Weights.DEFAULT
+            model = models.segmentation.lraspp_mobilenet_v3_large(weights=weights)
+        except:
+            model = models.segmentation.lraspp_mobilenet_v3_large(pretrained=True)
     
     model.eval()
     if torch.cuda.is_available():
@@ -130,7 +136,7 @@ article = """
 # Example images
 examples = [
     ["demo1.jpg", "DeepLabV3+ (ResNet50)"],
-    ["demo2.png", "FCN (ResNet50)"]
+    ["demo2.png", "LRASPP (MobileNetV3)"]
 ]
 
 demo = gr.Interface(
@@ -138,7 +144,7 @@ demo = gr.Interface(
     inputs=[
         gr.Image(type="pil", label="Input Image"),
         gr.Dropdown(
-            choices=["DeepLabV3+ (ResNet50)", "DeepLabV3+ (ResNet101)", "FCN (ResNet50)"], 
+            choices=["DeepLabV3+ (ResNet50)", "DeepLabV3+ (ResNet101)", "DeepLabV3+ (MobileNetV3)", "LRASPP (MobileNetV3)"], 
             value="DeepLabV3+ (ResNet50)", 
             label="Select Pre-trained Model"
         )