XAI/PEEB
7
1import os2import io3 4import torch5import json6import base647import gradio as gr8import numpy as np9from pathlib import Path10from PIL import Image11 12from plots import get_pre_define_colors13from utils.load_model import load_xclip14from utils.predict import xclip_pred15 16 17#! Huggingface does not allow load model to main process, so we need to load the model when needed, it may not help in improve the speed of the app.18DEVICE = "cuda" if torch.cuda.is_available() else "cpu"19print(f"Not at Huggingface demo, load model to main process.")20XCLIP, OWLVIT_PRECESSOR = load_xclip(DEVICE)21 22print(f"Device: {DEVICE}")23 24XCLIP_DESC_PATH = "data/jsons/bs_cub_desc.json"25XCLIP_DESC = json.load(open(XCLIP_DESC_PATH, "r"))26IMAGES_FOLDER = "data/images"27# XCLIP_RESULTS = json.load(open("data/jsons/xclip_org.json", "r"))28IMAGE2GT = json.load(open("data/jsons/image2gt.json", 'r'))29CUB_DESC_EMBEDS = torch.load('data/text_embeddings/cub_200_desc.pt')30CUB_IDX2NAME = json.load(open('data/jsons/cub_desc_idx2name.json', 'r'))31CUB_IDX2NAME = {int(k): v for k, v in CUB_IDX2NAME.items()}32 33IMAGE_FILE_LIST = json.load(open("data/jsons/file_list.json", "r"))34IMAGE_GALLERY = [Image.open(os.path.join(IMAGES_FOLDER, 'org', file_name)).convert('RGB') for file_name in IMAGE_FILE_LIST]35 36ORG_PART_ORDER = ['back', 'beak', 'belly', 'breast', 'crown', 'forehead', 'eyes', 'legs', 'wings', 'nape', 'tail', 'throat']37ORDERED_PARTS = ['crown', 'forehead', 'nape', 'eyes', 'beak', 'throat', 'breast', 'belly', 'back', 'wings', 'legs', 'tail']38COLORS = get_pre_define_colors(12, cmap_set=['Set2', 'tab10'])39SACHIT_COLOR = "#ADD8E6"40# CUB_BOXES = json.load(open("data/jsons/cub_boxes_owlvit_large.json", "r"))41VISIBILITY_DICT = json.load(open("data/jsons/cub_vis_dict_binary.json", 'r'))42VISIBILITY_DICT['Eastern_Bluebird.jpg'] = dict(zip(ORDERED_PARTS, [True]*12))43 44# --- Image related functions ---45def img_to_base64(img):46 img_pil = Image.fromarray(img) if isinstance(img, np.ndarray) else img47 buffered = io.BytesIO()48 img_pil.save(buffered, format="JPEG")49 img_str = base64.b64encode(buffered.getvalue())50 return img_str.decode()51 52def create_blank_image(width=500, height=500, color=(255, 255, 255)):53 """Create a blank image of the given size and color."""54 return np.array(Image.new("RGB", (width, height), color))55 56# Convert RGB colors to hex57def rgb_to_hex(rgb):58 return f"#{''.join(f'{x:02x}' for x in rgb)}"59 60def load_part_images(file_name: str) -> dict:61 part_images = {}62 # start_time = time.time()63 for part_name in ORDERED_PARTS:64 base_name = Path(file_name).stem65 part_image_path = os.path.join(IMAGES_FOLDER, "boxes", f"{base_name}_{part_name}.jpg")66 if not Path(part_image_path).exists():67 continue68 image = np.array(Image.open(part_image_path))69 part_images[part_name] = img_to_base64(image)70 # print(f"Time cost to load 12 images: {time.time() - start_time}")71 # This takes less than 0.01 seconds. So the loading time is not the bottleneck.72 return part_images73 74def generate_xclip_explanations(result_dict:dict, visibility: dict, part_mask: dict = dict(zip(ORDERED_PARTS, [1]*12))):75 """76 The result_dict needs three keys: 'descriptions', 'pred_scores', 'file_name'77 descriptions: {part_name1: desc_1, part_name2: desc_2, ...}78 pred_scores: {part_name1: score_1, part_name2: score_2, ...}79 file_name: str80 """81 82 descriptions = result_dict['descriptions']83 image_name = result_dict['file_name']84 part_images = PART_IMAGES_DICT[image_name]85 MAX_LENGTH = 5086 exp_length = 40087 fontsize = 1588 89 # Start the SVG inside a div90 svg_parts = [f'<div style="width: {exp_length}px; height: 450px; background-color: white;">',91 "<svg width=\"100%\" height=\"100%\">"]92 93 # Add a row for each visible bird part94 y_offset = 095 for part in ORDERED_PARTS:96 if visibility[part] and part_mask[part]:97 # Calculate the length of the bar (scaled to fit within the SVG)98 part_score = max(result_dict['pred_scores'][part], 0)99 bar_length = part_score * exp_length100 101 # Modify the overlay image's opacity on mouseover and mouseout102 mouseover_action1 = f"document.getElementById('overlayImage').src = 'data:image/jpeg;base64,{part_images[part]}'; document.getElementById('overlayImage').style.opacity = 1;"103 mouseout_action1 = "document.getElementById('overlayImage').style.opacity = 0;"104 105 combined_mouseover = f"javascript: {mouseover_action1};"106 combined_mouseout = f"javascript: {mouseout_action1};"107 108 # Add the description109 num_lines = len(descriptions[part]) // MAX_LENGTH + 1110 for line in range(num_lines):111 desc_line = descriptions[part][line*MAX_LENGTH:(line+1)*MAX_LENGTH]112 y_offset += fontsize113 svg_parts.append(f"""114 <text x="0" y="{y_offset}" font-size="{fontsize}" 115 onmouseover="{combined_mouseover}"116 onmouseout="{combined_mouseout}">117 {desc_line}118 </text>119 """)120 121 # Add the bars122 svg_parts.append(f"""123 <rect x="0" y="{y_offset +3}" width="{bar_length}" height="{fontsize*0.7}" fill="{PART_COLORS[part]}"124 onmouseover="{combined_mouseover}"125 onmouseout="{combined_mouseout}">126 </rect>127 """)128 # Add the scores129 svg_parts.append(f'<text x="{exp_length - 50}" y="{y_offset+fontsize+3}" font-size="{fontsize}" fill="{PART_COLORS[part]}">{part_score:.2f}</text>')130 131 y_offset += fontsize + 3132 svg_parts.extend(("</svg>", "</div>"))133 # Join everything into a single string134 html = "".join(svg_parts)135 136 137 return html138 139 140 141def generate_sachit_explanations(result_dict:dict):142 descriptions = result_dict['descriptions']143 scores = result_dict['scores']144 MAX_LENGTH = 50145 exp_length = 400146 fontsize = 15147 148 descriptions = zip(scores, descriptions)149 descriptions = sorted(descriptions, key=lambda x: x[0], reverse=True)150 151 # Start the SVG inside a div152 svg_parts = [f'<div style="width: {exp_length}px; height: 450px; background-color: white;">',153 "<svg width=\"100%\" height=\"100%\">"]154 155 # Add a row for each visible bird part156 y_offset = 0157 for score, desc in descriptions:158 159 # Calculate the length of the bar (scaled to fit within the SVG)160 part_score = max(score, 0)161 bar_length = part_score * exp_length162 163 # Split the description into two lines if it's too long164 num_lines = len(desc) // MAX_LENGTH + 1165 for line in range(num_lines):166 desc_line = desc[line*MAX_LENGTH:(line+1)*MAX_LENGTH]167 y_offset += fontsize168 svg_parts.append(f"""169 <text x="0" y="{y_offset}" font-size="{fontsize}" fill="black">170 {desc_line}171 </text>172 """)173 174 # Add the bar175 svg_parts.append(f"""176 <rect x="0" y="{y_offset+3}" width="{bar_length}" height="{fontsize*0.7}" fill="{SACHIT_COLOR}">177 </rect>178 """)179 180 # Add the score181 svg_parts.append(f'<text x="{exp_length - 50}" y="{y_offset+fontsize+3}" font-size="fontsize" fill="{SACHIT_COLOR}">{part_score:.2f}</text>') # Added fill color182 183 y_offset += fontsize + 3184 185 186 svg_parts.extend(("</svg>", "</div>"))187 # Join everything into a single string188 html = "".join(svg_parts)189 190 191 return html192 193# --- Constants created by the functions above ---194BLANK_OVERLAY = img_to_base64(create_blank_image())195PART_COLORS = {part: rgb_to_hex(COLORS[i]) for i, part in enumerate(ORDERED_PARTS)}196blank_image = np.array(Image.open('data/images/final.png').convert('RGB'))197PART_IMAGES_DICT = {file_name: load_part_images(file_name) for file_name in IMAGE_FILE_LIST}198 199# --- Gradio Functions ---200def update_selected_image(event: gr.SelectData):201 image_height = 400202 index = event.index203 204 image_name = IMAGE_FILE_LIST[index]205 current_image.state = image_name206 org_image = Image.open(os.path.join(IMAGES_FOLDER, 'org', image_name)).convert('RGB')207 img_base64 = f"""208 <div style="position: relative; height: {image_height}px; display: inline-block;">209 <img id="birdImage" src="data:image/jpeg;base64,{img_to_base64(org_image)}" style="height: {image_height}px; width: auto;">210 <img id="overlayImage" src="data:image/jpeg;base64,{BLANK_OVERLAY}" style="position:absolute; top:0; left:0; width:auto; height: {image_height}px; opacity: 0;">211 </div>212 """213 gt_label = IMAGE2GT[image_name]214 gt_class.state = gt_label215 216 # --- for initial value only ---217 out_dict = xclip_pred(new_desc=None, 218 new_part_mask=None, 219 new_class=None, 220 org_desc=XCLIP_DESC_PATH, 221 image=Image.open(os.path.join(IMAGES_FOLDER, 'org', current_image.state)).convert('RGB'), 222 model=XCLIP, 223 owlvit_processor=OWLVIT_PRECESSOR, 224 device=DEVICE, 225 image_name=current_image.state,226 cub_embeds=CUB_DESC_EMBEDS,227 cub_idx2name=CUB_IDX2NAME,228 descriptors=XCLIP_DESC)229 xclip_label = out_dict['pred_class']230 clip_pred_scores = out_dict['pred_score']231 xclip_part_scores = out_dict['pred_desc_scores']232 result_dict = {'descriptions': dict(zip(ORG_PART_ORDER, out_dict["descriptions"])), 'pred_scores': xclip_part_scores, 'file_name': current_image.state}233 xclip_exp = generate_xclip_explanations(result_dict, VISIBILITY_DICT[current_image.state], part_mask=dict(zip(ORDERED_PARTS, [1]*12)))234 # --- end of intial value ---235 236 xclip_color = "green" if xclip_label.strip() == gt_label.strip() else "red"237 xclip_pred_markdown = f"""238 ### <span style='color:{xclip_color}'>XCLIP: {xclip_label} {clip_pred_scores:.4f}</span>239 """240 241 gt_label = f"""242 ## {gt_label}243 """244 current_predicted_class.state = xclip_label245 246 # Populate the textbox with current descriptions247 custom_class_name = "class name: custom"248 descs = XCLIP_DESC[xclip_label]249 descs = {k: descs[i] for i, k in enumerate(ORG_PART_ORDER)}250 descs = {k: descs[k] for k in ORDERED_PARTS}251 custom_text = [custom_class_name] + list(descs.values())252 descriptions = ";\n".join(custom_text)253 # textbox = gr.Textbox.update(value=descriptions, lines=12, visible=True, label="XCLIP descriptions", interactive=True, info='Please use ";" to separate the descriptions for each part, and keep the format of {part name}: {descriptions}', show_label=False)254 textbox = gr.Textbox(value=descriptions, 255 lines=12, 256 visible=True, 257 label="XCLIP descriptions", 258 interactive=True, 259 info='Please use ";" to separate the descriptions for each part, and keep the format of {part name}: {descriptions}', 260 show_label=False)261 # modified_exp = gr.HTML().update(value="", visible=True)262 return gt_label, img_base64, xclip_pred_markdown, xclip_exp, current_image, textbox263 264def on_edit_button_click_xclip():265 # empty_exp = gr.HTML.update(visible=False)266 empty_exp = gr.HTML(visible=False)267 268 # Populate the textbox with current descriptions269 descs = XCLIP_DESC[current_predicted_class.state]270 descs = {k: descs[i] for i, k in enumerate(ORG_PART_ORDER)}271 descs = {k: descs[k] for k in ORDERED_PARTS}272 custom_text = ["class name: custom"] + list(descs.values())273 descriptions = ";\n".join(custom_text)274 # textbox = gr.Textbox.update(value=descriptions, lines=12, visible=True, label="XCLIP descriptions", interactive=True, info='Please use ";" to separate the descriptions for each part, and keep the format of {part name}: {descriptions}', show_label=False)275 textbox = gr.Textbox(value=descriptions,276 lines=12,277 visible=True,278 label="XCLIP descriptions",279 interactive=True,280 info='Please use ";" to separate the descriptions for each part, and keep the format of {part name}: {descriptions}',281 show_label=False)282 283 return textbox, empty_exp284 285def convert_input_text_to_xclip_format(textbox_input: str):286 287 # Split the descriptions by newline to get individual descriptions for each part288 descriptions_list = textbox_input.split(";\n")289 # the first line should be "class name: xxx"290 class_name_line = descriptions_list[0]291 new_class_name = class_name_line.split(":")[1].strip()292 293 descriptions_list = descriptions_list[1:]294 295 # construct descripion dict with part name as key296 descriptions_dict = {}297 for desc in descriptions_list:298 if desc.strip() == "":299 continue300 part_name, _ = desc.split(":")301 descriptions_dict[part_name.strip()] = desc302 # fill with empty string if the part is not in the descriptions303 part_mask = {}304 for part in ORDERED_PARTS:305 if part not in descriptions_dict:306 descriptions_dict[part] = ""307 part_mask[part] = 0308 else:309 part_mask[part] = 1310 return descriptions_dict, part_mask, new_class_name311 312def on_predict_button_click_xclip(textbox_input: str):313 descriptions_dict, part_mask, new_class_name = convert_input_text_to_xclip_format(textbox_input)314 315 # Get the new predictions and explanations316 out_dict = xclip_pred(new_desc=descriptions_dict, 317 new_part_mask=part_mask, 318 new_class=new_class_name, 319 org_desc=XCLIP_DESC_PATH, 320 image=Image.open(os.path.join(IMAGES_FOLDER, 'org', current_image.state)).convert('RGB'), 321 model=XCLIP, 322 owlvit_processor=OWLVIT_PRECESSOR, 323 device=DEVICE, 324 image_name=current_image.state,325 cub_embeds=CUB_DESC_EMBEDS,326 cub_idx2name=CUB_IDX2NAME,327 descriptors=XCLIP_DESC)328 xclip_label = out_dict['pred_class']329 xclip_pred_score = out_dict['pred_score']330 xclip_part_scores = out_dict['pred_desc_scores']331 custom_label = out_dict['modified_class']332 custom_pred_score = out_dict['modified_score']333 custom_part_scores = out_dict['modified_desc_scores']334 335 # construct a result dict to generate xclip explanations336 result_dict = {'descriptions': dict(zip(ORG_PART_ORDER, out_dict["descriptions"])), 'pred_scores': xclip_part_scores, 'file_name': current_image.state}337 xclip_explanation = generate_xclip_explanations(result_dict, VISIBILITY_DICT[current_image.state], part_mask)338 modified_result_dict = {'descriptions': dict(zip(ORG_PART_ORDER, out_dict["modified_descriptions"])), 'pred_scores': custom_part_scores, 'file_name': current_image.state}339 modified_explanation = generate_xclip_explanations(modified_result_dict, VISIBILITY_DICT[current_image.state], part_mask)340 341 xclip_color = "green" if xclip_label.strip() == gt_class.state.strip() else "red"342 xclip_pred_markdown = f"""343 ### <span style='color:{xclip_color}'> {xclip_label} {xclip_pred_score:.4f}</span>344 """345 custom_color = "green" if custom_label.strip() == gt_class.state.strip() else "red"346 custom_pred_markdown = f"""347 ### <span style='color:{custom_color}'> {custom_label} {custom_pred_score:.4f}</span>348 """349 # textbox = gr.Textbox.update(visible=False)350 textbox = gr.Textbox(visible=False)351 # return textbox, xclip_pred_markdown, xclip_explanation, custom_pred_markdown, modified_explanation352 353 # modified_exp = gr.HTML().update(value=modified_explanation, visible=True)354 modified_exp = gr.HTML(value=modified_explanation, visible=True)355 return textbox, xclip_pred_markdown, xclip_explanation, custom_pred_markdown, modified_exp356 357 358custom_css = """359 html, body {360 margin: 0;361 padding: 0;362 }363 364 #container {365 position: relative;366 width: 400px;367 height: 400px;368 border: 1px solid #000;369 margin: 0 auto; /* This will center the container horizontally */370 }371 372 #canvas {373 position: absolute;374 top: 0;375 left: 0;376 width: 100%;377 height: 100%;378 object-fit: cover;379 }380 381"""382 383# Define the Gradio interface384with gr.Blocks(theme=gr.themes.Soft(), css=custom_css, title="PEEB") as demo:385 current_image = gr.State("")386 current_predicted_class = gr.State("")387 gt_class = gr.State("")388 389 with gr.Column():390 title_text = gr.Markdown("# PEEB - demo")391 gr.Markdown(392 """393 - In this demo a demo for PEEB paper (NAACL finding 2024). 394 - paper: https://arxiv.org/abs/2403.05297395 - code: https://github.com/anguyen8/peeb/tree/inspect_ddp396 """397 )398 399 # display the gallery of images400 with gr.Column():401 402 gr.Markdown("## Select an image to start!")403 image_gallery = gr.Gallery(value=IMAGE_GALLERY, label=None, preview=False, allow_preview=False, columns=10, height=250)404 gr.Markdown("### Custom descritions: \n The first row should be **class name: {some name};**, where you can name your descriptions. \n For the remianing descriptions, please use **;** to separate the descriptions for each part, and use the format **{part name}: {descriptions}**. \n Note that you can delete a part completely, in such cases, all descriptions will remove the corresponding part.")405 406 with gr.Row():407 with gr.Column():408 image_label = gr.Markdown("### Class Name")409 org_image = gr.HTML()410 411 with gr.Column():412 with gr.Row():413 # xclip_predict_button = gr.Button(label="Predict", value="Predict")414 xclip_predict_button = gr.Button(value="Predict")415 xclip_pred_label = gr.Markdown("### PEEB:")416 xclip_explanation = gr.HTML()417 418 with gr.Column():419 # xclip_edit_button = gr.Button(label="Edit", value="Reset Descriptions")420 xclip_edit_button = gr.Button(value="Reset Descriptions")421 custom_pred_label = gr.Markdown(422 "### Custom Descritpions:"423 )424 xclip_textbox = gr.Textbox(lines=12, placeholder="Edit the descriptions here", visible=False)425 # ai_explanation = gr.Image(type="numpy", visible=True, show_label=False, height=500)426 custom_explanation = gr.HTML()427 428 gr.HTML("<br>")429 430 image_gallery.select(update_selected_image, inputs=None, outputs=[image_label, org_image, xclip_pred_label, xclip_explanation, current_image, xclip_textbox])431 xclip_edit_button.click(on_edit_button_click_xclip, inputs=[], outputs=[xclip_textbox, custom_explanation])432 xclip_predict_button.click(on_predict_button_click_xclip, inputs=[xclip_textbox], outputs=[xclip_textbox, xclip_pred_label, xclip_explanation, custom_pred_label, custom_explanation])433 434demo.launch()