# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: sg
#     language: python
#     name: sg
# ---

# +
import pandas as pd
import os.path as osp

from shapeglot.vis_utils import stack_images_horizontally


# -

def random_zsl_example(row, top_image_dir):
    zsl_class = row.stimulus_class
    im_loc_a = osp.join(top_image_dir, zsl_class, row.stimulus_a, 'model.png')
    im_loc_b = osp.join(top_image_dir, zsl_class, row.stimulus_b, 'model.png')
    im_loc_c = osp.join(top_image_dir, zsl_class, row.stimulus_c, 'model.png')
    images = [im_loc_a, im_loc_b, im_loc_c]    
    print('Utterance:', row.utterance)
    print('Target:', row.target)
    return stack_images_horizontally(images)


zsl_referential_data_file = '../../data/zero_shot_data_non_chairs/zsl_annotated_triplets.csv'
top_image_dir = '../../data/zero_shot_data_non_chairs/images'

df = pd.read_csv(zsl_referential_data_file)
print(df.stimulus_class.unique())

random_zsl_example(df.sample(1).iloc[0], top_image_dir)

StarcoderdataJupyterScriptsDedupFiltered