Update README.md

README.md

@@ -10,12 +10,12 @@ pipeline_tag: image-classification
 
 [**MambaVision: A Hybrid Mamba-Transformer Vision Backbone**](https://arxiv.org/abs/2407.08083).
 
-### Model Overview
+## Model Overview
 
 We introduce a novel mixer block by creating a symmetric path without SSM to enhance the modeling of global context. MambaVision has a hierarchical architecture that employs both self-attention and mixer blocks. 
 
 
-### Model Performance
+## Model Performance
 
 MambaVision demonstrates a strong performance by achieving a new SOTA Pareto-front in
 terms of Top-1 accuracy and throughput. 
@@ -26,13 +26,7 @@ class="center">
 </p>
 
 
-### Model Usage
-
-You must first login into HuggingFace to pull the model:
-
-```Bash
-huggingface-cli login
-```
+## Model Usage
 
 It is highly recommended to install the requirements for MambaVision by running the following:
 
@@ -43,59 +37,82 @@ pip install mambavision
 
 For each model, we offer two variants for image classification and feature extraction that can be imported with 1 line of code. 
 
-The model can be simply imported according to:
+### Image Classification
 
-```Python
-from transformers import AutoModelForImageClassification
-model = AutoModelForImageClassification.from_pretrained("nvidia/MambaVision-T-1K", trust_remote_code=True)
-```
+In the following example, we demonstrate how MambaVision can be used for image classification. 
+
+Given the following image from [COCO dataset](https://cocodataset.org/#home)  val set as an input:
 
-The model outputs logits when an image is passed. If label is additionally provided, cross entropy loss between the output prediction and label is computed. 
 
-The following demonstrates a minimal example of how to use the model:
+<p align="center">
+<img src="https://cdn-uploads.huggingface.co/production/uploads/64414b62603214724ebd2636/4duSnqLf4lrNiAHczSmAN.jpeg" width=42% height=42% 
+class="center">
+</p>
+
 
+The following snippet can be used for image classification:
 
 ```Python
 from transformers import AutoModelForImageClassification
 from PIL import Image
+from timm.data.transforms_factory import create_transform
 import requests
-import torch
-import timm
-
-# import mambavision model
 
 model = AutoModelForImageClassification.from_pretrained("nvidia/MambaVision-T-1K", trust_remote_code=True)
 
 # eval mode for inference
-model.eval()
+model.cuda().eval()
 
 # prepare image for the model
-url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
+url = 'http://images.cocodataset.org/val2017/000000020247.jpg'
 image = Image.open(requests.get(url, stream=True).raw)
+input_resolution = (3, 224, 224)  # MambaVision supports any input resolutions but has been originally trained on (3, 224, 224)
 
-# define a transform
+transform = create_transform(input_size=input_resolution, is_training=False, mean=model.config.mean, std=model.config.std, crop_mode=model.config.crop_mode, crop_pct=model.config.crop_pct)
 
-transforms = timm.data.create_transform((3, 224, 224))
+inputs = transform(image).unsqueeze(0).cuda()
+# model inference
+outputs = model(inputs)
+logits = outputs['logits'] 
+predicted_class_idx = logits.argmax(-1).item()
+print("Predicted class:", model.config.id2label[predicted_class_idx])
+```
 
-image = transforms(image).unsqueeze(0)
+The predicted label is brown bear, bruin, Ursus arctos.
 
-# put both model and image on cuda
+### Feature Extraction
 
-model = model.cuda()
-image = image.cuda()
+MambaVision can also be used as a generic feature extractor. 
 
-# forward pass
-outputs = model(image)
+Specifically, we can extract the outputs of each stage of model (4 stages) as well as the final averaged-pool features that are flattened. 
 
-# You can then extract the predicted probabilities by applying softmax: 
+The following snippet can be used for feature extraction:
 
-probabilities = torch.nn.functional.softmax(outputs['logits'], dim=0)
+```Python
+from transformers import AutoModel
+from PIL import Image
+from timm.data.transforms_factory import create_transform
+import requests
+
+model = AutoModel.from_pretrained("nvidia/MambaVision-T-1K", trust_remote_code=True)
 
-# In order to find the top 5 predicted class indexes and their corresponding values:
+# eval mode for inference
+model.cuda().eval()
 
-values, indices = torch.topk(probabilities, 5)
+# prepare image for the model
+url = 'http://images.cocodataset.org/val2017/000000020247.jpg'
+image = Image.open(requests.get(url, stream=True).raw)
+input_resolution = (3, 224, 224)  # MambaVision supports any input resolutions but has been originally trained on (3, 224, 224)
 
+transform = create_transform(input_size=input_resolution, is_training=False, mean=model.config.mean, std=model.config.std, crop_mode=model.config.crop_mode, crop_pct=model.config.crop_pct)
 
+inputs = transform(image).unsqueeze(0).cuda()
+# model inference
+out_avg_pool, features = model(inputs)
+print("Size of the averaged pool features:", out_avg_pool.size())  # torch.Size([1, 640])
+print("Number of stages in extracted features:", len(features)) # 4 stages
+print("Size of extracted features in stage 1:", features[0].size()) # torch.Size([1, 80, 56, 56])
+print("Size of extracted features in stage 4:", features[3].size()) # torch.Size([1, 640, 7, 7])
 ```