---
language: en
license: mit
tags:
  - vision
  - image-captioning
  - image-classification
  - bird-species
datasets:
  - cub-200-2011
---

# Bird Captioning and Classification Model (CUB-200-2011)

This is a fine-tuned VisionEncoderDecoderModel based on `nlpconnect/vit-gpt2-image-captioning`, trained on the CUB-200-2011 dataset for bird species classification and image captioning.

## Model Description
- **Base Model**: ViT-GPT2 (`nlpconnect/vit-gpt2-image-captioning`)
- **Tasks**:
  - Generates descriptive captions for bird images, including species and attributes.
  - Classifies images into one of 200 bird species.
- **Dataset**: CUB-200-2011 (11,788 images, 200 bird species)
- **Training**: 10 epochs, batch size 16, mixed precision, AdamW optimizer (lr=3e-5), combined loss (caption + 0.5 * classification).
- **Best Validation Loss**: 0.0690 (Epoch 3)

## Files
- `model.safetensors`: Trained model weights
- `config.json`: Model configuration
- `preprocessor_config.json`: ViTImageProcessor settings
- `tokenizer_config.json`, `vocab.json`: GPT2 tokenizer files
- `species_mapping.txt`: Mapping of class indices to bird species names
- `cub200_captions.csv`: Generated captions for the dataset
- `model.py`: Custom `BirdCaptioningModel` class definition

## Usage
### Prerequisites
```bash
pip install transformers torch huggingface_hub
```

### Load Model and Dependencies
```python
from transformers import ViTImageProcessor, AutoTokenizer
from huggingface_hub import PyTorchModelHubMixin
import torch
from model import BirdCaptioningModel  # Save model.py locally

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load model
model = BirdCaptioningModel.from_pretrained("INVERTO/bird-captioning-cub200").to(device)
image_processor = ViTImageProcessor.from_pretrained("INVERTO/bird-captioning-cub200")
tokenizer = AutoTokenizer.from_pretrained("INVERTO/bird-captioning-cub200")
model.eval()

# Load species mapping
species_mapping = {}
with open("species_mapping.txt", "r") as f:
    for line in f:
        idx, name = line.strip().split(",", 1)
        species_mapping[int(idx)] = name
```

### Inference
```python
from PIL import Image

def predict_bird_image(image_path):
    image = Image.open(image_path).convert("RGB")
    pixel_values = image_processor(image, return_tensors="pt").pixel_values.to(device)
    with torch.no_grad():
        output_ids = model.base_model.generate(pixel_values, max_length=75, num_beams=4)
        _, class_logits = model(pixel_values)
        predicted_class_idx = torch.argmax(class_logits, dim=1).item()
        confidence = torch.nn.functional.softmax(class_logits, dim=1)[0, predicted_class_idx].item() * 100
        caption = tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
        species = species_mapping.get(predicted_class_idx, "Unknown")
    return caption, species, confidence

# Example
caption, species, confidence = predict_bird_image("/kaggle/input/cub2002011/CUB_200_2011/images/006.Least_Auklet/Least_Auklet_0007_795123.jpg")
print(f"Caption: {caption}")
print(f"Species: {species}")
print(f"Confidence: {confidence:.2f}%")
```

## Dataset
- **CUB-200-2011**: 11,788 images of 200 bird species with attribute annotations.
- Captions were generated based on species names and attributes (e.g., bill shape, wing color).

## Training Details
- **Loss**: Combined captioning (CrossEntropy) and classification (CrossEntropy) loss.
- **Optimizer**: AdamW (lr=3e-5)
- **Scheduler**: CosineAnnealingLR
- **Hardware**: GPU (CUDA)
- **Training Time**: ~5 min/epoch

## Limitations
- May overfit after Epoch 3 (validation loss increases).
- Captions are limited to species and up to 5 attributes.
- Classification accuracy not explicitly reported.

## License
MIT License

## Contact
For issues, contact INVERTO on Hugging Face.