Initial deployment

app.py

@@ -0,0 +1,117 @@
+import gradio as gr
+import torch
+import clip
+from torchvision import transforms
+from style_net import StyleNet
+from dataset import FlickrStreamer
+
+# --- CONFIGURATION ---
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+INDEX_PATH = "flickr_embeddings.pt"
+SEARCH_LIMIT = 5000  # Must match the indexer limit!
+# ---------------------
+
+print(f"System starting on {DEVICE}...")
+
+# 1. LOAD MODELS
+# A. Editing Model
+style_net = StyleNet().to(DEVICE)
+
+# B. Retrieval Model (CLIP)
+clip_model, preprocess = clip.load("ViT-B/32", device=DEVICE)
+
+# 2. LOAD SEARCH INDEX
+print("Loading Search Index...")
+try:
+    image_features = torch.load(INDEX_PATH).to(DEVICE)
+    print(f"Index loaded: {image_features.shape[0]} images.")
+except FileNotFoundError:
+    print("WARNING: 'flickr_embeddings.pt' not found. Run indexer.py first!")
+    image_features = None
+
+# 3. CONNECT TO DATASET (For fetching the actual images)
+# Note: No transform here because we want to display the original PIL image
+dataset_raw = FlickrStreamer(limit=SEARCH_LIMIT, transform=None)
+
+# --- FUNCTION: SEARCH ---
+def search_flickr(text_query):
+    if image_features is None:
+        return []
+    
+    with torch.no_grad():
+        # 1. Encode Text
+        text_tokens = clip.tokenize([text_query]).to(DEVICE)
+        text_features = clip_model.encode_text(text_tokens)
+        text_features /= text_features.norm(dim=-1, keepdim=True)
+        
+        # 2. Math: Cosine Similarity (Dot Product)
+        # (1, 512) x (5000, 512).T  -> (1, 5000) similarity scores
+        similarity = (100.0 * text_features @ image_features.T).softmax(dim=-1)
+        
+        # 3. Get Top 5 matches
+        values, indices = similarity[0].topk(5)
+        
+        # 4. Retrieve actual images from dataset
+        # (Deep Lake might be slightly slow here as it fetches specific indices)
+        results = []
+        for idx in indices:
+            img_idx = idx.item()
+            results.append(dataset_raw[img_idx])
+            
+        return results
+
+# --- FUNCTION: EDIT ---
+tf_edit = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.ToTensor()
+])
+
+def edit_image(image, style_choice):
+    if image is None: return None
+    
+    weights_map = {
+        "Sketch": "model_sketch.pth",
+        "Van Gogh": "model_painting.pth",
+        "Cyberpunk": "model_cyberpunk.pth"
+    }
+    
+    try:
+        style_net.load_state_dict(torch.load(weights_map.get(style_choice), map_location=DEVICE))
+    except Exception:
+        return image 
+        
+    input_tensor = tf_edit(image).unsqueeze(0).to(DEVICE)
+    with torch.no_grad():
+        output_tensor = style_net(input_tensor)
+        
+    output_img = transforms.ToPILImage()(output_tensor.squeeze().cpu())
+    return output_img
+
+# --- THE UI ---
+with gr.Blocks() as demo:
+    gr.Markdown("# Group 19: CLIP Search & Edit Engine")
+    
+    with gr.Tab("Retrieval"):
+        gr.Markdown(f"Search through {SEARCH_LIMIT} Flickr images by concept.")
+        with gr.Row():
+            txt_query = gr.Textbox(label="Type a query (e.g., 'a dog on a boat', 'neon lights')")
+            btn_search = gr.Button("Search Library")
+        
+        # Output is a Gallery (Grid of images)
+        gallery_out = gr.Gallery(label="Top Results", columns=5, height=300)
+        
+        btn_search.click(search_flickr, inputs=txt_query, outputs=gallery_out)
+        
+    with gr.Tab("Image Editing"):
+        gr.Markdown("Transform any image using our trained CLIPstyler models.")
+        with gr.Row():
+            im_in = gr.Image(type="pil", label="Input Image")
+            style = gr.Dropdown(["Sketch", "Van Gogh", "Cyberpunk"], label="Choose Style")
+        
+        btn_edit = gr.Button("Apply")
+        im_out = gr.Image(label="Result")
+        
+        btn_edit.click(edit_image, inputs=[im_in, style], outputs=im_out)
+
+if __name__ == "__main__":
+    demo.launch(share=True)

