pages/annotation.py

import os
import io
import random
import pandas as pd
import streamlit as st
from tempfile import NamedTemporaryFile
from streamlit_tags import st_tags
from datasets import load_dataset
from huggingface_hub import HfApi, Repository

HF_TOKEN = os.environ.get("HF_TOKEN")
API = HfApi()


def save_data():
    st.session_state["user"]["part_two"] = {
        "aesthetic": [aesthetic],
        "real_buildings": [real_buildings],
        "confidence": [confidence],
        "expectation": [expectation],
        "coherent": [coherent],
        "coherent_others": [coherent_others],
        "not_coherent": [not_coherent],
        "not_coherent_others": [not_coherent_others],
        "feedback": [feedback],
        "storicity": [storicity],
        "perspective": [perspective],
        "image_idx": [image_idx]
    }
    st.session_state["n_annotated_images"] += 1

    results = {
        "user": [st.session_state["user"]["nickname"]],
        **st.session_state["user"]["part_one"],
        **st.session_state["user"]["part_two"]
    }
    user_results = pd.DataFrame.from_dict(results, orient='index').T
    print(user_results)
    updated_annotation = pd.concat([annotations, user_results])
    print(updated_annotation)
    # Create a NamedTemporaryFile
    with NamedTemporaryFile(mode='w+', suffix='.csv') as tmp:
        updated_annotation.to_csv(tmp.name, index=False)
        tmp.seek(0)
        API.upload_file(
            path_or_fileobj=tmp.name,
            path_in_repo="results.csv",
            repo_id="Nicorb/annotation_results",
            repo_type="dataset",
        )


@st.cache_resource
def get_img_dataset(path):
    return load_dataset(
        path=path,
        split="train",
        token=HF_TOKEN
    )


def get_dataset(path):
    return load_dataset(
        path=path,
        split="train",
        token=HF_TOKEN
    )


images = get_img_dataset("Nicorb/test_img")
annotations = get_dataset("Nicorb/annotation_results").to_pandas()

if "n_annotated_images" not in st.session_state:
    st.session_state["n_annotated_images"] = 0
if "selectbox_disable" not in st.session_state:
    st.session_state["selectbox_disable"] = False

st.title("Studio sulla City Identity")

storicity = st.selectbox(
    "Scegli un periodo storico:",
    [
        "Medioevo", "Ottocento", "Novecento"
    ],
    index=None,
    key="storicity",
    disabled=st.session_state["selectbox_disable"]
)
perspective = st.selectbox(
    "Scegli il punto di vista:",
    [
        "Frontale", "Prospettiva centrale", "Panoramico"
    ],
    index=None,
    key="perspective",
    disabled=st.session_state["selectbox_disable"]
)

# GENERARE TRE IMMAGINI PER OGNI SCELTA
random_idxs = random.choices(list(range(len(images))), k=3)

with st.form("valutazione_immagine"):
    if st.session_state["storicity"] and st.session_state["perspective"] and st.session_state["n_annotated_images"] < 3:
        st.session_state["selectbox_disable"] = True
        image_idx = st.session_state["n_annotated_images"]
        st.image(
            images[random_idxs[image_idx]]["image"]
        )

        aesthetic = st.radio(
            "L'immagine mi piace",
            [
                "1 - Per nulla",
                "2",
                "3",
                "4",
                "5 - Totalmente"
            ],
            horizontal=True
        )

        real_buildings = st.radio(
            "Penso che gli edifici nell'immagine esistano realmente:",
            [
                "Si", "No"
            ]
        )

        confidence = st.radio(
            "Sono convinto/a della mia risposta",
            [
                "1 - Per nulla",
                "2",
                "3",
                "4",
                "5 - Totalmente"
            ],
            horizontal=True
        )

        expectation = st.radio(
            "Gli edifici nell'immagine rispettano la mia aspettativa sulla città di Genova",
            [
                "1 - Per nulla d'accordo",
                "2",
                "3",
                "4",
                "5 - Pienamente d'accordo"
            ],
            horizontal=True
        )

        coherent = st.multiselect(
            "Penso che le seguenti caratteristiche dell'immagine siano coerenti con le mie aspettative:",
            [
                "colore", "nitidezza", "aspetto realistico", "aspetto surreale", "stile architettonico",
                "difetti dell'immagine", "presenza del contesto", "correttezza del contesto (se presente)",
                "coerenza con la città", "tratti caratteristici della città"  # aggiungere ripetizione? troppi elementi?
            ],
            placeholder="Seleziona gli elementi"
        )
        coherent_others = st_tags(
            label="Altre caratteristiche:",
            text="Premere invio per aggiungere un nuovo elemento",
            value=[],
            maxtags=-1,
            key="coherent"
        )

        not_coherent = st.multiselect(
            "Penso che le seguenti caratteristiche dell'immagine *NON* siano coerenti con le mie aspettative:",
            [
                "colore", "nitidezza", "aspetto realistico", "aspetto surreale", "stile architettonico",
                "difetti dell'immagine", "presenza del contesto", "correttezza del contesto (se presente)",
                "coerenza con la città", "tratti caratteristici della città"  # aggiungere ripetizione? troppi elementi?
            ],
            placeholder="Seleziona gli elementi"
        )
        not_coherent_others = st_tags(
            label="Altre caratteristiche:",
            text="Premere invio per aggiungere un nuovo elemento",
            value=[],
            maxtags=-1,
            key="not_coherent"
        )

        feedback = st.text_area("Note, idee, suggerimenti e osservazioni")

        submit = st.form_submit_button("Salva", on_click=save_data)

        if submit:
            st.success("Salvataggio completato")
            st.page_link("pages/annotation.py", label="Continua")
    else:
        st.session_state["selectbox_disable"] = False