Update app.py

app.py

@@ -1,137 +1,58 @@
-import gradio as gr
-import torch
-from transformers import ViTModel
+!pip install -q torch torchvision transformers huggingface_hub pillow gradio
+
+import torch, json, gradio as gr
 from PIL import Image
 from torchvision import transforms
-import json
-import os
+from transformers import ViTModel
 from huggingface_hub import hf_hub_download
 
-# ---------------------
-# Config
-# ---------------------
-IMG_SIZE = 224
-SEQ_LEN = 32
-VOCAB_SIZE = 75460
-REPO_ID = "hackergeek/RADIOCAP13"             # your HF repo
-WEIGHTS_FILENAME = "pytorch_model.bin"
-VOCAB_FILENAME   = "vocab.json"
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# ---------------------
-# Download model files (if not present)
-# ---------------------
-# Download weights
-weights_path = hf_hub_download(repo_id=REPO_ID, filename=WEIGHTS_FILENAME)
-# Download vocab
-vocab_path   = hf_hub_download(repo_id=REPO_ID, filename=VOCAB_FILENAME)
-
-# ---------------------
-# Preprocessing & Tokenizer
-# ---------------------
-transform = transforms.Compose([
-    transforms.Resize((IMG_SIZE, IMG_SIZE)),
-    transforms.ToTensor(),
-])
+# --------------------- تنظیمات ---------------------
+SEQ_LEN=32
+REPO="hackergeek/RADIOCAP13"
+device="cuda" if torch.cuda.is_available() else "cpu"
 
-def preprocess_image(img):
-    if img is None:
-        raise ValueError("Image is None")
-    if not isinstance(img, Image.Image):
-        img = Image.fromarray(img)
-    if img.mode != "RGB":
-        img = img.convert("RGB")
-    return transform(img)
+# دانلود وزن‌ها و vocab
+w=hf_hub_download(REPO,"pytorch_model.bin")
+v=hf_hub_download(REPO,"vocab.json")
 
-class SimpleTokenizer:
-    def __init__(self, word2idx=None):
-        self.word2idx = word2idx or {}
-        self.idx2word = {v: k for k, v in self.word2idx.items()}
+# --------------------- Tokenizer ---------------------
+with open(v) as f: vocab=json.load(f)
+idx2word={v:k for k,v in vocab.items()}
+pad=vocab.get("<PAD>",0); sos=vocab.get("<SOS>",1); eos=vocab.get("<EOS>",2)
 
-    @classmethod
-    def load(cls, path):
-        with open(path, "r") as f:
-            word2idx = json.load(f)
-        return cls(word2idx)
-
-# ---------------------
-# Decoder
-# ---------------------
-class BiasDecoder(torch.nn.Module):
-    def __init__(self, feature_dim=768, vocab_size=VOCAB_SIZE):
+# --------------------- مدل‌ها ---------------------
+vit=ViTModel.from_pretrained("google/vit-base-patch16-224-in21k").to(device).eval()
+class D(torch.nn.Module):
+    def __init__(self):
         super().__init__()
-        self.token_emb    = torch.nn.Embedding(vocab_size, feature_dim)
-        self.pos_emb      = torch.nn.Embedding(SEQ_LEN-1, feature_dim)
-        self.final_layer  = torch.nn.Linear(feature_dim, vocab_size)
-
-    def forward(self, img_feat, target_seq):
-        x   = self.token_emb(target_seq)
-        pos = torch.arange(x.size(1), device=x.device).clamp(max=self.pos_emb.num_embeddings-1)
-        x   = x + self.pos_emb(pos)
-        x   = x + img_feat.unsqueeze(1)
+        self.token_emb=torch.nn.Embedding(75460,768)
+        self.pos_emb=torch.nn.Embedding(SEQ_LEN-1,768)
+        self.final_layer=torch.nn.Linear(768,75460)
+    def forward(self,f,s):
+        x=self.token_emb(s)
+        x+=self.pos_emb(torch.arange(x.size(1),device=x.device).clamp(max=self.pos_emb.num_embeddings-1))
+        x+=f.unsqueeze(1)
         return self.final_layer(x)