dataset.py

@@ -0,0 +1,31 @@
+import deeplake
+import torch
+from PIL import Image
+import numpy as np
+from torch.utils.data import Dataset
+
+class FlickrStreamer(Dataset):
+    def __init__(self, limit=1000, transform=None):
+        # Streams data from cloud. No massive download needed.
+        print(f"Connecting to Flickr30k (First {limit} images)...")
+        # Windows sometimes has permission issues with cache, so we explicitly set access_method
+        self.ds = deeplake.load('hub://activeloop/flickr30k', access_method="stream")
+        self.limit = limit
+        self.transform = transform
+
+    def __len__(self):
+        return self.limit
+
+    def __getitem__(self, idx):
+        sample = self.ds[idx]
+        image_data = sample.image.numpy()
+        
+        try:
+            image = Image.fromarray(image_data)
+        except:
+            image = Image.fromarray(image_data.astype('uint8'))
+            
+        if self.transform:
+            image = self.transform(image)
+            
+        return image

flickr_embeddings.pt

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

model_cyberpunk.pth

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

model_painting.pth

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

model_sketch.pth

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

requirements.txt

@@ -0,0 +1,8 @@
+torch
+torchvision
+ftfy
+regex
+tqdm
+deeplake
+gradio
+git+https://github.com/openai/CLIP.git

style_net.py

@@ -0,0 +1,63 @@
+import torch
+import torch.nn as nn
+
+class ResidualBlock(nn.Module):
+    def __init__(self, channels):
+        super(ResidualBlock, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(channels, channels, 3, padding=1),
+            nn.InstanceNorm2d(channels),
+            nn.ReLU(),
+            nn.Conv2d(channels, channels, 3, padding=1),
+            nn.InstanceNorm2d(channels)
+        )
+
+    def forward(self, x):
+        return x + self.conv(x)
+    
+class UpsampleConvLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride):
+        super().__init__()
+        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        # Reflection padding reduces "border" artifacts
+        self.reflection_pad = nn.ReflectionPad2d(kernel_size // 2)
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=1)
+
+    def forward(self, x):
+        x = self.upsample(x)
+        x = self.reflection_pad(x)
+        return self.conv(x)
+
+class StyleNet(nn.Module):
+    def __init__(self):
+        super(StyleNet, self).__init__()
+        # Encoder
+        self.encoder = nn.Sequential(
+            nn.ReflectionPad2d(4),           # Pad first...
+            nn.Conv2d(3, 32, 9, padding=0),  # ...then Conv (padding=0)
+            nn.InstanceNorm2d(32), nn.ReLU(),
+            
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(32, 64, 3, stride=2, padding=0),
+            nn.InstanceNorm2d(64), nn.ReLU(),
+            
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(64, 128, 3, stride=2, padding=0),
+            nn.InstanceNorm2d(128), nn.ReLU()
+        )
+        # Bottleneck
+        self.bottleneck = nn.Sequential(*[ResidualBlock(128) for _ in range(3)])
+        # Decoder
+        self.decoder = nn.Sequential(
+            UpsampleConvLayer(128, 64, kernel_size=3, stride=1),
+            nn.InstanceNorm2d(64), nn.ReLU(),
+            UpsampleConvLayer(64, 32, kernel_size=3, stride=1),
+            nn.InstanceNorm2d(32), nn.ReLU(),
+            nn.Conv2d(32, 3, 9, padding=4),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        features = self.encoder(x)
+        features = self.bottleneck(features)
+        return self.decoder(features)