main.py

main.py

@@ -1,31 +1,27 @@
 import re
 from collections import Counter
+import pickle
 
 import torch
 import torch.nn as nn
-from torch.utils.data import Dataset, DataLoader, random_split
+from torch.utils.data import Dataset, DataLoader
 from PIL import Image
 import torchvision.transforms as transforms
 import torchvision.models as models
-from datasets import load_dataset
-import pandas as pd
 from langdetect import detect
-import pickle
 
-# Global config
+# ========== Global config ==========
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 MAX_LEN = 20
 
-# Image transform used for both training and inference
+# Fixed image transform
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor()
 ])
 
 
-# =========================
-# Dataset & preprocessing
-# =========================
+# ========== Dataset & preprocessing (for local training only) ==========
 class VQADataset(Dataset):
     def __init__(self, df, transform):
         self.df = df.reset_index(drop=True)
@@ -44,10 +40,12 @@ class VQADataset(Dataset):
 
 
 def prepare_data(max_answers=50, min_word_count=3, max_len=MAX_LEN):
-    """Load VQA-RAD, clean, build vocab, return dataset + mappings."""
+    """For local training only – not used in Space."""
+    from datasets import load_dataset
+    import pandas as pd
+
     dataset = load_dataset("flaviagiammarino/vqa-rad")
-    df = pd.DataFrame(dataset["train"])
-    df = df[["image", "question", "answer"]]
+    df = pd.DataFrame(dataset["train"])[["image", "question", "answer"]]
 
     def clean_text(text):
         text = str(text).lower()
@@ -95,17 +93,17 @@ def prepare_data(max_answers=50, min_word_count=3, max_len=MAX_LEN):
     }
 
 
-# =========================
-# Model
-# =========================
+# ========== Model ==========
 class VQAModel(nn.Module):
     def __init__(self, vocab_size, embed_dim, hidden_dim, num_answers):
         super().__init__()
-        # Use pretrained ResNet18
+        # Use lightweight ResNet18 backbone
         self.cnn = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
         self.cnn.fc = nn.Identity()
+
         self.embedding = nn.Embedding(vocab_size, embed_dim)
         self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
+
         self.fc1 = nn.Linear(512 + hidden_dim, 256)
         self.relu = nn.ReLU()
         self.fc2 = nn.Linear(256, num_answers)
@@ -115,18 +113,15 @@ class VQAModel(nn.Module):
         q_embed = self.embedding(question)
         _, (h, _) = self.lstm(q_embed)
         q_feat = h.squeeze(0)
-        combined = torch.cat((img_feat, q_feat), dim=1)
-        x = self.relu(self.fc1(combined))
-        out = self.fc2(x)
-        return out
+        x = torch.cat((img_feat, q_feat), dim=1)
+        x = self.relu(self.fc1(x))
+        return self.fc2(x)
 
 
-# =========================
-# Training (local, not used on Space)
-# =========================
+# ========== Training (local only) ==========
 def train_model(train_dataset, vocab, idx_to_answer,
                 epochs=20, batch_size=32, lr=1e-3, save_prefix="vqa_custom"):
-    """Train model locally and save artifacts."""
+    """Run this only on Colab / local, not in Space."""
     vocab_size = len(vocab)
     num_answers = len(idx_to_answer)
 
