Initial clean commit for Gradio app

.gitattributes

@@ -0,0 +1 @@
+*.pth filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,77 @@
+# CLIP-Style Image Search Engine (Textbook Implementation)
+
+This project provides a complete, modular, and end-to-end implementation of a CLIP-style model for text-to-image search. The architecture and training methodology are a faithful reproduction of the approach described in Chapter 14 of the textbook, "Building an Image Search Engine Using CLIP: a Multimodal Approach".
+
+The project is structured for clarity and maintainability, making it an ideal portfolio piece to showcase skills in PyTorch, model implementation, and MLOps practices like deployment with Streamlit and Hugging Face.
+
+## Key Features
+
+- **Faithful "Book Version" Architecture:** Implements the specific design choices from the textbook:
+  - **Frozen Vision Encoder:** Uses a pre-trained `ResNet50` as a fixed feature extractor.
+  - **Frozen Text Encoder:** Uses a pre-trained `DistilBERT` as a fixed feature extractor.
+  - **Projection Heads:** Maps both image and text features into a shared 256-dimensional space.
+  - **Custom Contrastive Loss:** Implements the unique loss function described in the book.
+- **Modular & Professional Code Structure:** The code is separated into logical files (`config.py`, `dataset.py`, `model.py`, `train.py`, `app.py`) for better organization and scalability.
+- **End-to-End MLOps Pipeline:**
+  - **Training:** A dedicated script to train the model and save the weights.
+  - **Inference:** A standalone Streamlit web application for interactive text-to-image search.
+  - **Hub Integration:** Detailed instructions for uploading the trained model and hosting the app on the Hugging Face Hub.
+
+## Project Structure
+your-clip-project/
+│
+├── data/
+│   ├── images/
+│   └── captions.txt
+│
+├── app.py
+├── config.py
+├── dataset.py
+├── model.py
+├── train.py
+│
+├── requirements.txt
+└── README.md
+
+
+## Setup and Installation
+
+**1. Clone the Repository:**
+```bash
+git clone <your-repo-url>
+cd your-clip-project
+2. Create a Python Virtual Environment:
+
+Bash
+
+python -m venv venv
+source venv/bin/activate  # On Windows, use `venv\Scripts\activate`
+3. Install Dependencies:
+
+Bash
+
+pip install -r requirements.txt
+4. Download the Flickr8k Dataset:
+
+Request the dataset from the official source: https://illinois.edu/fb/sec/1713398.
+
+Download and extract Flickr8k_Dataset.zip into the data/images/ folder.
+
+Find a captions.txt file (commonly available on Kaggle versions of the dataset) and place it at data/captions.txt.
+
+How to Run
+Step 1: Train the Model
+First, you must train the model. This will create a clip_book_model.pth file containing the learned weights of the projection heads.
+
+Run the training script from your terminal:
+
+Bash
+
+python train.py
+Step 2: Launch the Web Application
+Once the model is trained, launch the interactive search engine with Streamlit:
+
+Bash
+
+streamlit run app.py
+This will open a new tab in your browser with the application running.

app.py

@@ -0,0 +1,95 @@
+import gradio as gr
+import torch
+import torch.nn.functional as F
+from PIL import Image
+import requests
+from io import BytesIO
+from torchvision import transforms
+from transformers import DistilBertTokenizer
+import config # Your config file
+from inference_model import CLIPModel # Your model class file
+
+# --- 1. Load Model and Tokenizer (runs only once) ---
+# This section loads your trained model and tokenizer when the app starts.
+device = config.DEVICE
+
+# Load model with dimensions from config
+model = CLIPModel(
+    image_embedding_dim=config.IMAGE_EMBEDDING_DIM,
+    text_embedding_dim=config.TEXT_EMBEDDING_DIM,
+    projection_dim=config.PROJECTION_DIM
+).to(device)
+
+# Load the trained model weights from your .pth file
+model.load_state_dict(torch.load(config.MODEL_PATH, map_location=device))
+model.eval()
+
+# Load tokenizer
+tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
+print("Model and Tokenizer loaded successfully.")
+
+# --- 2. Image Preprocessing Function (reused from your code) ---
+def preprocess_image(image):
+    """Preprocess the image for the model."""
+    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 transform(image).unsqueeze(0)
+
+# --- 3. The Main Gradio Function ---
+# This is the core function that Gradio will build a UI around.
+# It takes the inputs from the UI and returns the outputs to the UI.
+def find_best_match(image_input, text_queries_input):
+    """
+    Takes an image and a block of text queries, and returns a dictionary
+    of queries and their similarity scores.
+    """
+    if image_input is None:
+        return "Please provide an image."
+    if not text_queries_input:
+        return "Please provide text descriptions."
+
+    # Process the image to get a tensor
+    image_tensor = preprocess_image(image_input).to(device)
+
+    # Process the text queries into a clean list
+    queries = [q.strip() for q in text_queries_input.split('\n') if q.strip()]
+    if not queries:
+        return "Please provide valid text descriptions."
+
+    # Process the text queries to get tokens
+    text_inputs = tokenizer(queries, padding=True, truncation=True, return_tensors="pt").to(device)
+
+    # Get model embeddings
+    with torch.no_grad():
+        image_embedding, text_embeddings = model(
+            image_features=image_tensor,
+            text_input_ids=text_inputs['input_ids'],
+            text_attention_mask=text_inputs['attention_mask']
+        )
+
+    # Calculate cosine similarity and format for Gradio's Label component
+    image_embedding_norm = F.normalize(image_embedding, p=2, dim=-1)
+    text_embeddings_norm = F.normalize(text_embeddings, p=2, dim=-1)
+    similarity_scores = (image_embedding_norm @ text_embeddings_norm.T).squeeze(0)
+    
+    # Create a results dictionary: { "query text": score, ... }
+    results = {query: score.item() for query, score in zip(queries, similarity_scores)}
+    
+    return results
+
+# --- 4. Create and Launch the Gradio Interface ---
+iface = gr.Interface(
+    fn=find_best_match,
+    inputs=[
+        gr.Image(type="pil", label="Upload or Drag an Image"),
+        gr.Textbox(lines=5, label="Text Descriptions (one per line)", placeholder="a person on a beach\na black cat\na city skyline at night")
+    ],
+    outputs=gr.Label(num_top_classes=5, label="Results"),
+    title="🖼️ CLIP Image-Text Search",
+    description="Provide an image and several text descriptions. The app will use a trained CLIP model to find the best textual match for the image."
+)
+
+iface.launch()

clip_book_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e372639e3a5bfa25a166d5825293e682c684628e1d95eccb6da17fbbac82b522
+size 363273830

config.py

@@ -0,0 +1,20 @@
+# config.py
+import torch
+
+# --- Project Paths ---
+IMAGE_DIR = "data/Flicker8k_Dataset"
+CAPTION_FILE = "data/captions.txt"
+MODEL_PATH = "clip_book_model.pth"
+
+# --- Model Dimensions ---
+IMAGE_EMBEDDING_DIM = 2048  # ResNet50 output dimension
+TEXT_EMBEDDING_DIM = 768    # DistilBERT output dimension
+PROJECTION_DIM = 256        # Shared embedding space dimension
+
+# --- Training Parameters ---
+BATCH_SIZE = 32
+NUM_EPOCHS = 3
+LEARNING_RATE = 1e-3
+
+# --- System ---
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

dataset.py

@@ -0,0 +1,51 @@
+# dataset.py
+import os
+import torch
+from torch.utils.data import Dataset
+from PIL import Image
+import pandas as pd
+import torchvision.transforms as transforms
+
+class Flickr8kDataset(Dataset):
+    """
+    Custom PyTorch Dataset for the Flickr8k data.
+    It loads images and their corresponding captions, tokenizing the text
+    on initialization for efficiency.
+    """
+    def __init__(self, image_dir, caption_file, tokenizer):
+        self.image_dir = image_dir
+        self.tokenizer = tokenizer
+
+        self.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]),
+        ])
+
+        df = pd.read_csv(caption_file)
+        self.image_paths = [os.path.join(self.image_dir, fname) for fname in df['image']]
+        self.captions = df['caption'].tolist()
+
+        print("Tokenizing all captions... (This may take a moment)")
+        self.caption_encodings = self.tokenizer(
+            self.captions,
+            truncation=True,
+            padding='max_length',
+            max_length=200,
+            return_tensors="pt"
+        )
+        print("Tokenization complete.")
+
+    def __len__(self):
+        return len(self.captions)
+
+    def __getitem__(self, idx):
+        item = {key: val[idx] for key, val in self.caption_encodings.items()}
+        try:
+            img = Image.open(self.image_paths[idx]).convert("RGB")
+            item['image'] = self.transform(img)
+        except (FileNotFoundError):
+            print(f"Warning: Could not load image at {self.image_paths[idx]}. Returning a black image.")
+            item['image'] = torch.zeros((3, 224, 224))
+        item["caption_text"] = self.captions[idx]
+        return item

inference_model.py

@@ -0,0 +1,73 @@
+# inference_model.py
+import torch
+import torch.nn as nn
+from torchvision.models import resnet50
+from transformers import DistilBertModel
+
+# --- Copy these classes from your original file ---
+class VisionEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        # Note: Using the newer 'weights' parameter is recommended
+        pretrained_resnet50 = resnet50(weights='IMAGENET1K_V1')
+        self.model = nn.Sequential(*list(pretrained_resnet50.children())[:-1])
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+    def forward(self, x):
+        x = self.model(x)
+        return x.view(x.size(0), -1)
+
+class TextEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = DistilBertModel.from_pretrained('distilbert-base-uncased')
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+    def forward(self, input_ids, attention_mask=None):
+        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
+        return outputs.last_hidden_state[:, 0, :]
+
+class ProjectionHead(nn.Module):
+    def __init__(self, embedding_dim, projection_dim=256, dropout=0.1):
+        super().__init__()
+        self.projection = nn.Linear(embedding_dim, projection_dim)
+        self.gelu = nn.GELU()
+        self.fc = nn.Linear(projection_dim, projection_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.layer_norm = nn.LayerNorm(projection_dim)
+
+    def forward(self, x):
+        projected = self.projection(x)
+        x = self.gelu(projected)
+        x = self.fc(x)
+        x = self.dropout(x)
+        x = x + projected
+        x = self.layer_norm(x)
+        return x
+
+# --- This is the MODIFIED CLIPModel for inference ---
+class CLIPModel(nn.Module):
+    def __init__(self, image_embedding_dim, text_embedding_dim, projection_dim):
+        super().__init__()
+        self.vision_encoder = VisionEncoder()
+        self.text_encoder = TextEncoder()
+        self.image_projection = ProjectionHead(embedding_dim=image_embedding_dim, projection_dim=projection_dim)
+        self.text_projection = ProjectionHead(embedding_dim=text_embedding_dim, projection_dim=projection_dim)
+
+    def forward(self, image_features=None, text_input_ids=None, text_attention_mask=None):
+        image_embedding = None
+        if image_features is not None:
+            image_features = self.vision_encoder(image_features)
+            image_embedding = self.image_projection(image_features)
+
+        text_embedding = None
+        if text_input_ids is not None:
+            text_features = self.text_encoder(
+                input_ids=text_input_ids,
+                attention_mask=text_attention_mask
+            )
+            text_embedding = self.text_projection(text_features)
+
+        return image_embedding, text_embedding

requirements.txt

@@ -0,0 +1,7 @@
+gradio
+torch
+torchvision
+transformers
+Pillow
+requests
+huggingface-hub

train.py

@@ -0,0 +1,54 @@
+# train.py
+import torch
+from torch.utils.data import DataLoader
+from transformers import DistilBertTokenizer
+from tqdm import tqdm
+
+import config
+from dataset import Flickr8kDataset
+from model import CLIPModel
+
+def main():
+    print(f"--- Starting Training ---")
+    print(f"Using device: {config.DEVICE}")
+
+    tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
+    dataset = Flickr8kDataset(config.IMAGE_DIR, config.CAPTION_FILE, tokenizer)
+    dataloader = DataLoader(dataset, batch_size=config.BATCH_SIZE, shuffle=True)
+
+    model = CLIPModel(
+        image_embedding_dim=config.IMAGE_EMBEDDING_DIM,
+        text_embedding_dim=config.TEXT_EMBEDDING_DIM,
+        projection_dim=config.PROJECTION_DIM
+    ).to(config.DEVICE)
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=config.LEARNING_RATE)
+
+    for epoch in range(config.NUM_EPOCHS):
+        print(f"\nEpoch {epoch+1}/{config.NUM_EPOCHS}")
+        model.train()
+        total_loss = 0
+        progress_bar = tqdm(dataloader, desc=f"Training Epoch {epoch+1}")
+
+        for batch in progress_bar:
+            batch = {k: v.to(config.DEVICE) for k, v in batch.items() if isinstance(v, torch.Tensor)}
+            
+            loss = model(batch)
+            
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            
+            total_loss += loss.item()
+            progress_bar.set_postfix({"loss": f"{loss.item():.4f}"})
+
+        avg_loss = total_loss / len(dataloader)
+        print(f"Epoch {epoch+1} - Average Loss: {avg_loss:.4f}")
+
+    print("\nTraining complete.")
+    torch.save(model.state_dict(), config.MODEL_PATH)
+    print(f"Model saved to {config.MODEL_PATH}")
+    print("\nTo upload to Hugging Face Hub, run the upload_to_hub.py script.")
+
+if __name__ == '__main__':
+    main()

training_model.py

@@ -0,0 +1,73 @@
+# model.py
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models import resnet50
+from transformers import DistilBertModel
+
+class VisionEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        pretrained_resnet50 = resnet50(weights='IMAGENET1K_V1')
+        self.model = nn.Sequential(*list(pretrained_resnet50.children())[:-1])
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+    def forward(self, x):
+        x = self.model(x)
+        return x.view(x.size(0), -1)
+
+class TextEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = DistilBertModel.from_pretrained('distilbert-base-uncased')
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+    def forward(self, input_ids, attention_mask=None):
+        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
+        return outputs.last_hidden_state[:, 0, :]
+
+class ProjectionHead(nn.Module):
+    def __init__(self, embedding_dim, projection_dim=256, dropout=0.1):
+        super().__init__()
+        self.projection = nn.Linear(embedding_dim, projection_dim)
+        self.gelu = nn.GELU()
+        self.fc = nn.Linear(projection_dim, projection_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.layer_norm = nn.LayerNorm(projection_dim)
+
+    def forward(self, x):
+        projected = self.projection(x)
+        x = self.gelu(projected)
+        x = self.fc(x)
+        x = self.dropout(x)
+        x = x + projected
+        x = self.layer_norm(x)
+        return x
+
+class CLIPModel(nn.Module):
+    def __init__(self, image_embedding_dim, text_embedding_dim, projection_dim):
+        super().__init__()
+        self.vision_encoder = VisionEncoder()
+        self.text_encoder = TextEncoder()
+        self.image_projection = ProjectionHead(embedding_dim=image_embedding_dim, projection_dim=projection_dim)
+        self.text_projection = ProjectionHead(embedding_dim=text_embedding_dim, projection_dim=projection_dim)
+
+    def forward(self, batch):
+        image_features = self.vision_encoder(batch["image"])
+        text_features = self.text_encoder(
+            input_ids=batch["input_ids"],
+            attention_mask=batch["attention_mask"]
+        )
+        image_embeddings = self.image_projection(image_features)
+        text_embeddings = self.text_projection(text_features)
+
+        # Textbook's specific loss calculation
+        logits = text_embeddings @ image_embeddings.T
+        images_similarity = image_embeddings @ image_embeddings.T
+        texts_similarity = text_embeddings @ text_embeddings.T
+        targets = F.softmax((images_similarity + texts_similarity) / 2, dim=-1)
+        texts_loss = F.cross_entropy(logits, targets)
+        images_loss = F.cross_entropy(logits.T, targets.T)
+        return (images_loss + texts_loss) / 2.0