nvidia

/

mit-b0

like 33

Follow

NVIDIA 18.5k

@amyeroberts

I've tried with TF weight, but got error.

from transformers import SegformerFeatureExtractor, SegformerForImageClassification
from PIL import Image
import requests

url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)

feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/mit-b0")
model = SegformerForImageClassification.from_pretrained("nvidia/mit-b0")

inputs = feature_extractor(images=image, return_tensors="tf")
outputs = model(**inputs)
logits = outputs.logits
# model predicts one of the 1000 ImageNet classes
predicted_class_idx = logits.argmax(-1).item()
print("Predicted class:", model.config.id2label[predicted_class_idx])

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/tmp/ipykernel_33/1248207528.py in <module>
     10 
     11 inputs = feature_extractor(images=image, return_tensors="tf")
---> 12 outputs = model(**inputs)
     13 logits = outputs.logits
     14 # model predicts one of the 1000 ImageNet classe

TypeError: conv2d() received an invalid combination of arguments - got (tensorflow.python.framework.ops.EagerTensor, Parameter, Parameter, tuple, tuple, tuple, int), but expected one of:
 * (Tensor input, Tensor weight, Tensor bias, tuple of ints stride, tuple of ints padding, tuple of ints dilation, int groups)
      didn't match because some of the arguments have invalid types: (!tensorflow.python.framework.ops.EagerTensor!, !Parameter!, !Parameter!, !tuple!, !tuple!, !tuple!, int)
 * (Tensor input, Tensor weight, Tensor bias, tuple of ints stride, str padding, tuple of ints dilation, int groups)
      didn't match because some of the arguments have invalid types: (!tensorflow.python.framework.ops.EagerTensor!, !Parameter!, !Parameter!, !tuple!, !tuple!, !tuple!, int)

This issue is arising because SegformerForImageClassification is a PyTorch model and the inputs are tensorflow tensors. To load the TF model, you want to import TFSegformerForImageClassification.

uh, got it. Thanks.