Spaces:

5Grains
/

Week_10_Second

Sleeping

5Grains commited on Nov 12, 2023

Commit

7edebe4

1 Parent(s): e58d7a5

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -7,12 +7,11 @@ from PIL import Image
 import tensorflow as tf
 from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation
-feature_extractor = SegformerFeatureExtractor.from_pretrained(
-    "prem-timsina/segformer-b0-finetuned-food"
-)
-model = TFSegformerForSemanticSegmentation.from_pretrained(
-    "prem-timsina/segformer-b0-finetuned-food"
-)
 def ade_palette():
     """ADE20K palette that maps each class to RGB values."""

 import tensorflow as tf
 from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation
+# Load model directly
+from transformers import AutoFeatureExtractor, SegformerForSemanticSegmentation
+extractor = AutoFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
+model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
 def ade_palette():
     """ADE20K palette that maps each class to RGB values."""