| 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) |
