File size: 1,823 Bytes
ea52636 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | readme_content = """# Image Captioning Model
This is an Image Captioning model trained on a custom dataset with 6000 training images and 1000 validation images. It uses ResNet50 for image feature extraction and an LSTM with fine-tuned GloVe embeddings for caption generation.
## Model Details
- **Image Encoder**: ResNet50
- **Text Decoder**: LSTM with fine-tuned GloVe embeddings
- **Vocabulary Size**: 3478
- **Embedding Dimension**: 512
- **Hidden Size**: 1024
- **Number of Layers**: 2
- **Max Caption Length**: 30
## Usage
Load the model and tokenizer using PyTorch and the provided `model.py`.
### Example
```python
import torch
from PIL import Image
import pickle
from huggingface_hub import hf_hub_download
from model import image_captioning_model
# Download files
repo_id = "Pumponchik/image_caption_weak"
model_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin")
config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
tokenizer_path = hf_hub_download(repo_id=repo_id, filename="tokenizer.pkl")
transform_path = hf_hub_download(repo_id=repo_id, filename="image_transforms.pkl")
# Load config
with open(config_path, "r") as f:
config = json.load(f)
# Load model
model = image_captioning_model(**config)
model.load_state_dict(torch.load(model_path))
model.eval()
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
# Load tokenizer and transforms
with open(tokenizer_path, "rb") as f:
tokenizer = pickle.load(f)
with open(transform_path, "rb") as f:
transform = pickle.load(f)
# Generate caption
image = Image.open("example.png").convert("RGB")
image_tensor = transform(image).unsqueeze(0).to(device)
caption = model.generate_caption(image_tensor, tokenizer, max_length=config["max_len"], device=device)
print("Generated Caption:", caption) |