Spaces:
Running on Zero
Running on Zero
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import torchvision.models as models | |
| import torchvision.transforms as T | |
| import gradio as gr | |
| from PIL import Image, ImageOps | |
| import os | |
| # On ZeroGPU Spaces the `spaces` package exists; elsewhere fall back to a no-op decorator | |
| try: | |
| import spaces | |
| except ImportError: # not on a ZeroGPU Space: provide a do-nothing stand-in | |
| class spaces: | |
| def GPU(fn): | |
| return fn | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| # Must match the values used in training | |
| EMBED_DIM, ATTENTION_DIM, DECODER_DIM = 256, 256, 512 | |
| CKPT_PATH = "best_model.pt" | |
| class EncoderCNN(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| # weights=None: the trained ResNet weights come from the checkpoint | |
| resnet = models.resnet50(weights=None) | |
| self.resnet = nn.Sequential(*list(resnet.children())[:-2]) | |
| def forward(self, images): | |
| features = self.resnet(images) | |
| b, c, h, w = features.shape | |
| return features.view(b, c, h * w).permute(0, 2, 1) | |
| class BahdanauAttention(nn.Module): | |
| def __init__(self, encoder_dim, decoder_dim, attention_dim): | |
| super().__init__() | |
| self.W_encoder = nn.Linear(encoder_dim, attention_dim) | |
| self.W_decoder = nn.Linear(decoder_dim, attention_dim) | |
| self.V = nn.Linear(attention_dim, 1) | |
| def forward(self, encoder_out, decoder_hidden): | |
| wh = self.W_encoder(encoder_out) | |
| ws = self.W_decoder(decoder_hidden).unsqueeze(1) | |
| scores = self.V(torch.tanh(wh + ws)).squeeze(2) | |
| weights = F.softmax(scores, dim=1) | |
| context = (encoder_out * weights.unsqueeze(2)).sum(dim=1) | |
| return context, weights | |
| class DecoderWithAttention(nn.Module): | |
| def __init__(self, attention_dim, embed_dim, decoder_dim, vocab_size, encoder_dim=2048): | |
| super().__init__() | |
| self.vocab_size = vocab_size | |
| self.attention = BahdanauAttention(encoder_dim, decoder_dim, attention_dim) | |
| self.embedding = nn.Embedding(vocab_size, embed_dim) | |
| self.dropout = nn.Dropout(0.5) | |
| self.lstm_cell = nn.LSTMCell(embed_dim + encoder_dim, decoder_dim) | |
| self.fc = nn.Linear(decoder_dim, vocab_size) | |
| class ImageCaptioner(nn.Module): | |
| def __init__(self, attention_dim, embed_dim, decoder_dim, vocab_size): | |
| super().__init__() | |
| self.encoder = EncoderCNN() | |
| self.decoder = DecoderWithAttention(attention_dim, embed_dim, decoder_dim, vocab_size) | |
| self.init_h = nn.Linear(2048, decoder_dim) | |
| self.init_c = nn.Linear(2048, decoder_dim) | |
| def init_hidden_state(self, encoder_out): | |
| mean = encoder_out.mean(dim=1) | |
| return torch.tanh(self.init_h(mean)), torch.tanh(self.init_c(mean)) | |
| # ---- load checkpoint ---- | |
| if not os.path.exists(CKPT_PATH): # e.g. running from the GitHub repo without the weights | |
| from huggingface_hub import hf_hub_download | |
| CKPT_PATH = hf_hub_download(repo_id="birukale/Image-Caption", filename="best_model.pt", repo_type="space") | |
| ckpt = torch.load(CKPT_PATH, map_location="cpu") | |
| itos = ckpt["itos"] | |
| stoi = {w: i for i, w in itos.items()} | |
| model = ImageCaptioner(ATTENTION_DIM, EMBED_DIM, DECODER_DIM, len(itos)).to(device) | |
| model.load_state_dict(ckpt["model"]) | |
| model.eval() | |
| transform = T.Compose([ | |
| T.Resize((224, 224)), T.ToTensor(), | |
| T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), | |
| ]) | |
| def caption_image(image, max_len=30): | |
| if image is None: | |
| return "Please upload an image." | |
| image = ImageOps.exif_transpose(image).convert("RGB") | |
| img = transform(image).unsqueeze(0).to(device) | |
| enc = model.encoder(img) | |
| h, c = model.init_hidden_state(enc) | |
| word = torch.tensor([stoi["<start>"]], device=device) | |
| out = [] | |
| for _ in range(max_len): | |
| emb = model.decoder.embedding(word) | |
| context, _ = model.decoder.attention(enc, h) | |
| h, c = model.decoder.lstm_cell(torch.cat([emb, context], dim=1), (h, c)) | |
| word = model.decoder.fc(h).argmax(dim=1) | |
| if word.item() == stoi["<end>"]: | |
| break | |
| out.append(itos[word.item()]) | |
| return " ".join(out) | |
| demo = gr.Interface( | |
| fn=caption_image, | |
| inputs=gr.Image(type="pil", label="Upload an image"), | |
| outputs=gr.Textbox(label="Generated caption"), | |
| title="Image Captioning: ResNet-50 + Bahdanau Attention + LSTM", | |
| description="Trained from scratch on Flickr8k. Works best on everyday scenes: people, dogs, sports, streets, outdoors.", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |