Spaces:
Sleeping
Sleeping
File size: 4,443 Bytes
98b6c0d ff63d1b 98b6c0d ff63d1b 98b6c0d ff63d1b 98b6c0d ff63d1b 98b6c0d ff63d1b 5a14c00 ff63d1b 98b6c0d ff63d1b 98b6c0d ff63d1b 98b6c0d ff63d1b 98b6c0d ff63d1b 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 129 130 131 132 | import gradio as gr
import torch
import clip
from torchvision import transforms
from style_net import StyleNet
from dataset import FlickrStreamer
from clip_styler import style_transfer
# --- 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, 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
# --- FUNCTION: EDIT ---
tf_edit = transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor()
])
def edit_image(image, style_choice, num_steps):
if image is None: return None
weights_map = {
"Sketch": "model_sketch.pth",
"Van Gogh": "model_painting.pth",
"Cyberpunk": "model_cyberpunk.pth"
}
detail_prompt = {
"Sketch": "charcoal sketch",
"Van Gogh": "Van Gogh painting",
"Cyberpunk": "neon cyberpunk",
}
'''
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
'''
return style_transfer(image, detail_prompt[style_choice], num_steps=num_steps)
# --- 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")
steps_slider = gr.Slider(
minimum=50,
maximum=200,
value=100,
step=10,
label="Number of Steps (more = better quality but slower)"
)
btn_edit = gr.Button("Apply")
im_out = gr.Image(label="Result")
btn_edit.click(edit_image, inputs=[im_in, style, steps_slider], outputs=im_out)
if __name__ == "__main__":
demo.launch(share=True) |