+decoder=D().to(device); decoder.load_state_dict(torch.load(w,map_location=device)); decoder.eval()
 
-# ---------------------
-# Load models
-# ---------------------
-vit = ViTModel.from_pretrained("google/vit-base-patch16-224-in21k").to(device)
-vit.eval()
-
-decoder = BiasDecoder().to(device)
-decoder.load_state_dict(torch.load(weights_path, map_location=device))
-decoder.eval()
+# --------------------- پیش‌پردازش تصویر ---------------------
+def preprocess(img):
+    if img.mode!="RGB": img=img.convert("RGB")
+    t=transforms.Compose([transforms.Resize((224,224)),transforms.ToTensor()])
+    return t(img).unsqueeze(0).to(device)
 
-tokenizer = SimpleTokenizer.load(vocab_path)
-pad_idx   = tokenizer.word2idx["<PAD>"]
-
-# ---------------------
-# Caption generation
-# ---------------------
+# --------------------- تولید کپشن ---------------------
 @torch.no_grad()
-def generate_caption(img, max_len=SEQ_LEN, beam_size=3):
-    img_tensor = preprocess_image(img).unsqueeze(0).to(device)
-    img_feat   = vit(pixel_values=img_tensor).pooler_output
-
-    beams = [([tokenizer.word2idx["<SOS>"]], 0.0)]
-
-    for _ in range(max_len - 1):
-        candidates = []
-        for seq, score in beams:
-            inp    = torch.tensor(seq + [pad_idx] * (SEQ_LEN - len(seq)), device=device).unsqueeze(0)
-            logits = decoder(img_feat, inp)
-            probs  = torch.nn.functional.log_softmax(logits[0, len(seq)-1], dim=-1)
-            top_p, top_i = torch.topk(probs, beam_size)
-
-            for i in range(beam_size):
-                candidates.append((seq + [top_i[i].item()], score + top_p[i].item()))
-
-        beams = sorted(candidates, key=lambda x: x[1], reverse=True)[:beam_size]
-
-        if all(s[-1] == tokenizer.word2idx["<EOS>"] for s, _ in beams):
-            break
-
-    words = [tokenizer.idx2word.get(i, "<UNK>") for i in beams[0][0][1:] if i != pad_idx]
-    return " ".join(words)
-
-# ---------------------
-# Gradio interface
-# ---------------------
-with gr.Blocks() as demo:
-    gr.Markdown("# RADIOCAP13 — Image Captioning Demo")
-    gr.Markdown(f"**Device:** {'GPU 🚀' if torch.cuda.is_available() else 'CPU 🐢'}")
-
-    img_in = gr.Image(type="pil", label="Upload an Image")
-    out    = gr.Textbox(label="Generated Caption")
-    btn    = gr.Button("Generate Caption")
-    status = gr.Markdown("Ready.")
-
-    def wrapped(img):
-        status.update("Processing…")
-        caption = generate_caption(img)
-        status.update("Done ✔️")
-        return caption
-
-    btn.click(wrapped, inputs=img_in, outputs=out)
-
-if __name__ == "__main__":
-    demo.launch()
+def caption(img):
+    f=vit(pixel_values=preprocess(img)).pooler_output
+    seq=[sos]
+    for _ in range(SEQ_LEN-1):
+        inp=torch.tensor(seq+[pad]*(SEQ_LEN-len(seq)),device=device).unsqueeze(0)
+        logits=decoder(f,inp)
+        nxt=torch.argmax(logits[0,len(seq)-1]).item()
+        seq.append(nxt)
+        if nxt==eos: break
+    return " ".join(idx2word.get(i,"<UNK>") for i in seq[1:] if i not in [pad,sos,eos])
+
+# --------------------- رابط Gradio ---------------------
+gr.Interface(fn=caption, inputs=gr.Image(type="pil"), outputs="text", title="RADIOCAP13 Captioning").launch()