Update README.md

README.md

@@ -16,7 +16,29 @@ SegFormer model fine-tuned on SegmentsAI [`sidewalk-semantic`](https://huggingfa
 ## Model description
 SegFormer consists of a hierarchical Transformer encoder and a lightweight all-MLP decode head to achieve great results on semantic segmentation benchmarks such as ADE20K and Cityscapes. The hierarchical Transformer is first pre-trained on ImageNet-1k, after which a decode head is added and fine-tuned altogether on a downstream dataset.
 
-## Notebook and Code 
+## Code and Notebook 
+Here is how to use this model to classify an image of the sidewalk dataset:
+
+```python
+from transformers import SegformerFeatureExtractor, SegformerForImageClassification
+from PIL import Image
+import requests
+
+url = "https://segmentsai-prod.s3.eu-west-2.amazonaws.com/assets/admin-tobias/439f6843-80c5-47ce-9b17-0b2a1d54dbeb.jpg"
+image = Image.open(requests.get(url, stream=True).raw)
+
+feature_extractor = SegformerFeatureExtractor.from_pretrained("zoheb/mit-b5-finetuned-sidewalk-semantic")
+model = SegformerForImageClassification.from_pretrained("zoheb/mit-b5-finetuned-sidewalk-semantic")
+
+inputs = feature_extractor(images=image, return_tensors="pt")
+outputs = model(**inputs)
+logits = outputs.logits
+
+# model predicts one of the 35 Sidewalk classes
+predicted_class_idx = logits.argmax(-1).item()
+print("Predicted class:", model.config.id2label[predicted_class_idx])
+```
+
 You can go through its detailed notebook [here](https://github.com/ZohebAbai/Deep-Learning-Projects/blob/master/09_HF_Image_Segmentation_using_Transformers.ipynb).
 
 For more code examples, refer to the [documentation](https://huggingface.co/transformers/model_doc/segformer.html#).