@@ -161,53 +156,7 @@ def train_model(train_dataset, vocab, idx_to_answer,
     return model
 
 
-# =========================
-# BLIP + translator (optional, heavy)
-# =========================
-def load_blip_and_translator():
-    from transformers import Blip2Processor, Blip2ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM
-
-    processor = Blip2Processor.from_pretrained("Salesforce/blip2-flan-t5-xl")
-    blip_model = Blip2ForConditionalGeneration.from_pretrained(
-        "Salesforce/blip2-flan-t5-xl",
-        device_map="auto",
-        torch_dtype=torch.float16,
-    )
-
-    translator_tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
-    translator_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
-
-    lang_code_map = {
-        "en": "eng_Latn",
-        "hi": "hin_Deva",
-        "te": "tel_Telu",
-        "ta": "tam_Taml",
-        "kn": "kan_Knda",
-        "ml": "mal_Mlym",
-    }
-
-    def open_vqa(image_path, question):
-        image = Image.open(image_path).convert("RGB")
-        inputs = processor(image, question, return_tensors="pt").to(blip_model.device)
-        out = blip_model.generate(**inputs, max_new_tokens=50)
-        return processor.decode(out[0], skip_special_tokens=True)
-
-    def translate(text, src, tgt):
-        translator_tokenizer.src_lang = lang_code_map[src]
-        inputs = translator_tokenizer(text, return_tensors="pt")
-        tokens = translator_model.generate(
-            **inputs,
-            forced_bos_token_id=translator_tokenizer.convert_tokens_to_ids(lang_code_map[tgt]),
-            max_length=50,
-        )
-        return translator_tokenizer.decode(tokens[0], skip_special_tokens=True)
-
-    return processor, blip_model, translator_tokenizer, translator_model, open_vqa, translate
-
-
-# =========================
-# Load artifacts + inference
-# =========================
+# ========== Load artifacts + inference ==========
 def load_artifacts(prefix="vqa_custom", map_location=None):
     with open(f"{prefix}_vocab.pkl", "rb") as f:
         vocab = pickle.load(f)
@@ -220,8 +169,7 @@ def load_artifacts(prefix="vqa_custom", map_location=None):
     model.eval()
 
     def encode_question_infer(q, max_len=MAX_LEN):
-        q = str(q).lower()
-        tokens = q.split()
+        tokens = str(q).lower().split()
         enc = [vocab.get(w, vocab["<UNK>"]) for w in tokens]
         enc = enc[:max_len] + [vocab["<PAD>"]] * (max_len - len(enc))
         return torch.tensor(enc).unsqueeze(0)
@@ -236,12 +184,9 @@ def load_artifacts(prefix="vqa_custom", map_location=None):
         return idx_to_answer[pred.item()]
 
     def final_pipeline(image_path, question, open_vqa_fn=None, translate_fn=None):
+        # Keep it simple in Space: English only, no BLIP / translator unless passed in
         lang = detect(question)
-
-        if lang != "en" and translate_fn is not None:
-            q_en = translate_fn(question, lang, "en")
-        else:
-            q_en = question
+        q_en = question if lang == "en" or translate_fn is None else translate_fn(question, lang, "en")
 
         if open_vqa_fn is not None and ("what is" in q_en.lower() or "this place" in q_en.lower()):
             answer_en = open_vqa_fn(image_path, q_en)
@@ -250,8 +195,7 @@ def load_artifacts(prefix="vqa_custom", map_location=None):
 
         if lang != "en" and translate_fn is not None:
             return translate_fn(answer_en, "en", lang)
-        else:
-            return answer_en
+        return answer_en
 
     return final_pipeline, predict_custom_vqa, vocab, idx_to_answer, model, encode_question_infer
 
@@ -260,24 +204,21 @@ def load_artifacts_and_helpers(prefix="vqa_custom", map_location=None):
     return load_artifacts(prefix=prefix, map_location=map_location)
 
 
-# =========================
-# Optional CLI (local only)
-# =========================
 if __name__ == "__main__":
+    # Local CLI only; never runs in Space
     import argparse
-
-    parser = argparse.ArgumentParser(description="VQA pipeline (prepare/train/infer)")
+    parser = argparse.ArgumentParser()
     parser.add_argument("--prepare", action="store_true")
     parser.add_argument("--train", action="store_true")
     parser.add_argument("--epochs", type=int, default=20)
     parser.add_argument("--prefix", default="vqa_custom")
-    parser.add_argument("--image", help="Path to image for inference")
+    parser.add_argument("--image")
     parser.add_argument("--question", default="What is in the image?")
     args = parser.parse_args()
 
     if args.prepare or args.train:
         artifacts = prepare_data()
-        print("Prepared dataset with", len(artifacts["answer_to_idx"]), "answer classes.")
+        print("Prepared dataset with", len(artifacts["answer_to_idx"]), "answers")
 
     if args.train:
         train_model(
@@ -289,5 +230,5 @@ if __name__ == "__main__":
         )
 
     if args.image:
-        final_pipeline, _, _, _, _, _ = load_artifacts(prefix=args.prefix)
+        final_pipeline, *_ = load_artifacts(prefix=args.prefix)
         print(final_pipeline(args.image, args.question, open_vqa_fn=None, translate_fn=None))