Upload README.md with huggingface_hub

README.md

@@ -20,10 +20,53 @@ Table of Contents:
 - [Dataset](#dataset)
 
 ## Load trained model
+
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/qubvel/segmentation_models.pytorch/blob/main/examples/segformer_inference_pretrained.ipynb)
+
+1. Install requirements.
+
+```bash
+pip install -U segmentation_models_pytorch albumentations
+```
+
+2. Run inference.
+
 ```python
+import torch
+import requests
+import numpy as np
+import albumentations as A
 import segmentation_models_pytorch as smp
 
-model = smp.from_pretrained("smp-hub/segformer-b1-1024x1024-city-160k")
+from PIL import Image
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load pretrained model and preprocessing function
+checkpoint = "smp-hub/segformer-b1-1024x1024-city-160k"
+model = smp.from_pretrained(checkpoint).eval().to(device)
+preprocessing = A.Compose.from_pretrained(checkpoint)
+
+# Load image
+url = "https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg"
+image = Image.open(requests.get(url, stream=True).raw)
+
+# Preprocess image
+np_image = np.array(image)
+normalized_image = preprocessing(image=np_image)["image"]
+input_tensor = torch.as_tensor(normalized_image)
+input_tensor = input_tensor.permute(2, 0, 1).unsqueeze(0)  # HWC -> BCHW
+input_tensor = input_tensor.to(device)
+
+# Perform inference
+with torch.no_grad():
+    output_mask = model(input_tensor)
+
+# Postprocess mask
+mask = torch.nn.functional.interpolate(
+    output_mask, size=(image.height, image.width), mode="bilinear", align_corners=False
+)
+mask = mask.argmax(1).cpu().numpy()  # argmax over predicted classes (channels dim)
 ```
 
 ## Model init parameters
@@ -40,11 +83,8 @@ model_init_params = {
 }
 ```
 
-## Model metrics
-[More Information Needed]
-
 ## Dataset
-Dataset name: [More Information Needed]
+Dataset name: [Cityscapes](https://paperswithcode.com/dataset/cityscapes)
 
 ## More Information
 - Library: https://github.com/qubvel/segmentation_models.pytorch