app.py

import streamlit as st
import pandas as pd
import time
import os
import plotly.graph_objects as go


st.set_page_config(layout="wide")

data_source = st.sidebar.radio("Source", ["Danbooru", "Gelbooru", "Rule 34"], index=0)
if data_source == "Danbooru":
    parquet_file = os.getenv('PARQUET_FILE1')
elif data_source == "Gelbooru":
    parquet_file = os.getenv('PARQUET_FILE2')
elif data_source == "Rule 34":
    parquet_file = os.getenv('PARQUET_FILE3')

@st.cache_resource
def load_and_preprocess_data(parquet_file):
    start_time = time.time()
    df = pd.read_parquet(parquet_file)
    df = df.sort_values(by='post_id', ascending=False)
    df["tags"] = df["tags"].apply(lambda x: set(x))
    df.set_index('post_id', inplace=True)

    sorted_indices = {
        'Post ID (Descending)': df.index,
        'Post ID (Ascending)': df.index[::-1],
        'AVA Score': df['ava_score'].sort_values(ascending=False).index,
        'Aesthetic Score': df['aesthetic_score'].sort_values(ascending=False).index,
    }
    print(f"Data loaded and preprocessed: {time.time() - start_time:.2f} seconds")
    return df, sorted_indices

st.title(f'{data_source} Images')
data, sorted_indices = load_and_preprocess_data(parquet_file)

score_range = st.sidebar.slider('Select AVA Score range', min_value=0.0, max_value=10.0, value=(5.5, 10.0), step=0.1, help='Filter images based on their AVA Score range.')
score_range_v2 = st.sidebar.slider('Select Aesthetic Score range', min_value=0.0, max_value=10.0, value=(9.0, 10.0), step=0.1, help='Filter images based on their Aesthetic Score range.')

min_post_id = int(data.index.min()) if not data.empty else 0
max_post_id = int(data.index.max()) if not data.empty else 100000
post_id_range = st.sidebar.slider('Select Post ID range', 
                                  min_value=min_post_id, 
                                  max_value=max_post_id, 
                                  value=(min_post_id, max_post_id), 
                                  step=1000, 
                                  help='Filter images based on Post ID range.')

available_ratings = sorted(data['rating'].unique().tolist()) if 'rating' in data.columns else ['general']

selected_ratings = st.sidebar.multiselect(
    'Select ratings to include',
    options=available_ratings,
    default=[],
    help='Filter images by their rating category'
)

page_number = st.sidebar.number_input('Page', min_value=1, value=1, step=1, help='Navigate through the pages of filtered results.')
sort_option = st.sidebar.selectbox('Sort by (slow)', options=['Post ID (Descending)', 'Post ID (Ascending)', 'AVA Score', 'Aesthetic Score'], index=0, help='Select sorting option for the results.')

# user input
user_input_tags = st.text_input('Enter tags (space-separated)', value='', help='Filter images based on tags. Use "-" to exclude tags.')   
selected_tags = set([tag.strip() for tag in user_input_tags.split() if tag.strip() and not tag.strip().startswith('-')])
undesired_tags = set([tag[1:] for tag in  user_input_tags.split() if tag.startswith('-')])
print(f"Selected tags: {selected_tags}, Undesired tags: {undesired_tags}")

# Function to filter data based on user input
def filter_data(df, score_range, score_range_v2, post_id_range, selected_tags, sort_option, selected_ratings):
    start_time = time.time()
    
    filtered_data = df[
        (df['ava_score'] >= score_range[0]) & 
        (df['ava_score'] <= score_range[1]) &
        (df['aesthetic_score'] >= score_range_v2[0]) & 
        (df['aesthetic_score'] <= score_range_v2[1]) &
        (df.index >= post_id_range[0]) &
        (df.index <= post_id_range[1])
    ]
    
    if selected_ratings and 'rating' in df.columns:
        filtered_data = filtered_data[filtered_data['rating'].isin(selected_ratings)]
    
    print(f"Data filtered based on scores, post ID and ratings: {time.time() - start_time:.2f} seconds")
    
    if sort_option != "Post ID (Descending)":
        sorted_index = sorted_indices[sort_option]
        sorted_index = sorted_index[sorted_index.isin(filtered_data.index)]
        filtered_data = filtered_data.loc[sorted_index]
        print(f"Applying indcies: {time.time() - start_time:.2f} seconds")
    
    if selected_tags or undesired_tags:
        filtered_data = filtered_data[filtered_data['tags'].apply(lambda x: selected_tags.issubset(x) and not undesired_tags.intersection(x))]
    
    print(f"Data filtered: {time.time() - start_time:.2f} seconds")
    return filtered_data

# Filter data
filtered_data = filter_data(data, score_range, score_range_v2, post_id_range, selected_tags, sort_option, selected_ratings)
st.sidebar.write(f"Total filtered images: {len(filtered_data)}")

# Pagination 
items_per_page = 50
start_idx = (page_number - 1) * items_per_page
end_idx = start_idx + items_per_page
current_data = filtered_data.iloc[start_idx:end_idx]

# Display the data
columns_per_row = 5
rows = [current_data.iloc[i:i + columns_per_row] for i in range(0, len(current_data), columns_per_row)]
for row in rows:
    cols = st.columns(columns_per_row)
    for col, (_, row_data) in zip(cols, row.iterrows()):
        with col:
            post_id = row_data.name
            if data_source == "Danbooru":
                link = f"https://danbooru.donmai.us/posts/{post_id}"
            elif data_source == "Gelbooru": 
                link = f"https://gelbooru.com/index.php?page=post&s=view&id={post_id}"
            elif data_source == "Rule 34":
                link = f"https://rule34.xxx/index.php?page=post&s=view&id={post_id}"
                
            st.image(row_data['large_file_url'], caption=f"ID: {row_data.name}, AVA Score: {row_data['ava_score']:.2f}, Aesthetic Score: {row_data['aesthetic_score']:.2f}\n{link}", use_container_width=True)
            
def histogram_slider(df, column1, column2):
    sample_data = df.sample(min(10000, len(df)))
    
    fig = go.Figure()
    fig.add_trace(go.Histogram(x=sample_data[column1], nbinsx=50, name=column1, opacity=0.75))
    fig.add_trace(go.Histogram(x=sample_data[column2], nbinsx=50, name=column2, opacity=0.75))
    fig.update_layout(
        barmode='overlay',
        bargap=0.1,
        height=200,
        xaxis=dict(showticklabels=True),
        yaxis=dict(showticklabels=True),
        margin=dict(l=0, r=0, t=0, b=0),
        legend=dict(orientation='h', yanchor='bottom', y=-0.4, xanchor='center', x=0.5),
    )
    st.sidebar.plotly_chart(fig, use_container_width=True, config={'displayModeBar': False})

# histogram
start_time = time.time()
histogram_slider(filtered_data, 'ava_score', 'aesthetic_score')
print(f"Histogram displayed: {time.time() - start_time:.2f} seconds")