b5

app.py

@@ -9,8 +9,8 @@ import gradio as gr
 def seg(image):
     ## first resize the image !!
     image.resize((200,200))
-    feature_extractor = AutoFeatureExtractor.from_pretrained("nvidia/mit-b0")
-    model = SegformerForSemanticSegmentation.from_pretrained("nvidia/mit-b0")
+    feature_extractor = AutoFeatureExtractor.from_pretrained("nvidia/mit-b5")
+    model = SegformerForSemanticSegmentation.from_pretrained("nvidia/mit-b5")
     print(model)
     
     inputs = feature_extractor(images=image, return_tensors="pt")