Update README.md

README.md

@@ -1,10 +1,43 @@
 ---
+license: mit
+datasets:
+- imageomics/fish-vista
 tags:
-- model_hub_mixin
-- pytorch_model_hub_mixin
+- swin
+- swin-transformer
+- fish
+- fish-classification
 ---
 
-This model has been pushed to the Hub using the [PytorchModelHubMixin](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin) integration:
-- Code: [More Information Needed]
-- Paper: [More Information Needed]
-- Docs: [More Information Needed]
+# Swin Transformer for Fish Classification
+
+This swin-tiny transformer was trained to classify fish species from the fish-visa dataset. The dataset contains massive imbalance and has varying
+degrees of accuracy depending on the class frequency. I recommend to use this model weights as a baseline for building a classification or segmentation
+model.
+
+## Usage
+Load and use the model :
+
+```python
+import torch
+from PIL import Image
+from transformers import AutoImageProcessor
+from transformers import AutoModel
+
+# 1. Load the model and processor from the Hub
+repo_id = "Aelyos/swin-fish-classification"
+
+# This downloads the code from the Hub and initializes the classes automatically
+processor = AutoImageProcessor.from_pretrained(repo_id, trust_remote_code=True)
+model = AutoModel.from_pretrained(repo_id, trust_remote_code=True, id2label=id2label, label2id=label2id)
+```
+
+## Training
+The model was trained on 80 epochs on a A10 GPU, first by training the classifier head only for 30 epochs and then by training the entire model for 50 more epochs.
+
+![loss_curve](https://cdn-uploads.huggingface.co/production/uploads/6661f6eeb8e23e099edebafb/N3vKSOhykFIy2cMr6oD-F.png)
+
+## Results
+Following the original fish-vista paper, I computed the accuracy accross bins (ultra rare 2-10, minority 10-100, neutral 100-500, majority 500+).
+
+![image](https://cdn-uploads.huggingface.co/production/uploads/6661f6eeb8e23e099edebafb/ItwMxmZz2J6HhKr7cNLt8.png)