test

app.py

@@ -1,16 +1,24 @@
 import streamlit as st
 import pandas as pd
 import time
-import json
 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():
+def load_and_preprocess_data(parquet_file):
     start_time = time.time()
-    df = pd.read_parquet(os.getenv('PARQUET_FILE'))
+    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)
@@ -18,40 +26,62 @@ def load_and_preprocess_data():
     sorted_indices = {
         'Post ID (Descending)': df.index,
         'Post ID (Ascending)': df.index[::-1],
-        'Clip Score': df['clip_aesthetic'].sort_values(ascending=False).index,
-        'Siglip Score': df['clip_aesthetic_2_5'].sort_values(ascending=False).index,
+        '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('Danbooru Images')
-data, sorted_indices = load_and_preprocess_data()
+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'
+)
 
-# isdebar
-st.sidebar.header('Filter Options')
-st.sidebar.write('Adjust the filter options to refine the results.')
-score_range = st.sidebar.slider('Select clip score range', min_value=0.0, max_value=10.0, value=(0.0, 10.0), step=0.1, help='Filter images based on their CLIP score range.')
-score_range_v2 = st.sidebar.slider('Select siglip score range', min_value=0.0, max_value=10.0, value=(6.0, 10.0), step=0.1, help='Filter images based on their SigLIP score range.')
 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)', 'Clip Score', 'Siglip Score'], help='Select sorting option for the 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)', help='Filter images based on tags. Use "-" to exclude tags.')   
+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, selected_tags, sort_option):
+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['clip_aesthetic'] >= score_range[0]) & 
-        (df['clip_aesthetic'] <= score_range[1]) &
-        (df['clip_aesthetic_2_5'] >= score_range_v2[0]) & 
-        (df['clip_aesthetic_2_5'] <= score_range_v2[1])
+        (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])
     ]
-    print(f"Data filtered based on scores: {time.time() - start_time:.2f} seconds")
+    
+    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]
@@ -66,11 +96,11 @@ def filter_data(df, score_range, score_range_v2, selected_tags, sort_option):
     return filtered_data
 
 # Filter data
-filtered_data = filter_data(data, score_range, score_range_v2, selected_tags, sort_option)
+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 = 30
+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]
@@ -82,11 +112,18 @@ for row in rows:
     cols = st.columns(columns_per_row)
     for col, (_, row_data) in zip(cols, row.iterrows()):
         with col:
-            st.image(row_data['large_file_url'], caption=f"ID: {row_data.name}, CLIP: {row_data['clip_aesthetic']:.2f}, SigLIP: {row_data['clip_aesthetic_2_5']:.2f}", use_column_width=True)
-            # st.markdown("<div style='margin: 2px;'></div>", unsafe_allow_html=True)
+            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(5000, len(df)))
+    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))
@@ -104,5 +141,5 @@ def histogram_slider(df, column1, column2):
 
 # histogram
 start_time = time.time()
-histogram_slider(filtered_data, 'clip_aesthetic', 'clip_aesthetic_2_5')
+histogram_slider(filtered_data, 'ava_score', 'aesthetic_score')
 print(f"Histogram displayed: {time.time() - start_time:.2f} seconds")