tf

app.py

@@ -5,12 +5,12 @@ import matplotlib.pyplot as plt
 import numpy as np
 from PIL import Image
 import tensorflow as tf
-from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation
+from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation
 
 feature_extractor = SegformerFeatureExtractor.from_pretrained(
     "nvidia/segformer-b2-finetuned-cityscapes-1024-1024"
 )
-model = SegformerForSemanticSegmentation.from_pretrained(
+model = TFSegformerForSemanticSegmentation.from_pretrained(
     "nvidia/segformer-b2-finetuned-cityscapes-1024-1024"
 )
 
@@ -101,12 +101,6 @@ def sepia(input_img):
     fig = draw_plot(pred_img, seg)
     return fig
 
-#def prepro_img(img):
-     #print("Preprocessing image...")
-     #img = img.resize((1024,1024))
-     #print("Image preprocessing completed.")
-     #return img
-
 demo = gr.Interface(fn=sepia,
                     inputs=gr.Image(shape=(400, 600)),
                     #inputs=gr.Image(type='pil', prepeocessing_function=prepro_img),