Upload 2 files

app.py

@@ -0,0 +1,94 @@
+import gradio as gr
+import clip_chat
+
+
+def logit2sentence(logit, slider_value):
+    sentence = ""
+    if logit < slider_value / 2.5:
+        sentence = "Nope. Not at all."
+    elif slider_value / 2.5 < logit < slider_value / 1.56:
+        sentence = "Not really..."
+    elif slider_value / 1.56 < logit < slider_value / 1.36:
+        sentence = "Close but not there."
+    elif slider_value / 1.36 < logit < slider_value / 1.14:
+        sentence = "That's quite close."
+    elif slider_value / 1.14 < logit < slider_value:
+        sentence = "Almost guessed."
+    elif logit >= slider_value:
+        sentence = "YES!!"
+    return sentence
+
+
+def give_up():
+    image = clip_chat.image_org
+    return image, None, "You lost... (Press \"Reset\" to play again)"
+
+
+def update_difficulty(x):
+    if not has_started:
+        clip_chat.goal = x
+        return clip_chat.goal
+    return clip_chat.goal
+
+
+has_started = False
+best_guess = None
+
+
+def respond(message, chat_history, label_value, image_value):
+    global has_started, best_guess
+
+    if not has_started:
+        has_started = True
+
+    logits, is_better = clip_chat.answer(message)
+    bot_message = logit2sentence(logits, clip_chat.goal)
+
+    if is_better == 3:
+        best_guess = {f"Best Guess: \"{message}\"": float(logits) / clip_chat.goal}
+    if float(logits) >= clip_chat.goal:
+        bot_message = "YES!"
+        best_guess = "YOU WIN! (Press \"Reset\" to play again)"
+        image_value = clip_chat.image_org
+    else:
+        if is_better == -1:
+            bot_message += ""
+        elif is_better == 0:
+            bot_message += "You did worse than the last one."
+        elif is_better == 1 or is_better == 3:
+            bot_message += "You did better than the last one."
+
+    label_value = best_guess
+
+    chat_history.append((message, bot_message))
+    return "", chat_history, label_value, image_value
+
+
+def reset_everything():
+    global has_started, best_guess
+    clip_chat.reset_everything()
+    has_started = False
+    best_guess = None
+    return clip_chat.goal, None, "This is a \"Guess the Image\" game. I'm thinking of a picture and you have to guess using the chat above.", None
+
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            chatbot = gr.Chatbot()
+            msg = gr.Textbox()
+            slider = gr.inputs.Slider(minimum=18, maximum=27, default=21, label="Difficulty (18 - Easy, 27 - Expert)")
+        with gr.Column():
+            label = gr.Label("This is a \"Guess the Image\" game. I'm thinking of a picture and you have to guess using the chat above.")
+            image_output = gr.outputs.Image(type="pil")
+            show_image_button = gr.Button("Give Up...")
+            reset_button = gr.Button("Reset")
+
+    msg.submit(respond, [msg, chatbot], [msg, chatbot, label, image_output])
+    slider.release(update_difficulty, inputs=[slider], outputs=[slider])
+    show_image_button.click(give_up, outputs=[image_output, chatbot, label], queue=False)
+    reset_button.click(reset_everything, outputs=[slider, image_output, label, chatbot], queue=False)
+
+if __name__ == "__main__":
+    demo.title = "CLIP Guess the Image"
+    demo.launch(share=False)

clip_chat.py

@@ -0,0 +1,82 @@
+import torch
+import clip
+from PIL import Image
+import glob
+import os
+from random import choice
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model, preprocess = clip.load("ViT-L/14@336px", device=device)
+
+COCO = glob.glob(os.path.join(os.getcwd(), "images", "*"))
+available_models = ['RN50', 'RN101', 'RN50x4', 'RN50x16', 'RN50x64', 'ViT-B/32', 'ViT-B/16', 'ViT-L/14', 'ViT-L/14@336px']
+
+
+def load_random_image():
+    image_path = choice(COCO)
+    image = Image.open(image_path)
+    return image
+
+
+def next_image():
+    global image_org, image
+    image_org = load_random_image()
+    image = preprocess(Image.fromarray(image_org)).unsqueeze(0).to(device)
+
+
+def calculate_logits(image_features, text_features):
+    image_features = image_features / image_features.norm(dim=1, keepdim=True)
+    text_features = text_features / text_features.norm(dim=1, keepdim=True)
+
+    logit_scale = model.logit_scale.exp()
+    return logit_scale * image_features @ text_features.t()
+
+
+last = -1
+best = -1
+
+goal = 21
+
+image_org = load_random_image()
+image = preprocess(image_org).unsqueeze(0).to(device)
+with torch.no_grad():
+    image_features = model.encode_image(image)
+
+
+def answer(message):
+    global last, best
+
+    text = clip.tokenize([message]).to(device)
+
+    with torch.no_grad():
+        text_features = model.encode_text(text)
+        logits_per_image, _ = model(image, text)
+        logits = calculate_logits(image_features, text_features).cpu().numpy().flatten()[0]
+
+    if last == -1:
+        is_better = -1
+    elif last > logits:
+        is_better = 0
+    elif last < logits:
+        is_better = 1
+    elif logits > goal:
+        is_better = 2
+    else:
+        is_better = -1
+
+    last = logits
+    if logits > best:
+        best = logits
+        is_better = 3
+
+    return logits, is_better
+
+
+def reset_everything():
+    global last, best, goal, image, image_org
+    last = -1
+    best = -1
+    goal = 21
+    image_org = load_random_image()
+    image = preprocess(image_org).unsqueeze(0).to(device)