Spaces:
Runtime error
Runtime error
import gradio as gr | |
from response_db import StResponseDb | |
from create_cache import Game_Cache | |
import numpy as np | |
from PIL import Image | |
import pandas as pd | |
import torch | |
import pickle | |
import uuid | |
db = StResponseDb() | |
css = """ | |
.chatbot {display:flex;flex-direction:column} | |
.msg {padding:4px;margin-bottom:4px;border-radius:4px;width:80%} | |
.msg.user {background-color:cornflowerblue;color:white;align-self:self-end} | |
.msg.bot {background-color:lightgray} | |
.na_button {background-color:red;color:red} | |
""" | |
from model.run_question_asking_model import return_modules, return_modules_yn | |
question_model, response_model_simul, _, caption_model = return_modules() | |
question_model_yn, response_model_simul_yn, _, caption_model_yn = return_modules_yn() | |
class Game_Session: | |
def __init__(self, taskid, yn, hard_setting): | |
self.yn = yn | |
self.hard_setting = hard_setting | |
global question_model, response_model_simul, caption_model | |
global question_model_yn, response_model_simul_yn, caption_model_yn | |
self.question_model = question_model | |
self.response_model_simul = response_model_simul | |
self.caption_model = caption_model | |
self.question_model_yn = question_model_yn | |
self.response_model_simul_yn = response_model_simul_yn | |
self.caption_model_yn = caption_model_yn | |
global image_files, images_np, p_y_x, p_r_qy, p_y_xqr, captions, questions, target_questions | |
self.image_files, self.image_np, self.p_y_x, self.p_r_qy, self.p_y_xqr = None, None, None, None, None | |
self.captions, self.questions, self.target_questions = None, None, None | |
self.history = [] | |
self.game_id = str(uuid.uuid4()) | |
self.set_curr_models() | |
def set_curr_models(self): | |
if self.yn: | |
self.curr_question_model, self.curr_caption_model, self.curr_response_model_simul = self.question_model_yn, self.caption_model_yn, self.response_model_simul_yn | |
else: | |
self.curr_question_model, self.curr_caption_model, self.curr_response_model_simul = self.question_model, self.caption_model, self.response_model_simul | |
def get_next_question(self): | |
return self.curr_question_model.select_best_question(self.p_y_x, self.questions, self.images_np, self.captions, self.curr_response_model_simul) | |
def ask_a_question(input, taskid, gs): | |
gs.history.append(input) | |
gs.p_r_qy = gs.curr_response_model_simul.get_p_r_qy(input, gs.history[-2], gs.images_np, gs.captions) | |
gs.p_y_xqr = gs.p_y_x*gs.p_r_qy | |
gs.p_y_xqr = gs.p_y_xqr/torch.sum(gs.p_y_xqr)if torch.sum(gs.p_y_xqr) != 0 else torch.zeros_like(gs.p_y_xqr) | |
gs.p_y_x = gs.p_y_xqr | |
gs.questions.remove(gs.history[-2]) | |
db.add(gs.game_id, taskid, len(gs.history)//2-1, gs.history[-2], gs.history[-1]) | |
gs.history.append(gs.get_next_question()) | |
top_prob = torch.max(gs.p_y_x).item() | |
top_pred = torch.argmax(gs.p_y_x).item() | |
if top_prob > 0.8: | |
gs.history = gs.history[:-1] | |
db.add(gs.game_id, taskid, len(gs.history)//2, f"Guess: Image {top_pred}", "") | |
# write some HTML | |
html = "<div class='chatbot'>" | |
for m, msg in enumerate(gs.history): | |
if msg=="nothing": msg="n/a" | |
cls = "bot" if m%2 == 0 else "user" | |
html += "<div class='msg {}'> {}</div>".format(cls, msg) | |
html += "</div>" | |
### Game finished: | |
if top_prob > 0.8: | |
html += f"<p>The model identified <b>Image {top_pred+1}</b> as the image. Please select a new task ID to continue.</p>" | |
return html, gs, gr.Textbox.update(visible=False), gr.Button.update(visible=False), gr.Button.update(visible=False), gr.Number.update(visible=True), gr.Button.update(visible=True), gr.Number.update(visible=False), gr.Button.update(visible=False) | |
else: | |
if not gs.yn: | |
return html, gs, gr.Textbox.update(visible=True), gr.Button.update(visible=True), gr.Button.update(visible=True), gr.Number.update(visible=False), gr.Button.update(visible=False), gr.Number.update(visible=False), gr.Button.update(visible=False) | |
else: | |
return html, gs, gr.Textbox.update(visible=False), gr.Button.update(visible=False), gr.Button.update(visible=False), gr.Number.update(visible=False), gr.Button.update(visible=False), gr.Number.update(visible=True), gr.Button.update(visible=True) | |
def set_images(taskid): | |
pilot_study = pd.read_csv("pilot-study.csv") | |
taskid_original = taskid | |
taskid = pilot_study['mscoco-id'].tolist()[int(taskid)] | |
with open(f'cache/{int(taskid)}.p', 'rb') as fp: | |
game_cache = pickle.load(fp) | |
gs = Game_Session(int(taskid), game_cache.yn, game_cache.hard_setting) | |
id1 = f"./mscoco-images/val2014/{game_cache.image_files[0]}" | |
id2 = f"./mscoco-images/val2014/{game_cache.image_files[1]}" | |
id3 = f"./mscoco-images/val2014/{game_cache.image_files[2]}" | |
id4 = f"./mscoco-images/val2014/{game_cache.image_files[3]}" | |
id5 = f"./mscoco-images/val2014/{game_cache.image_files[4]}" | |
id6 = f"./mscoco-images/val2014/{game_cache.image_files[5]}" | |
id7 = f"./mscoco-images/val2014/{game_cache.image_files[6]}" | |
id8 = f"./mscoco-images/val2014/{game_cache.image_files[7]}" | |
id9 = f"./mscoco-images/val2014/{game_cache.image_files[8]}" | |
id10 = f"./mscoco-images/val2014/{game_cache.image_files[9]}" | |
gs.image_files = [id1, id2, id3, id4, id5, id6, id7, id8, id9, id10] | |
gs.image_files = [x[15:] for x in gs.image_files] | |
gs.images_np = [np.asarray(Image.open(f"./mscoco-images/{i}")) for i in gs.image_files] | |
gs.images_np = [np.dstack([i]*3) if len(i.shape)==2 else i for i in gs.images_np] | |
gs.p_y_x = (torch.ones(10)/10).to(gs.curr_question_model.device) | |
gs.captions = gs.curr_caption_model.get_captions(gs.image_files) | |
gs.questions, gs.target_questions = gs.curr_question_model.get_questions(gs.image_files, gs.captions, 0) | |
gs.curr_question_model.reset_question_bank() | |
gs.curr_question_model.question_bank = game_cache.question_dict | |
first_question = gs.curr_question_model.select_best_question(gs.p_y_x, gs.questions, gs.images_np, gs.captions, gs.curr_response_model_simul) | |
first_question_html = f"<div class='chatbot'><div class='msg bot'>{first_question}</div></div>" | |
gs.history.append(first_question) | |
html = f"<p>Current Task ID: <b>{int(taskid_original)}</b></p>" | |
if not gs.yn: | |
return id1, id2, id3, id4, id5, id6, id7, id8, id9, id10, gs, first_question_html, gr.HTML.update(value=html, visible=True), gr.Textbox.update(visible=True, value=''), gr.Button.update(visible=True), gr.Button.update(visible=True), gr.Number.update(visible=False), gr.Button.update(visible=False), gr.Button.update(visible=False), gr.Button.update(visible=False) | |
else: | |
return id1, id2, id3, id4, id5, id6, id7, id8, id9, id10, gs, first_question_html, gr.HTML.update(value=html, visible=True), gr.Textbox.update(visible=False), gr.Button.update(visible=False), gr.Button.update(visible=False), gr.Number.update(visible=False), gr.Button.update(visible=False), gr.Button.update(visible=True), gr.Button.update(visible=True) | |
with gr.Blocks(title="Image Q&A Guessing Game", css=css) as demo: | |
gr.HTML("<h1>Image Q&A Guessing Game</h1>\ | |
<p style='font-size:120%;'>\ | |
Imagine you are playing 20-questions with an AI model.<br>\ | |
The AI model plays the role of the question asker. You play the role of the responder. <br>\ | |
There are 10 images. <b>Your image is Image 1</b>. The other images are distraction images.\ | |
The model can see all 10 images and all the questions and answers for the current set of images. It will ask a question based on the available information.<br>\ | |
<span style='color: #0000ff'>The goal of the model is to accurately guess the correct image (i.e. <b><span style='color: #0000ff'>Image 1</span></b>) in as few turns as possible.<br>\ | |
Your goal is to help the model guess the image by answering as clearly and accurately as possible.</span><br><br>\ | |
<b>Guidelines:</b><br>\ | |
<ol style='font-size:120%;'>\ | |
<li>It is best to keep your answers short (a single word or a short phrase). No need to answer in full sentences.</li>\ | |
<li>If you feel that the question cannot be answered or does not apply to Image 1, please select N/A.</li>\ | |
</ol> \ | |
<br>\ | |
(Note: We are testing multiple game settings. In some instances, the game will be open-ended, while in other instances, the answer choices will be limited to yes/no.)<br></p>\ | |
<br>\ | |
<h2>Please enter a TaskID to start</h2>") | |
with gr.Column(): | |
with gr.Row(): | |
taskid = gr.Number(label="Task ID (Enter a number from 0 to 160)", value=0) | |
start_button = gr.Button("Enter") | |
with gr.Row(): | |
task_text = gr.HTML() | |
with gr.Column() as img_block: | |
with gr.Row(): | |
img1 = gr.Image(label="Image 1", show_label=True) | |
img2 = gr.Image(label="Image 2", show_label=True) | |
img3 = gr.Image(label="Image 3", show_label=True) | |
img4 = gr.Image(label="Image 4", show_label=True) | |
img5 = gr.Image(label="Image 5", show_label=True) | |
with gr.Row(): | |
img6 = gr.Image(label="Image 6", show_label=True) | |
img7 = gr.Image(label="Image 7", show_label=True) | |
img8 = gr.Image(label="Image 8", show_label=True) | |
img9 = gr.Image(label="Image 9", show_label=True) | |
img10 = gr.Image(label="Image 10", show_label=True) | |
conversation = gr.HTML() | |
game_session_state = gr.State() | |
answer = gr.Textbox(placeholder="Insert answer here.", label="Answer the given question.", visible=False) | |
null_answer = gr.Textbox("nothing", visible=False) | |
yes_answer = gr.Textbox("yes", visible=False) | |
no_answer = gr.Textbox("no", visible=False) | |
with gr.Column(): | |
with gr.Row(): | |
yes_box = gr.Button("Yes", visible=False) | |
no_box = gr.Button("No", visible=False) | |
with gr.Column(): | |
with gr.Row(): | |
na_box = gr.Button("N/A", visible=False, elem_classes="na_button") | |
submit = gr.Button("Submit", visible=False) | |
### Button click events | |
start_button.click(fn=set_images, inputs=taskid, outputs=[img1, img2, img3, img4, img5, img6, img7, img8, img9, img10, game_session_state, conversation, task_text, answer, na_box, submit, taskid, start_button, yes_box, no_box]) | |
submit.click(fn=ask_a_question, inputs=[answer, taskid, game_session_state], outputs=[conversation, game_session_state, answer, na_box, submit, taskid, start_button, yes_box, no_box]) | |
na_box.click(fn=ask_a_question, inputs=[null_answer, taskid, game_session_state], outputs=[conversation, game_session_state, answer, na_box, submit, taskid, start_button, yes_box, no_box]) | |
yes_box.click(fn=ask_a_question, inputs=[yes_answer, taskid, game_session_state], outputs=[conversation, game_session_state, answer, na_box, submit, taskid, start_button, yes_box, no_box]) | |
no_box.click(fn=ask_a_question, inputs=[no_answer, taskid, game_session_state], outputs=[conversation, game_session_state, answer, na_box, submit, taskid, start_button, yes_box, no_box]) | |
demo.launch() | |