feat: add HTML and CSS

app.py

@@ -6,6 +6,11 @@ from transformers import ViTModel
 from PIL import Image
 import pickle
 import re
+import os
+
+# ==========================================
+# 1. CLASS DEFINITIONS (Must Match Training)
+# ==========================================
 
 class Vocabulary:
     def __init__(self, freq_threshold=5):
@@ -36,11 +41,9 @@ class Encoder(nn.Module):
     def __init__(self, embed_dim, freeze=False):
         super().__init__()
         self.vit = ViTModel.from_pretrained("facebook/vit-mae-base")
-
         if freeze:
             for param in self.vit.parameters():
                 param.requires_grad = False
-
         self.linear = nn.Sequential(
             nn.Linear(self.vit.config.hidden_size, embed_dim),
             nn.ReLU(),
@@ -61,9 +64,7 @@ class MultiHeadAttention(nn.Module):
         self.num_heads = num_heads
         self.hidden_dim = hidden_dim
         self.head_dim = hidden_dim // num_heads
-        
         assert hidden_dim % num_heads == 0, "hidden_dim must be divisible by num_heads"
-
         self.query = nn.Linear(hidden_dim, hidden_dim)
         self.key = nn.Linear(encoder_dim, hidden_dim)
         self.value = nn.Linear(encoder_dim, hidden_dim)
@@ -74,7 +75,6 @@ class MultiHeadAttention(nn.Module):
         Q = self.query(hidden).view(B, self.num_heads, self.head_dim)
         K = self.key(encoder_outputs).view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
         V = self.value(encoder_outputs).view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
-
         scores = torch.matmul(Q.unsqueeze(2), K.transpose(-2, -1)) / (self.head_dim ** 0.5)
         attn = torch.softmax(scores, dim=-1)
         context = torch.matmul(attn, V)
@@ -93,7 +93,6 @@ class Decoder(nn.Module):
     def generate(self, features, max_len=50, start_index=1, end_index=2, beam_size=3, beam_search=True):
         B = features.size(0)
         device = features.device
-        
         states = (torch.zeros(self.lstm.num_layers, B, self.lstm.hidden_size, device=device),
                   torch.zeros(self.lstm.num_layers, B, self.lstm.hidden_size, device=device))
 
@@ -124,16 +123,13 @@ class Decoder(nn.Module):
                     logits = self.fc(out.squeeze(1))
                     log_probs = torch.log_softmax(logits, dim=1)
                     top_log_probs, top_indices = log_probs.topk(beam_size, dim=1)
-                    
                     for k in range(beam_size):
                         next_seq = seq + [top_indices[0, k].item()]
                         next_log_prob = log_prob + top_log_probs[0, k].item()
                         new_beams.append((next_seq, next_log_prob, (h_new, c_new)))
-                
                 new_beams = sorted(new_beams, key=lambda x: x[1], reverse=True)[:beam_size]
                 beams = new_beams
                 if all(seq[-1] == end_index for seq, _, _ in beams): break
-            
             best_seq = beams[0][0]
             if best_seq[0] == start_index: best_seq = best_seq[1:]
             return best_seq
@@ -149,48 +145,138 @@ class Model(nn.Module):
         captions = self.decoder.generate(features, max_len=max_len, beam_search=True)
         return captions
 
+# ==========================================
+# 2. SETUP AND LOADING
+# ==========================================
+
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 EMBED_DIM = 256
 HIDDEN_DIM = 512
 VOCAB_PATH = "vocab-v4.pkl"
-MODEL_PATH = "vit_lstm_best-v4.pth" 
-
-print("Loading Vocabulary...")
-try:
-    with open(VOCAB_PATH, "rb") as f:
-        vocab = pickle.load(f)
-    print(f"Vocabulary Loaded. Size: {len(vocab)}")
-except FileNotFoundError:
-    raise RuntimeError("vocab.pkl not found! Please upload it.")
-
-print("Initializing Model...")
-encoder = Encoder(EMBED_DIM, freeze=True)
-decoder = Decoder(EMBED_DIM, HIDDEN_DIM, len(vocab))
-model = Model(encoder, decoder).to(DEVICE)
-
-print("Loading Weights...")
-try:
-    checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)
-    if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
-        model.load_state_dict(checkpoint['model_state_dict'])
-    else:
-        model.load_state_dict(checkpoint)
-    model.eval()
-    print("Model Loaded Successfully!")
-except FileNotFoundError:
-    raise RuntimeError("vit_lstm.pth not found! Please upload it.")
-except Exception as e:
-    print(f"Warning loading weights: {e}")
-
-inference_transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-])
-
-def generate_caption(image):
+MODEL_PATH = "vit_lstm_best-v4.pth"
+
+# Global Variables
+vocab = None
+model = None
+inference_transform = None
+
+def load_system():
+    global vocab, model, inference_transform
+    print("Loading Vocabulary...")
+    try:
+        with open(VOCAB_PATH, "rb") as f:
+            vocab = pickle.load(f)
+    except Exception as e:
+        return f"Error loading vocab: {e}"
+
+    print("Initializing Model...")
+    encoder = Encoder(EMBED_DIM, freeze=True)
+    decoder = Decoder(EMBED_DIM, HIDDEN_DIM, len(vocab))
+    model = Model(encoder, decoder).to(DEVICE)
+
+    print("Loading Weights...")
+    try:
+        checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)
+        if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
+            model.load_state_dict(checkpoint['model_state_dict'])
+        else:
+            model.load_state_dict(checkpoint)
+        model.eval()
+    except Exception as e:
+        return f"Error loading model weights: {e}"
+
+    inference_transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+    return "System Loaded"
+
+# Load on startup
+load_status = load_system()
+
+# ==========================================
+# 3. HTML FORMATTERS
+# ==========================================
+
+def format_loading_html():
+    return """
+    <div class="loading-box" style="
+        text-align: center; 
+        padding: 40px; 
+        border: 2px solid #6B7280; 
+        border-radius: 15px; 
+        background-color: #27272A;
+    ">
+        <h3 style="color: #6B7280; margin-bottom: 5px;">Status</h3>
+        <h2 style="color: #F3F4F6; font-size: 24px; margin: 10px 0;">Analyzing Image...</h2>
+        
+        <div style="font-size: 48px; font-weight: bold; color: #6B7280;">
+            ...
+        </div>
+        
+        <p style="color: #9CA3AF; font-weight: bold; margin-top: 5px;">Please wait...</p>
+    </div>
+    """
+
+def format_result_html(caption):
+    return f"""
+    <div style="
+        text-align: center; 
+        padding: 30px; 
+        border: 2px solid #4F46E5; 
+        border-radius: 15px; 
+        background-color: #27272A;
+        box-shadow: 0 4px 6px -1px rgba(79, 70, 229, 0.1);
+    ">
+        <h3 style="color: #818CF8; margin-bottom: 10px; text-transform: uppercase; letter-spacing: 2px;">Generated Caption</h3>
+        <div style="
+            font-size: 28px; 
+            font-weight: bold; 
+            color: #F9FAFB; 
+            margin: 20px 0;
+            line-height: 1.4;
+        ">
+            "{caption}"
+        </div>
+        <p style="color: #34D399; font-size: 14px;">✓ Analysis Complete</p>
+    </div>
+    """
+
+def format_initial_html():
+    return """
+    <div style="
+        text-align: center; 
+        padding: 40px; 
+        border: 2px dashed #4B5563; 
+        border-radius: 15px; 
+        background-color: #1F2937;
+        color: #9CA3AF;
+    ">
+        <h3>Output Area</h3>
+        <p>Your generated caption will appear here.</p>
+    </div>
+    """
+
+def format_error_html(error):
+    return f"""
+    <div style="text-align: center; padding: 20px; border: 2px solid #EF4444; border-radius: 15px; background-color: #450A0A;">
+        <h3 style="color: #F87171;">Error</h3>
+        <p style="color: #FECACA;">{error}</p>
+    </div>
+    """
+
+# ==========================================
+# 4. PREDICTION LOGIC
+# ==========================================
+
+def predict(image):
     if image is None:
