Spaces:
Sleeping
Sleeping
File size: 4,129 Bytes
98b6c0d 3601849 98b6c0d 3601849 98b6c0d 3601849 98b6c0d 3601849 98b6c0d 3601849 98b6c0d 3601849 98b6c0d 3601849 98b6c0d 442d9ae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | import gradio as gr
import torch
import clip
from torchvision import transforms
from style_net import StyleNet
from dataset import FlickrStreamer
from util import DEVICE
from clip_styler import ClipStyler
# --- CONFIGURATION ---
#DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
INDEX_PATH = "flickr_embeddings.pt"
SEARCH_LIMIT = 5000 # Must match the indexer limit!
# ---------------------
UI_EPOCHS = 100
STYLE_PRESETS = {
"Sketch": {
"prompt": "charcoal sketch style",
"source": "a Photo",
},
"Van Gogh": {
"prompt": "Van Gogh painting style",
"source": "a Photo",
},
"Cyberpunk": {
"prompt": "neon cyberpunk style",
"source": "a Photo",
},
}
print(f"System starting on {DEVICE}...")
# 1. LOAD MODELS
# A. Retrieval Model (CLIP)
clip_model, _ = clip.load("ViT-B/32", device=DEVICE)
# B. Editing Model (CLIP-guided optimizer)
clip_styler = ClipStyler(device=DEVICE, num_epochs=UI_EPOCHS)
# 2. LOAD SEARCH INDEX
print("Loading Search Index...")
try:
image_features = torch.load(INDEX_PATH, map_location=DEVICE).float()
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
def edit_image(image, style_choice):
if image is None or style_choice is None:
return None
preset = STYLE_PRESETS.get(style_choice)
if preset is None:
return image
try:
return clip_styler.stylize(
image=image,
prompt=preset["prompt"],
source=preset.get("source", "a Photo"),
num_epochs=UI_EPOCHS,
log_interval=0,
)
except Exception as exc:
print(f"[edit_image] Failed to stylize image: {exc}")
return image
# --- 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) |