5Grains commited on
Commit
8fdadcc
·
1 Parent(s): 297d3ef

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -8
app.py CHANGED
@@ -5,14 +5,11 @@ import matplotlib.pyplot as plt
5
  import numpy as np
6
  from PIL import Image
7
  import tensorflow as tf
8
- from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation
9
-
10
- feature_extractor = SegformerFeatureExtractor.from_pretrained(
11
- "nvidia/segformer-b0-finetuned-ade-512-512"
12
- )
13
- model = TFSegformerForSemanticSegmentation.from_pretrained(
14
- "nvidia/segformer-b0-finetuned-ade-512-512"
15
- )
16
 
17
  def ade_palette():
18
  """ADE20K palette that maps each class to RGB values."""
 
5
  import numpy as np
6
  from PIL import Image
7
  import tensorflow as tf
8
+ # Load model directly
9
+ from transformers import AutoFeatureExtractor, SegformerForSemanticSegmentation
10
+
11
+ extractor = AutoFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
12
+ model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
 
 
 
13
 
14
  def ade_palette():
15
  """ADE20K palette that maps each class to RGB values."""