Upload 3 files

README.md

@@ -1,14 +1,53 @@
----
-title: ImageCaption
-emoji: 📉
-colorFrom: yellow
-colorTo: pink
-sdk: gradio
-sdk_version: 5.29.1
-app_file: app.py
-pinned: false
-license: mit
-short_description: ImageCaption weak
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+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)

app.py

@@ -0,0 +1,131 @@
+import torch
+import torch.nn.functional as F
+import gradio as gr
+from PIL import Image
+import pickle
+import json
+from huggingface_hub import hf_hub_download
+from model import image_captioning_model
+
+# Загрузка файлов
+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")
+
+# Загрузка конфигурации
+with open(config_path, "r") as f:
+    config = json.load(f)
+
+# Загрузка модели
+model = image_captioning_model(**config)
+model.load_state_dict(torch.load(model_path, map_location="cpu"))
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+model.eval()
+
+# Загрузка токенизатора и трансформаций
+with open(tokenizer_path, "rb") as f:
+    tokenizer = pickle.load(f)
+with open(transform_path, "rb") as f:
+    transform = pickle.load(f)
+
+title_width = 43
+
+def wrap_text(text, max_width):
+    words = text.split(" ")
+    result = [[words[0]]]
+    for word in words[1:]:
+        if len(" ".join(result[-1])) + len(word) + 1 > max_width:
+            result[-1] = " ".join(result[-1])
+            result.append([])
+        result[-1].append(word)
+    result[-1] = " ".join(result[-1])
+    return "\n".join(result)
+
+def generate(model, image, tokenizer, max_seq_len=20, top_p=None, top_k=1):
+    assert top_p is None or top_k is None, "Не используйте одновременно top_p и top_k"
+
+    model.eval()
+    device = next(model.parameters()).device
+    image = image.to(device)
+    batch_size = image.size(0)
+
+    with torch.no_grad():
+        img_emb = model.img_features(image)  # [batch_size, 512]
+        h_0 = model.text_features.img_to_hidden(img_emb)  # [batch_size, hidden_size]
+        h_0 = h_0.unsqueeze(0).repeat(model.text_features.num_layers, 1, 1)  # [num_layers, batch_size, hidden_size]
+        c_0 = h_0.clone()
+
+        input_token = torch.full((batch_size,), tokenizer.tok_to_ind['<BOS>'], dtype=torch.long, device=device)
+        result_tokens = [[] for _ in range(batch_size)]
+        mask = torch.ones(batch_size, dtype=torch.bool, device=device)
+
+        for _ in range(max_seq_len):
+            input_emb = model.text_features.embed(input_token).unsqueeze(1)  # [batch_size, 1, embed_dim]
+            input_emb = model.text_features.text_to_hidden(input_emb)
+            output, (h_0, c_0) = model.text_features.rnn(input_emb, (h_0, c_0))  # [batch_size, 1, hidden_size]
+            logits = model.fc_out(output.squeeze(1))  # [batch_size, vocab_size]
+            probs = F.softmax(logits, dim=-1)  # [batch_size, vocab_size]
+
+            next_tokens = torch.zeros(batch_size, dtype=torch.long, device=device)
+            for i in range(batch_size):
+                if not mask[i]:
+                    continue
+                prob = probs[i]
+                if top_k is not None:
+                    sorted_probs, sorted_indices = torch.sort(prob, descending=True)
+                    sorted_probs = sorted_probs[:top_k]
+                    sorted_indices = sorted_indices[:top_k]
+                    sorted_probs = sorted_probs / sorted_probs.sum()
+                    next_token_idx = torch.multinomial(sorted_probs, 1)
+                    next_tokens[i] = sorted_indices[next_token_idx]
+                elif top_p is not None:
+                    sorted_probs, sorted_indices = torch.sort(prob, descending=True)
+                    cumulative_probs = torch.cumsum(sorted_probs, dim=0)
+                    mask_p = cumulative_probs <= top_p
+                    sorted_probs = sorted_probs[mask_p]
+                    sorted_indices = sorted_indices[mask_p]
+                    sorted_probs = sorted_probs / sorted_probs.sum()
+                    next_token_idx = torch.multinomial(sorted_probs, 1)
+                    next_tokens[i] = sorted_indices[next_token_idx]
+                else:
+                    next_tokens[i] = torch.multinomial(prob, 1)
+
+                token_id = next_tokens[i].item()
+                result_tokens[i].append(token_id)
+                if token_id == tokenizer.tok_to_ind['<EOS>']:
+                    mask[i] = False
+
+            input_token = next_tokens
+            if not mask.any():
+                break
+
+        # Декодирование токенов в текст
+        result_text = []
+        for tokens in result_tokens:
+            caption = " ".join([tokenizer.ind_to_tok.get(t, "<UNK>") for t in tokens 
+                               if t not in [tokenizer.tok_to_ind.get("<BOS>", -1), 
+                                           tokenizer.tok_to_ind.get("<EOS>", -1), 
+                                           tokenizer.tok_to_ind.get("<PAD>", -1)]])
+            result_text.append(caption)
+
+    return result_text[0] if batch_size == 1 else result_text
+
+def generate_caption(image):
+    image = image.convert("RGB")
+    image_tensor = transform(image).unsqueeze(0).to(device)
+    caption = generate(model, image_tensor, tokenizer, max_seq_len=config.get("max_len", 20), top_k=5)
+    caption = wrap_text(caption, title_width)
+    return caption
+
+# Создание Gradio-интерфейса
+iface = gr.Interface(
+    fn=generate_caption,
+    inputs=gr.Image(type="pil"),
+    outputs="text",
+    title="Image Captioning Demo",
+    description="Upload an image to generate a caption."
+)
+iface.launch()

