Spaces:

Amrender
/

segformer_b5

Sleeping

Amrender commited on Mar 9

Commit

8d56fd5

verified ·

1 Parent(s): 62206ca

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -85,7 +85,7 @@ class UnifiedCartographer(nn.Module):
     def __init__(self, num_classes=5):
         super().__init__()
         self.model = SegformerForSemanticSegmentation.from_pretrained(
-            "nvidia/segformer-b5-finetuned-cityscapes-1024-1024",
             num_labels=num_classes, ignore_mismatched_sizes=True
         )
     def forward(self, x):

     def __init__(self, num_classes=5):
         super().__init__()
         self.model = SegformerForSemanticSegmentation.from_pretrained(
+            "nvidia/segformer-b5-finetuned-ade-640-640",
             num_labels=num_classes, ignore_mismatched_sizes=True
         )
     def forward(self, x):