Update app.py

app.py

@@ -5,35 +5,119 @@ from torchvision import transforms
 
 import torch.nn as nn
 import torchvision.models as models
-
+# -----------------------
+# VQA Model
+# -----------------------
 class VQAModel(nn.Module):
-    def __init__(self, vocab_size):
+    def __init__(self, vocab_size, embedding_dim=256, lstm_units=256, cnn_output_dim=512, attention_dim=256, max_seq_len=30):
         super(VQAModel, self).__init__()
-        # Dùng ResNet làm CNN Encoder
-        self.cnn = models.resnet18(pretrained=True)
-        self.cnn.fc = nn.Linear(512, 256)  # Thay FC layer
-
-        # Dùng LSTM làm Text Encoder
-        self.embedding = nn.Embedding(vocab_size, 256)
-        self.lstm = nn.LSTM(256, 256, batch_first=True)
+        self.vocab_size = vocab_size
+        self.max_seq_len = max_seq_len
+
+        # CNN Encoder: Trích xuất đặc trưng ảnh
+        self.cnn = nn.Sequential(
+            nn.Conv2d(3, 32, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Conv2d(32, 64, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Conv2d(64, 128, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Conv2d(128, cnn_output_dim, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.AdaptiveAvgPool2d((1, 1))
+        )
+
+        # Text Embedding
+        self.embedding = nn.Embedding(vocab_size, embedding_dim)
+
+        # LSTM Encoders cho caption và question
+        self.caption_lstm = nn.LSTM(embedding_dim, lstm_units, batch_first=True)
+        self.question_lstm = nn.LSTM(embedding_dim, lstm_units, batch_first=True)
+
+        # Attention cho từng kênh
+        self.attention = Attention(cnn_output_dim, lstm_units, attention_dim)
+
+        # Decoder: sử dụng teacher forcing
+        # Context vector: kết hợp của attention từ caption, attention từ question và trạng thái cuối của question
+        # Kích thước context = lstm_units + lstm_units + lstm_units = 3 * lstm_units (ví dụ 768 nếu lstm_units=256)
+        # Kết hợp với embedding của câu trả lời (embedding_dim) => đầu vào của decoder = embedding_dim + 3*lstm_units
+        self.decoder_input_proj = nn.Linear(embedding_dim + 3 * lstm_units, lstm_units)
+        self.decoder_lstm = nn.LSTM(lstm_units, lstm_units, batch_first=True)
+        self.fc_out = nn.Linear(lstm_units, vocab_size)
+        self.dropout = nn.Dropout(0.5)
+
+    def forward(self, image, caption, question, answer_input):
+        # --- CNN Encoder ---
+        cnn_features = self.cnn(image)  # (batch, cnn_output_dim, 1, 1)
+        cnn_features = cnn_features.view(cnn_features.size(0), -1)  # (batch, cnn_output_dim)
+
+        # --- Text Encoders ---
+        cap_embed = self.embedding(caption)  # (batch, cap_seq_len, embedding_dim)
+        cap_output, _ = self.caption_lstm(cap_embed)  # (batch, cap_seq_len, lstm_units)
+
+        q_embed = self.embedding(question)  # (batch, q_seq_len, embedding_dim)
+        q_output, _ = self.question_lstm(q_embed)  # (batch, q_seq_len, lstm_units)
+
+        # --- Attention ---
+        cap_attended = self.attention(cnn_features.unsqueeze(1), cap_output)  # (batch, lstm_units)
+        q_attended = self.attention(cnn_features.unsqueeze(1), q_output)      # (batch, lstm_units)
+
+        q_last = q_output[:, -1, :]  # (batch, lstm_units)
+
+        # Context vector: (batch, 3*lstm_units)
+        context = torch.cat([cap_attended, q_attended, q_last], dim=-1)
+
+        # --- Decoder với Teacher Forcing ---
+        # answer_input: (batch, ans_seq_len)
+        answer_embed = self.embedding(answer_input)  # (batch, ans_seq_len, embedding_dim)
+        context_repeated = context.unsqueeze(1).repeat(1, answer_input.size(1), 1)  # (batch, ans_seq_len, 3*lstm_units)
+        decoder_in = torch.cat([answer_embed, context_repeated], dim=-1)  # (batch, ans_seq_len, embedding_dim + 3*lstm_units)
+        decoder_in = self.decoder_input_proj(decoder_in)  # (batch, ans_seq_len, lstm_units)
+
+        decoder_output, _ = self.decoder_lstm(decoder_in)  # (batch, ans_seq_len, lstm_units)
+        output = self.fc_out(self.dropout(decoder_output))  # (batch, ans_seq_len, vocab_size)
+        return output
 
-        # Fully Connected Layer để dự đoán câu trả lời
-        self.fc = nn.Linear(256, vocab_size)
+    def predict(self, image, question, word_to_idx, idx_to_word, device='cuda' if torch.cuda.is_available() else 'cpu'):
+        self.eval()
+        self.to(device)
 
-    def forward(self, image, question):
-        # Encode ảnh
-        img_features = self.cnn(image)
+        image = image.unsqueeze(0).to(device)
+        question_seq = [word_to_idx.get(word, word_to_idx['<PAD>']) for word in question.lower().split()]
+        question = torch.tensor(question_seq, dtype=torch.long).unsqueeze(0).to(device)
 
-        # Encode câu hỏi
+        # Encode image & question
+        cnn_features = self.cnn(image)
+        cnn_features = cnn_features.view(cnn_features.size(0), -1)
         q_embed = self.embedding(question)
-        _, (q_features, _) = self.lstm(q_embed)
-
-        # Kết hợp đặc trưng ảnh và câu hỏi
-        combined = img_features + q_features.squeeze(0)
-        output = self.fc(combined)
-
-        return output
-
+        q_output, _ = self.question_lstm(q_embed)
+        q_attended = self.attention(cnn_features.unsqueeze(1), q_output)
+        q_last = q_output[:, -1, :]
+        # Ở predict, ta tạo context vector từ q_attended lặp lại (chỉ dùng question cho ví dụ)
+        context = torch.cat([q_attended, q_attended, q_last], dim=-1)  # (1, 3*lstm_units)
+
+        # Khởi tạo câu trả lời với token <START>
+        answer_input = torch.tensor([[word_to_idx['<START>']]], dtype=torch.long).to(device)
+        answer_words = []
+
+        hidden = None
+        for _ in range(self.max_seq_len):
+            answer_embed = self.embedding(answer_input)  # (1, seq_len, embedding_dim)
+            context_repeated = context.unsqueeze(1).repeat(1, answer_input.size(1), 1)
+            decoder_in = torch.cat([answer_embed, context_repeated], dim=-1)
+            decoder_in = self.decoder_input_proj(decoder_in)
+            decoder_output, hidden = self.decoder_lstm(decoder_in, hidden)
+            output = self.fc_out(decoder_output[:, -1, :])
+            next_word_idx = output.argmax(dim=-1).item()
+            if next_word_idx == word_to_idx['<END>']:
+                break
+            answer_words.append(idx_to_word[next_word_idx])
+            answer_input = torch.cat([answer_input, torch.tensor([[next_word_idx]], dtype=torch.long).to(device)], dim=1)
+
+        return ' '.join(answer_words)
 
 
 # Load mô hình từ Hugging Face Model Hub hoặc local