""" Sidebar components for the embed_explore application. """ import streamlit as st import os import time import hashlib import numpy as np import pandas as pd from typing import Tuple, Optional from shared.services.embedding_service import EmbeddingService from shared.services.clustering_service import ClusteringService from shared.services.file_service import FileService from shared.lib.progress import StreamlitProgressContext from shared.components.clustering_controls import ( render_projection_controls, render_kmeans_controls, ) from shared.utils.backend import check_cuda_available, resolve_backend, is_oom_error from shared.utils.logging_config import get_logger logger = get_logger(__name__) def render_embedding_section() -> Tuple[bool, Optional[str], Optional[str], int, int]: """ Render the embedding section of the sidebar. Returns: Tuple of (embed_button_clicked, image_dir, model_name, n_workers, batch_size) """ with st.expander("Embed", expanded=True): image_dir = st.text_input("Image folder path") # Get available models dynamically available_models = EmbeddingService.get_model_options() model_name = st.selectbox("Model", available_models) col1, col2 = st.columns(2) with col1: n_workers = st.number_input( "N workers", min_value=1, max_value=64, value=16, step=1 ) with col2: batch_size = st.number_input( "Batch size", min_value=1, max_value=2048, value=32, step=1 ) embed_button = st.button("Run Embedding") # Handle embedding execution if embed_button and image_dir and os.path.isdir(image_dir): with StreamlitProgressContext(st.empty(), "Embedding complete!") as progress: try: embeddings, valid_paths = EmbeddingService.generate_embeddings( image_dir, model_name, batch_size, n_workers, progress_callback=progress ) if embeddings.shape[0] == 0: st.error("No valid image embeddings found.") logger.warning("Embedding generation returned 0 embeddings") st.session_state.embeddings = None st.session_state.valid_paths = None st.session_state.labels = None st.session_state.data = None st.session_state.selected_image_idx = None else: logger.info(f"Embeddings stored: shape={embeddings.shape}, dtype={embeddings.dtype}") st.success(f"Generated {embeddings.shape[0]} image embeddings.") st.session_state.embeddings = embeddings st.session_state.valid_paths = valid_paths st.session_state.last_image_dir = image_dir st.session_state.embedding_complete = True # Reset projection/clustering/selection state for the new embeddings st.session_state.labels = None st.session_state.kmeans_column = None st.session_state.data = None st.session_state.selected_image_idx = None except Exception as e: st.error(f"Error during embedding: {e}") logger.exception("Embedding generation failed") elif embed_button: st.error("Please provide a valid image directory path.") return embed_button, image_dir, model_name, n_workers, batch_size def render_projection_section(): """Render the 2D projection section.""" with st.expander("Project to 2D", expanded=False): embeddings = st.session_state.get("embeddings", None) valid_paths = st.session_state.get("valid_paths", None) if embeddings is None or valid_paths is None or len(valid_paths) < 2: st.info("Run embedding first to enable projection.") return n_samples, emb_dim = embeddings.shape st.markdown(f"**Ready to project:** {n_samples:,} images ({emb_dim}-dim embeddings)") reduction_method = st.selectbox( "Dimensionality Reduction", ["TSNE", "PCA", "UMAP"], help="Method to project high-dimensional embeddings to 2D for visualization.", ) dim_reduction_backend, seed = render_projection_controls() if st.button("Project to 2D", type="primary"): _run_projection(embeddings, valid_paths, reduction_method, dim_reduction_backend, seed) def render_kmeans_section(): """Render the optional KMeans clustering section.""" with st.expander("KMeans Clustering", expanded=False): df_plot = st.session_state.get("data", None) embeddings = st.session_state.get("embeddings", None) if df_plot is None or embeddings is None: st.info("Run projection first to enable KMeans.") return emb_dim = embeddings.shape[1] st.markdown(f"**{len(df_plot):,} points** ({emb_dim}-dim embeddings)") n_clusters = st.slider("Number of clusters", 2, min(100, max(2, len(df_plot) // 2)), 5) clustering_backend, n_workers, seed = render_kmeans_controls() if st.button("Run KMeans", type="primary"): _run_kmeans(embeddings, n_clusters, clustering_backend, n_workers, seed) def _run_projection(embeddings, valid_paths, reduction_method, dim_reduction_backend, seed): """Run dim reduction and create the 2D scatter plot dataframe.""" try: cuda_available, device_info = check_cuda_available() actual_backend = resolve_backend(dim_reduction_backend, "reduction") logger.info("=" * 60) logger.info("PROJECTION START") logger.info(f"Device: {device_info} (CUDA: {'Yes' if cuda_available else 'No'})") logger.info(f"Backend: {actual_backend} (requested: {dim_reduction_backend})") t_start = time.time() n_samples, emb_dim = embeddings.shape logger.info(f"Records: {n_samples:,} | Dim: {emb_dim}") with st.spinner(f"Running {reduction_method}..."): reduced = ClusteringService.run_dim_reduction_safe( embeddings, reduction_method, n_workers=8, dim_reduction_backend=actual_backend, seed=seed ) t_total = time.time() - t_start logger.info(f"Projection complete in {t_total:.2f}s") # Build plot dataframe (no cluster column) df_plot = pd.DataFrame({ "x": reduced[:, 0], "y": reduced[:, 1], "image_path": valid_paths, "file_name": [os.path.basename(p) for p in valid_paths], "idx": range(len(valid_paths)), }) # Carry over any prior KMeans columns from the previous df_plot (if length matches) prev_df = st.session_state.get("data") if prev_df is not None and len(prev_df) == len(df_plot): for col in prev_df.columns: if col.startswith("KMeans (k="): df_plot[col] = prev_df[col].values data_hash = hashlib.md5(f"{len(df_plot)}_{reduction_method}_{t_total}".encode()).hexdigest()[:8] st.session_state.data = df_plot st.session_state.data_version = data_hash st.session_state.selected_image_idx = None logger.info("=" * 60) st.success(f"Projected {n_samples:,} points to 2D using {reduction_method}.") except (RuntimeError, OSError) as e: if is_oom_error(e): st.error("**GPU Out of Memory**") st.info("Try: Reduce dataset size, use 'sklearn' backend, or try PCA.") logger.exception("GPU OOM during projection") else: st.error(f"Error during projection: {e}") logger.exception("Projection error") except MemoryError: st.error("**System Out of Memory** - Reduce dataset size") logger.exception("System memory exhausted during projection") except Exception as e: st.error(f"Error: {e}") logger.exception("Unexpected projection error") def _run_kmeans(embeddings, n_clusters, clustering_backend, n_workers, seed): """Run KMeans on already-extracted embeddings and add labels to df_plot.""" try: actual_backend = resolve_backend(clustering_backend, "clustering") logger.info(f"KMeans: k={n_clusters}, backend={actual_backend}") with st.spinner(f"Running KMeans (k={n_clusters})..."): labels = ClusteringService.run_kmeans_only_safe( embeddings, n_clusters, n_workers=n_workers, clustering_backend=actual_backend, seed=seed ) df_plot = st.session_state.data kmeans_col = f"KMeans (k={n_clusters})" df_plot[kmeans_col] = labels.astype(str) st.session_state.data = df_plot st.session_state.labels = labels st.session_state.kmeans_column = kmeans_col # Compute clustering summary on the full embedding space. # Cache by kmeans_col so multiple KMeans runs can each have their own # summary + representatives that the user can switch between. logger.info("Computing clustering summary statistics...") summary_df, representatives = ClusteringService.generate_clustering_summary( embeddings, labels, df_plot ) summaries = st.session_state.get("clustering_summaries", {}) reps_by_col = st.session_state.get("clustering_representatives_by_col", {}) summaries[kmeans_col] = summary_df reps_by_col[kmeans_col] = representatives st.session_state.clustering_summaries = summaries st.session_state.clustering_representatives_by_col = reps_by_col logger.info(f"Clustering summary computed for {kmeans_col}: {len(summary_df)} clusters") logger.info(f"KMeans complete: {len(np.unique(labels))} clusters") st.success(f"KMeans complete! {len(np.unique(labels))} clusters assigned.") except (RuntimeError, OSError) as e: if is_oom_error(e): st.error("**GPU Out of Memory**") logger.exception("GPU OOM during KMeans") else: st.error(f"Error during KMeans: {e}") logger.exception("KMeans error") except MemoryError: st.error("**System Out of Memory** - Reduce dataset size") logger.exception("System memory exhausted during KMeans") except Exception as e: st.error(f"Error: {e}") logger.exception("Unexpected KMeans error") def _get_available_kmeans_cols(df_plot) -> list: """Return KMeans columns in df_plot sorted by k value.""" if df_plot is None: return [] return sorted( [c for c in df_plot.columns if c.startswith("KMeans (k=")], key=lambda c: int(c.split("=")[1].rstrip(")")), ) def render_save_section(): """Render the save operations section of the sidebar. Both 'Save Images from Specific Cluster' and 'Repartition Images by Cluster' require at least one KMeans run. When multiple KMeans runs exist, the user picks which one to operate on via a shared selector at the top. """ df_plot = st.session_state.get("data", None) kmeans_cols = _get_available_kmeans_cols(df_plot) if not kmeans_cols: st.info("Run KMeans first to enable saving by cluster.") return # Shared selector: which KMeans run drives both save operations default_idx = len(kmeans_cols) - 1 # most recent run selected_kmeans_col = st.selectbox( "KMeans result", options=kmeans_cols, index=default_idx, key="save_kmeans_selector", help="Pick which KMeans run to use for save / repartition.", ) # --- Save images from a specific cluster utility --- save_status_placeholder = st.empty() with st.expander("Save Images from Specific Cluster", expanded=True): available_clusters = sorted(df_plot[selected_kmeans_col].unique(), key=lambda x: int(x)) selected_clusters = st.multiselect( "Select cluster(s) to save", available_clusters, default=available_clusters[:1] if available_clusters else [], key="save_cluster_select", ) save_dir = st.text_input( "Directory to save selected cluster images", value="cluster_selected_output", key="save_cluster_dir", ) save_cluster_button = st.button("Save images", key="save_cluster_btn") if save_cluster_button and selected_clusters: cluster_rows = df_plot[df_plot[selected_kmeans_col].isin(selected_clusters)].copy() # FileService expects a 'cluster' column cluster_rows["cluster"] = cluster_rows[selected_kmeans_col] max_workers = st.session_state.get("num_threads", 8) with StreamlitProgressContext( save_status_placeholder, f"Images from cluster(s) {', '.join(map(str, selected_clusters))} saved successfully!" ) as progress: try: save_summary_df, csv_path = FileService.save_cluster_images( cluster_rows, save_dir, max_workers, progress_callback=progress ) st.info(f"Summary CSV saved at {csv_path}") except Exception as e: save_status_placeholder.error(f"Error saving images: {e}") elif save_cluster_button: save_status_placeholder.warning("Please select at least one cluster.") # --- Repartition expander and status --- repartition_status_placeholder = st.empty() with st.expander("Repartition Images by Cluster", expanded=False): st.markdown("**Target directory for repartitioned images (will be created):**") repartition_dir = st.text_input( "Directory", value="repartitioned_output", key="repartition_dir", ) max_workers = st.number_input( "Number of threads (higher = faster, try 8-32)", min_value=1, max_value=64, value=8, step=1, key="num_threads", ) repartition_button = st.button("Repartition images by cluster", key="repartition_btn") if repartition_button: df_for_repartition = df_plot.copy() df_for_repartition["cluster"] = df_for_repartition[selected_kmeans_col] with StreamlitProgressContext( repartition_status_placeholder, f"Repartition complete! Images organized in {repartition_dir}", ) as progress: try: repartition_summary_df, csv_path = FileService.repartition_images_by_cluster( df_for_repartition, repartition_dir, max_workers, progress_callback=progress ) st.info(f"Summary CSV saved at {csv_path}") except Exception as e: repartition_status_placeholder.error(f"Error repartitioning images: {e}") def render_clustering_sidebar(): """Render the complete sidebar with embed / project / KMeans / save sections.""" tab_compute, tab_save = st.tabs(["Compute", "Save"]) with tab_compute: render_embedding_section() render_projection_section() render_kmeans_section() with tab_save: render_save_section()