model.py

@@ -0,0 +1,63 @@
+import torch
+import torch.nn as nn
+import torchvision.models as models
+
+class img_fe_class(nn.Module):
+    def __init__(self):
+        super(img_fe_class, self).__init__()
+        weights = models.ResNet50_Weights.DEFAULT
+        self.resnet = models.resnet50(weights=weights)
+        self.resnet = nn.Sequential(*list(self.resnet.children())[:-1])
+        for param in self.resnet.parameters():
+            param.requires_grad = False
+
+    def forward(self, imgs):
+        features = self.resnet(imgs)
+        return features.reshape(features.size(0), -1)
+
+class text_fe_class(nn.Module):
+    def __init__(self, vocab_size, embedding_dim, hidden_size, num_layers):
+        super(text_fe_class, self).__init__()
+        self.embedding = nn.Embedding(vocab_size, embedding_dim)
+        self.rnn = nn.LSTM(embedding_dim, hidden_size, num_layers, batch_first=True)
+
+    def forward(self, captions, hidden):
+        embeddings = self.embedding(captions)
+        outputs, hidden = self.rnn(embeddings, hidden)
+        return outputs, hidden
+
+class image_captioning_model(nn.Module):
+    def __init__(self, vocab_size, embedding_dim, hidden_size, num_layers):
+        super(image_captioning_model, self).__init__()
+        self.img_fe = img_fe_class()
+        self.text_fe = text_fe_class(vocab_size, embedding_dim, hidden_size, num_layers)
+        self.fc_out = nn.Linear(hidden_size, vocab_size)
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.img_to_hidden = nn.Linear(2048, hidden_size)  # 2048 - размер выхода ResNet50
+
+    def forward(self, imgs, captions):
+        img_features = self.img_fe(imgs)
+        hidden = self.img_to_hidden(img_features).unsqueeze(0).repeat(self.num_layers, 1, 1)
+        cell = torch.zeros_like(hidden)  # Для LSTM нужно также cell state
+        outputs, _ = self.text_fe(captions, (hidden, cell))
+        outputs = self.fc_out(outputs)
+        return outputs
+
+    def generate_caption(self, img, tokenizer, max_len=30, device="cuda"):
+        self.eval()
+        with torch.no_grad():
+            img = img.unsqueeze(0).to(device)
+            img_features = self.img_fe(img)
+            hidden = self.img_to_hidden(img_features).unsqueeze(0).repeat(self.num_layers, 1, 1)
+            cell = torch.zeros_like(hidden)
+            caption = [tokenizer.tok_to_ind["<BOS>"]]
+            for _ in range(max_len):
+                inputs = torch.LongTensor(caption[-1]).unsqueeze(0).unsqueeze(0).to(device)
+                outputs, (hidden, cell) = self.text_fe(inputs, (hidden, cell))
+                outputs = self.fc_out(outputs[:, -1, :])
+                next_token = outputs.argmax(dim=-1).item()
+                caption.append(next_token)
+                if next_token == tokenizer.tok_to_ind["<EOS>"]:
+                    break
+            return tokenizer.decode(caption)