-        return "Please upload an image."
+        yield format_error_html("No image uploaded"), ""
+        return
+
+    # Yield loading state
+    yield format_loading_html(), ""
     
     try:
         pil_image = image.convert("RGB")
@@ -202,24 +288,245 @@ def generate_caption(image):
         result_words = []
         for idx in output_indices:
             word = vocab.itos.get(idx, "<UNK>")
-            if word == "<EOS>":
-                break
+            if word == "<EOS>": break
             if word not in ("<SOS>", "<PAD>"):
                 result_words.append(word)
                 
         caption = " ".join(result_words)
-        return caption
+        
+        # Yield final result (HTML for display, Raw Text for clipboard)
+        yield format_result_html(caption), caption
         
     except Exception as e:
-        return f"Error occurred: {str(e)}"
-
-iface = gr.Interface(
-    fn=generate_caption,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.Textbox(label="Generated Caption"),
-    title="ViT + LSTM Image Captioning",
-    description="Upload an image to generate a caption using a Vision Transformer (Encoder) and LSTM (Decoder) architecture."
+        yield format_error_html(str(e)), ""
+
+# ==========================================
+# 5. JAVASCRIPT & CSS
+# ==========================================
+
+# JS for Copying text and Toggling Modal
+custom_js = """
+<script>
+function copyToClipboard() {
+    // Select the hidden textarea
+    const textarea = document.querySelector('#hidden_caption_output textarea');
+    if (!textarea || !textarea.value) {
+        alert("No caption to copy!");
+        return;
+    }
+    
+    // Copy logic
+    navigator.clipboard.writeText(textarea.value).then(function() {
+        // Show styled toast
+        const toast = document.createElement("div");
+        toast.innerText = "Caption Copied!";
+        toast.style.position = "fixed";
+        toast.style.bottom = "20px";
+        toast.style.right = "20px";
+        toast.style.backgroundColor = "#10B981";
+        toast.style.color = "white";
+        toast.style.padding = "10px 20px";
+        toast.style.borderRadius = "5px";
+        toast.style.zIndex = "9999";
+        document.body.appendChild(toast);
+        setTimeout(() => toast.remove(), 2000);
+    }, function(err) {
+        console.error('Async: Could not copy text: ', err);
+    });
+}
+
+function toggleApiModal() {
+    const modal = document.getElementById('api_modal');
+    if (modal.style.display === 'flex') {
+        modal.style.display = 'none';
+    } else {
+        modal.style.display = 'flex';
+    }
+}
+
+function closeApiModal(e) {
+    if (e.target.id === 'api_modal') {
+        document.getElementById('api_modal').style.display = 'none';
+    }
+}
+</script>
+"""
+
+# CSS for Layout, Dark Mode, and Modal
+custom_css = """
+body { background-color: #111827; }
+.container { max-width: 800px; margin: auto; padding-top: 20px; }
+.header { text-align: center; margin-bottom: 30px; }
+.header h1 { color: #818CF8; font-size: 2.5rem; }
+.header p { color: #9CA3AF; }
+
+/* Pulse Animation for Loading */
+@keyframes pulse {
+    0%, 100% { opacity: 1; }
+    50% { opacity: 0.5; }
+}
+.loading-box { animation: pulse 1.5s cubic-bezier(0.4, 0, 0.6, 1) infinite; }
+
+/* Modal Styles */
+#api_modal {
+    display: none; /* Hidden by default */
+    position: fixed;
+    z-index: 1000;
+    left: 0;
+    top: 0;
+    width: 100%;
+    height: 100%;
+    overflow: auto;
+    background-color: rgba(0,0,0,0.8);
+    justify-content: center;
+    align-items: center;
+}
+.modal-content {
+    background-color: #1F2937;
+    margin: auto;
+    padding: 30px;
+    border: 1px solid #374151;
+    border-radius: 10px;
+    width: 80%;
+    max-width: 600px;
+    color: #F3F4F6;
+    position: relative;
+    box-shadow: 0 10px 25px rgba(0,0,0,0.5);
+}
+.close-btn {
+    color: #9CA3AF;
+    float: right;
+    font-size: 28px;
+    font-weight: bold;
+    cursor: pointer;
+}
+.close-btn:hover { color: #F3F4F6; }
+code {
+    background-color: #111827;
+    padding: 2px 5px;
+    border-radius: 4px;
+    color: #F472B6;
+    font-family: monospace;
+}
+pre {
+    background-color: #111827;
+    padding: 15px;
+    border-radius: 8px;
+    overflow-x: auto;
+    color: #D1D5DB;
+}
+"""
+
+# ==========================================
+# 6. GRADIO INTERFACE CONSTRUCTION
+# ==========================================
+
+with gr.Blocks(css=custom_css, title="CogniCaption") as app:
+    # Inject JS helper functions
+    gr.HTML(custom_js)
+    
+    # Hidden Modal HTML Structure
+    gr.HTML("""
+    <div id="api_modal" onclick="closeApiModal(event)">
+        <div class="modal-content">
+            <span class="close-btn" onclick="toggleApiModal()">&times;</span>
+            <h2 style="margin-top:0; color: #818CF8;">Use CogniCaption as API</h2>
+            <hr style="border-color: #374151; margin: 15px 0;">
+            
+            <p>You can use this Hugging Face Space as an API via the <code>gradio_client</code>.</p>
+            
+            <h4>1. API Endpoint</h4>
+            <div style="display:flex; gap:10px; margin-bottom:15px;">
+                <input type="text" value="https://huggingface.co/spaces/YOUR_USERNAME/SPACE_NAME" readonly 
+                       style="width:100%; padding:10px; background:#111827; border:1px solid #374151; color:#9CA3AF; border-radius:5px;">
+            </div>
+            
+            <h4>2. How to Request</h4>
+            <p style="font-size:0.9rem; color:#9CA3AF;">
+                Send an image (filepath or URL) to the <code>predict</code> endpoint. The API returns a generated text caption.
+            </p>
+            
+            <h4>3. Python Example</h4>
+            <pre>
+from gradio_client import Client
+
+client = Client("YOUR_USERNAME/SPACE_NAME")
+result = client.predict(
+    image="https://example.com/image.jpg",
+    api_name="/predict"
 )
+print(result) # Outputs the caption tuple
+            </pre>
+            
+            <div style="text-align:right; margin-top:20px;">
+                <button onclick="toggleApiModal()" style="
+                    background-color: #4F46E5; 
+                    color: white; 
+                    border: none; 
+                    padding: 10px 20px; 
+                    border-radius: 5px; 
+                    cursor: pointer;">
+                    OK, Got it
+                </button>
+            </div>
+        </div>
+    </div>
+    """)
+
+    with gr.Column(elem_class="container"):
+        # Header
+        gr.HTML("""
+        <div class="header">
+            <h1>CogniCaption</h1>
+            <p>ViT + LSTM Image Captioning System</p>
+        </div>
+        """)
+        
+        # --- INPUT SECTION (TOP) ---
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Upload Image", elem_id="input_image")
+            submit_btn = gr.Button("Generate Caption", variant="primary", size="lg")
+
+        # Separator
+        gr.HTML("<hr style='border-color: #374151; margin: 30px 0;'>")
+
+        # --- OUTPUT SECTION (BOTTOM) ---
+        with gr.Column():
+            # The HTML Display for the user
+            output_display = gr.HTML(label="Result", value=format_initial_html())
+            
+            # Hidden Textbox to hold the raw text for the Copy Button to access
+            hidden_caption_storage = gr.Textbox(visible=False, elem_id="hidden_caption_output")
+            
+            # Action Buttons Row
+            with gr.Row():
+                copy_btn = gr.Button("Copy Caption to Clipboard", size="sm")
+                api_btn = gr.Button("Use CogniCaption As API", size="sm", variant="secondary")
+
+    # --- EVENT LISTENERS ---
+    
+    # 1. Prediction Event
+    submit_btn.click(
+        fn=predict,
+        inputs=[input_image],
+        outputs=[output_display, hidden_caption_storage]
+    )
+    
+    # 2. Copy Button Event (Trigger JS)
+    copy_btn.click(
+        fn=None,
+        inputs=None,
+        outputs=None,
+        js="copyToClipboard"
+    )
+    
+    # 3. API Button Event (Trigger JS Modal)
+    api_btn.click(
+        fn=None,
+        inputs=None,
+        outputs=None,
+        js="toggleApiModal"
+    )
 
 if __name__ == "__main__":
-    iface.launch()
+    app.launch()