michelecafagna26/High-Level-Dataset-explorer

1import streamlit as st2from datasets import load_dataset3import numpy as np4 5st.set_page_config(page_title="High-Level dataset")6 7FIELDS = ["scene", "action", "rationale", "object"]8QS = {9    "scene": "Where is the picture taken?",10    "action": "What is the subject doing?",11    "rationale": "Why is the subject doing it?"12}13SPLITS = ["test", "train"]14 15AVG_PURITY = 1.1016 17AVG_DIVERSITY = 0.87281918MIN_DIVERSITY = 019MAX_DIVERSITY = 10020 21@st.cache22def load_data(split):23 24    dataset = load_dataset("michelecafagna26/hl")25 26    coco2id = {int(dataset[split][i]['file_name'].replace("COCO_train2014_", "").replace(".jpg", "")): i for i in27               range(len(dataset[split]))}28 29    return dataset, coco2id30 31 32def write_obj(dataset, img_id, options, split, list_type="num", show_questions=False,33              show_conf=False):34 35    st.image(dataset[split][img_id]['image'])36 37    item_purity = np.mean([np.mean(dataset[split][img_id]['purity'][k]) for k in dataset[split][img_id]['purity']])38    item_diversity = np.mean(list(dataset[split][img_id]['diversity'].values()))39 40    # normalize41    item_diversity = 1-(item_diversity-MIN_DIVERSITY)/(MAX_DIVERSITY-MIN_DIVERSITY)42 43    col1, col2 = st.columns(2)44 45    col1.metric(label="Diversity score",46                value=round(item_diversity, 2),47                delta=round(item_diversity - AVG_DIVERSITY, 2),48                help="Item's internal lexical diversity.\n Positive delta means higher then the average")49 50    col2.metric(label="Purity score",51                value=round(item_purity, 2),52                delta=round(item_purity - AVG_PURITY, 2),53                help="Item's internal semantic similarity.\n Positive delta means higher then the average")54 55    for field in options:56 57        st.markdown(f"## {field.capitalize()}")58 59        if show_questions and field != "object":60            st.markdown(f" Question: _{QS[field]}_")61 62        for n, annotation in enumerate(dataset[split][img_id][field]):63 64            col1, col2 = st.columns(2)65 66            if list_type == "num":67                col1.markdown(f"{n + 1}. {annotation}")68            else:69                col1.markdown(f"{list_type} {annotation}")70 71            if show_conf and field != "object":72                col2.metric(label="confidence score",73                            value=dataset[split][img_id]['confidence'][field][n])74 75 76def main():77    st.title('High-Level Dataset')78 79    show_questions = st.sidebar.checkbox('Questions')80    show_conf = st.sidebar.checkbox('Confidence scores')81    options = st.sidebar.multiselect(82        'Choose the annotations',83        FIELDS,84        default=FIELDS)85 86    split = st.sidebar.selectbox(87        'Split',88        SPLITS)89 90    dataset, coco2id = load_data(split)91 92    # sidebar93    choosen_image = st.selectbox(94        'Select an image',95        list(coco2id.keys()),96        help="write a key like: 7603"97    )98 99    write_obj(dataset, coco2id[choosen_image], options=options, split=split, list_type="num",100              show_questions=show_questions, show_conf=show_conf)101 102 103if __name__ == "__main__":104    main()105