streamlit_evaluation_app.py

#!/usr/bin/env python3
"""
Streamlit Data Viewer and Model Evaluation System
Interactive dashboard for exploring validation results with threshold filtering
"""

import streamlit as st
import json
import pandas as pd
import numpy as np
from PIL import Image
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import os
from pathlib import Path
import subprocess
import sys
from rapidocr import RapidOCR
from matplotlib import pyplot as plt

# Page config
st.set_page_config(
    page_title="Pseudoable Classifier Evaluation Dashboard",
    page_icon="🔍",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for better styling
st.markdown("""
<style>
    .main-header {
        font-size: 3rem;
        font-weight: bold;
        color: #1f77b4;
        text-align: center;
        margin-bottom: 2rem;
        padding: 1rem;
        background: linear-gradient(90deg, #f0f8ff, #e6f3ff);
        border-radius: 10px;
    }
    .metric-card {
        background-color: #f8f9fa;
        padding: 1rem;
        border-radius: 8px;
        border-left: 4px solid #1f77b4;
        margin: 0.5rem 0;
    }
    .filter-section {
        background-color: #ffffff;
        padding: 1.5rem;
        border-radius: 10px;
        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
        margin-bottom: 2rem;
    }
    .image-container {
        border: 2px solid #e6e6e6;
        border-radius: 8px;
        padding: 10px;
        margin: 10px 0;
        background-color: #fafafa;
    }
    .prediction-badge {
        display: inline-block;
        padding: 0.25rem 0.5rem;
        border-radius: 15px;
        font-size: 0.8rem;
        font-weight: bold;
        margin: 0.2rem;
    }
    .correct-prediction {
        background-color: #d4edda;
        color: #155724;
    }
    .incorrect-prediction {
        background-color: #f8d7da;
        color: #721c24;
    }
</style>
""", unsafe_allow_html=True)

@st.cache_data
def load_task_config(config_path: str = './task_config.json'):
    """Load task configuration from JSON file"""
    if not os.path.exists(config_path):
        # Try to load from checkpoints directory
        checkpoint_config = './checkpoints/task_config.json'
        if os.path.exists(checkpoint_config):
            config_path = checkpoint_config
        else:
            return None
    
    with open(config_path, 'r') as f:
        config = json.load(f)
    return config

@st.cache_data
def load_validation_results(file_path: str = './validation_results.json'):
    """Load validation results from JSON file"""
    if not os.path.exists(file_path):
        return None
    
    with open(file_path, 'r') as f:
        data = json.load(f)
    return data

@st.cache_resource
def get_ocr_engine():
    """Initialize and cache OCR engine"""
    return RapidOCR()

@st.cache_data
def extract_text_from_image(image_path: str):
    """Extract text from image using OCR"""
    try:
        engine = get_ocr_engine()
        result = engine(image_path)
        
        # Handle new RapidOCR output format
        if result and hasattr(result, 'txts') and result.txts:
            texts = result.txts
            return {
                'text': ' '.join(texts) if texts else '',
                'num_text_blocks': len(texts),
                'has_text': len(texts) > 0
            }
        elif result and isinstance(result, (list, tuple)) and len(result) > 0:
            # Fallback for older format
            texts = []
            for item in result:
                if len(item) >= 2:
                    texts.append(item[1])
            
            return {
                'text': ' '.join(texts) if texts else '',
                'num_text_blocks': len(texts),
                'has_text': len(texts) > 0
            }
        else:
            return {
                'text': '',
                'num_text_blocks': 0,
                'has_text': False
            }
    except Exception as e:
        return {
            'text': f'OCR Error: {str(e)}',
            'num_text_blocks': 0,
            'has_text': False
        }

@st.cache_data
def process_validation_data(validation_data, task_config):
    """Process validation data into DataFrame for easier filtering"""
    if not validation_data or not task_config:
        return None
    
    rows = []
    tasks = {task['key']: task for task in task_config['tasks']}
    
    for result in validation_data['results']:
        row = {
            'idx': result['idx'],
            'caption': result['caption'],
            'image_path': result['image_path'],
            'url': result['url'],
            'hash': result['hash']
        }
        
        # Ground truth and predictions
        gt = result['ground_truth']
        pred = result['predictions']
        
        # Process each task dynamically
        for task_key, task_info in tasks.items():
            # Ground truth
            row[f'gt_{task_key}'] = gt[task_key]
            
            # Predictions
            row[f'pred_{task_key}'] = pred[f'{task_key}_prediction']
            row[f'conf_{task_key}'] = pred[f'{task_key}_confidence']
            
            # For binary tasks, also include probability for 'yes'
            if task_info['type'] == 'binary':
                row[f'prob_{task_key}_yes'] = pred.get(f'{task_key}_prob_yes', 0.5)
            
            # Correctness
            row[f'correct_{task_key}'] = gt[task_key] == pred[f'{task_key}_prediction']
        
        rows.append(row)
    
    return pd.DataFrame(rows)

def run_validation_if_needed():
    """Run validation if results don't exist"""
    if not os.path.exists('./validation_results.json'):
        st.warning("Validation results not found. Running validation...")
        
        # Check if model exists
        if not os.path.exists('./checkpoints/multi_head_siglip2_classifier.pth'):
            st.error("❌ Trained model not found! Please run the training pipeline first.")
            st.code("python stage_4.py")
            return False
        
        # Run validation
        with st.spinner("Running model on validation set... This may take a few minutes."):
            try:
                result = subprocess.run([sys.executable, 'validation_runner.py'], 
                                      capture_output=True, text=True)
                if result.returncode == 0:
                    st.success("✅ Validation completed successfully!")
                    st.rerun()
                else:
                    st.error(f"❌ Validation failed: {result.stderr}")
                    return False
            except Exception as e:
                st.error(f"❌ Error running validation: {e}")
                return False
    
    return True

def create_overview_metrics(df, validation_data, task_config):
    """Create overview metrics section"""
    st.markdown("## 📊 Overview Metrics")
    
    tasks = [task['key'] for task in task_config['tasks']]
    
    # Basic stats
    col1, col2, col3, col4 = st.columns(4)
    
    with col1:
        st.metric("Total Samples", len(df))
    
    with col2:
        avg_confidence = np.mean([df[f'conf_{task}'].mean() for task in tasks])
        st.metric("Avg Confidence", f"{avg_confidence:.3f}")
    
    with col3:
        overall_accuracy = np.mean([df[f'correct_{task}'].mean() for task in tasks])
        st.metric("Overall Accuracy", f"{overall_accuracy:.3f}")
    
    with col4:
        if validation_data and 'metadata' in validation_data:
            model_accuracies = validation_data['metadata']['validation_accuracies']
            model_avg = np.mean(list(model_accuracies.values()))
            st.metric("Model Accuracy", f"{model_avg:.3f}")
    
    # Detailed accuracies
    st.markdown("### 🎯 Accuracy per Classification Task")
    
    # Create dynamic columns based on number of tasks
    n_tasks = len(tasks)
    n_cols = min(5, n_tasks)  # Max 5 columns
    acc_cols = st.columns(n_cols)
    
    for i, task in enumerate(tasks):
        task_info = next(t for t in task_config['tasks'] if t['key'] == task)
        with acc_cols[i % n_cols]:
            accuracy = df[f'correct_{task}'].mean()
            st.metric(task_info['name'], f"{accuracy:.3f}")

def create_confidence_distribution_plot(df, task_config):
    """Create confidence distribution plots"""
    tasks = [task['key'] for task in task_config['tasks']]
    task_names = [task['name'] for task in task_config['tasks']]
    
    n_tasks = len(tasks)
    n_cols = 3
    n_rows = (n_tasks + n_cols - 1) // n_cols
    
    fig = make_subplots(
        rows=n_rows, cols=n_cols,
        subplot_titles=task_names,
        specs=[[{"secondary_y": False} for _ in range(n_cols)] for _ in range(n_rows)]
    )
    
    colors = plt.cm.Set1(np.linspace(0, 1, n_tasks))
    
    for i, (task_key, color) in enumerate(zip(tasks, colors)):
        row = (i // n_cols) + 1
        col = (i % n_cols) + 1
        
        fig.add_trace(
            go.Histogram(
                x=df[f'conf_{task_key}'],
                nbinsx=20,
                name=f'{task_key}',
                marker_color=f'rgba({color[0]*255:.0f},{color[1]*255:.0f},{color[2]*255:.0f},0.7)',
                opacity=0.7
            ),
            row=row, col=col
        )
    
    fig.update_layout(
        title="Confidence Score Distributions",
        showlegend=False,
        height=200 * n_rows + 100
    )
    
    return fig

def apply_filters(df, task_config):
    """Apply user-defined filters to the dataframe"""
    st.markdown("## 🔍 Filter Data")
    
    tasks = {task['key']: task for task in task_config['tasks']}
    
    # Create filter sidebar
    with st.sidebar:
        st.markdown("### Task Confidence Filters")
        
        # Confidence thresholds for each task
        confidence_filters = {}
        for task_key, task_info in tasks.items():
            if task_info['type'] == 'multi_class':
                # Only show confidence threshold for multi-class tasks
                confidence_filters[task_key] = st.slider(
                    f"{task_info['name']} Confidence", 
                    0.0, 1.0, 0.5, 0.01,
                    key=f"conf_{task_key}"
                )
        
        st.markdown("### Content Filters")
        
        # Category filters (for multi-class tasks)
        category_filters = {}
        for task_key, task_info in tasks.items():
            if task_info['type'] == 'multi_class':
                available_values = df[f'gt_{task_key}'].unique().tolist()
                selected_values = st.multiselect(
                    f"Ground Truth {task_info['name']}", 
                    available_values, 
                    default=available_values,
                    key=f"gt_{task_key}_filter"
                )
                category_filters[task_key] = selected_values
        
        # Binary prediction filters
        st.markdown("**Filter by Predictions:**")
        prediction_filters = {}
        for task_key, task_info in tasks.items():
            if task_info['type'] == 'binary':
                filter_value = st.selectbox(
                    f"{task_info['name']}:",
                    ["All", "Yes only", "No only"],
                    key=f"pred_{task_key}_filter"
                )
                prediction_filters[task_key] = filter_value
        
        # Correctness filter
        st.markdown("**Filter by Correctness:**")
        correctness_filter = st.selectbox(
            "Show only:",
            ["All predictions", "Correct predictions", "Incorrect predictions"]
        )
        
        # OCR filters (if screenshot task exists)
        has_screenshot_task = any(task['key'] == 'is_screenshot_with_text' for task in task_config['tasks'])
        if has_screenshot_task:
            st.markdown("**Filter by Text Content:**")
            ocr_filter = st.selectbox(
                "Text Content:",
                ["All images", "Images with text", "Images without text"],
                key="ocr_filter"
            )
            enable_ocr = st.checkbox("Enable OCR text extraction", value=True)
        else:
            ocr_filter = "All images"
            enable_ocr = False
    
    # Apply filters
    filtered_df = df.copy()
    
    # Confidence filters
    for task_key, threshold in confidence_filters.items():
        filtered_df = filtered_df[filtered_df[f'conf_{task_key}'] >= threshold]
    
    # Category filters
    for task_key, selected_values in category_filters.items():
        filtered_df = filtered_df[filtered_df[f'gt_{task_key}'].isin(selected_values)]
    
    # Binary prediction filters
    for task_key, filter_value in prediction_filters.items():
        if filter_value == "Yes only":
            filtered_df = filtered_df[filtered_df[f'pred_{task_key}'] == 'yes']
        elif filter_value == "No only":
            filtered_df = filtered_df[filtered_df[f'pred_{task_key}'] == 'no']
    
    # Correctness filter
    if correctness_filter == "Correct predictions":
        correct_mask = True
        for task_key in tasks.keys():
            correct_mask = correct_mask & filtered_df[f'correct_{task_key}']
        filtered_df = filtered_df[correct_mask]
    elif correctness_filter == "Incorrect predictions":
        correct_mask = True
        for task_key in tasks.keys():
            correct_mask = correct_mask & filtered_df[f'correct_{task_key}']
        filtered_df = filtered_df[~correct_mask]
    
    # Show filter results
    st.info(f"Filtered to {len(filtered_df)} samples (from {len(df)} total)")
    
    return filtered_df, ocr_filter, enable_ocr

def display_sample_images(df, task_config, ocr_filter="All images", enable_ocr=True):
    """Display sample images with predictions and ground truth"""
    st.markdown("## 🖼️ Sample Images")
    
    if len(df) == 0:
        st.warning("No images match the current filters.")
        return
    
    tasks = {task['key']: task for task in task_config['tasks']}
    
    # Add controls for image display
    col1, col2, col3 = st.columns([2, 1, 1])
    
    with col1:
        max_images = st.slider(
            "Number of images to display", 
            min_value=10, 
            max_value=min(200, len(df)), 
            value=min(50, len(df)), 
            step=10
        )
    
    with col2:
        sort_by = st.selectbox(
            "Sort by:",
            ["Original order", "Confidence (low to high)", "Confidence (high to low)"]
        )
    
    with col3:
        cols_per_row = st.selectbox("Images per row:", [2, 3, 4], index=1)
    
    # Sort dataframe if requested
    task_keys = list(tasks.keys())
    if sort_by == "Confidence (low to high)":
        avg_conf = sum(df[f'conf_{task}'] for task in task_keys) / len(task_keys)
        display_df = df.iloc[avg_conf.argsort()].head(max_images)
    elif sort_by == "Confidence (high to low)":
        avg_conf = sum(df[f'conf_{task}'] for task in task_keys) / len(task_keys)
        display_df = df.iloc[avg_conf.argsort()[::-1]].head(max_images)
    else:
        display_df = df.head(max_images)
    
    # Apply OCR filtering if needed
    if enable_ocr and ocr_filter != "All images":
        st.info("🔍 Applying OCR filtering... This may take a moment for many images.")
        
        ocr_results = []
        progress_bar = st.progress(0)
        
        for idx, (_, row) in enumerate(display_df.iterrows()):
            if os.path.exists(row['image_path']):
                ocr_result = extract_text_from_image(row['image_path'])
                ocr_results.append(ocr_result['has_text'])
            else:
                ocr_results.append(False)
            
            progress_bar.progress((idx + 1) / len(display_df))
        
        progress_bar.empty()
        
        # Filter based on OCR results
        if ocr_filter == "Images with text":
            mask = ocr_results
        else:  # "Images without text"
            mask = [not has_text for has_text in ocr_results]
        
        display_df = display_df[mask].reset_index(drop=True)
        st.success(f"OCR filtering complete. Found {len(display_df)} images matching criteria.")
    
    # Display images
    for i in range(0, len(display_df), cols_per_row):
        cols = st.columns(cols_per_row)
        
        for j in range(cols_per_row):
            if i + j < len(display_df):
                row = display_df.iloc[i + j]
                
                with cols[j]:
                    # Load and display image
                    try:
                        if os.path.exists(row['image_path']):
                            img = Image.open(row['image_path'])
                            st.image(img, caption=f"Sample {row['idx']}", use_column_width=True)
                        else:
                            st.error(f"Image not found: {row['image_path']}")
                            continue
                    except Exception as e:
                        st.error(f"Error loading image: {e}")
                        continue
                    
                    # Caption
                    st.markdown(f"**Caption:** {row['caption'][:100]}...")
                    
                    # OCR Text Extraction
                    if enable_ocr and 'is_screenshot_with_text' in tasks:
                        with st.expander("🔍 Extracted Text (OCR)", expanded=False):
                            ocr_result = extract_text_from_image(row['image_path'])
                            if ocr_result['has_text']:
                                st.markdown(f"**Text Blocks Found:** {ocr_result['num_text_blocks']}")
                                st.text_area(
                                    "Extracted Text:",
                                    value=ocr_result['text'],
                                    height=100,
                                    key=f"ocr_text_{row['idx']}",
                                    help="Text extracted from the image using OCR"
                                )
                                
                                text_length = len(ocr_result['text'])
                                word_count = len(ocr_result['text'].split())
                                st.caption(f"📊 Text stats: {text_length} chars, {word_count} words")
                                
                                if row['pred_is_screenshot_with_text'] == 'yes':
                                    st.success("✅ Screenshot prediction correlates with text presence")
                                elif ocr_result['num_text_blocks'] > 5:
                                    st.warning("⚠️ High text content but not predicted as screenshot")
                            else:
                                st.info("No text detected in this image")
                                if row['pred_is_screenshot_with_text'] == 'yes':
                                    st.warning("⚠️ Predicted as screenshot but no text found")
                    
                    # Predictions vs Ground Truth
                    st.markdown("**Predictions vs Ground Truth:**")
                    
                    # Display all tasks dynamically
                    for task_key, task_info in tasks.items():
                        pred_val = row[f'pred_{task_key}']
                        gt_val = row[f'gt_{task_key}']
                        conf_val = row[f'conf_{task_key}']
                        correct = pred_val == gt_val
                        
                        badge_class = "correct-prediction" if correct else "incorrect-prediction"
                        st.markdown(f"""
                        <div class="prediction-badge {badge_class}">
                            {task_info['name']}: {pred_val} | GT: {gt_val} | Conf: {conf_val:.3f}
                        </div>
                        """, unsafe_allow_html=True)
                    
                    st.markdown("---")
    
    if len(df) > max_images:
        st.info(f"Showing {max_images} of {len(df)} filtered images. Use the slider above to show more images.")

def create_confusion_matrices(df, task_config):
    """Create confusion matrices for each classification task"""
    st.markdown("## 📊 Model Performance Analysis")
    
    tasks = {task['key']: task for task in task_config['tasks']}
    binary_tasks = [t for t in tasks.values() if t['type'] == 'binary']
    multi_class_tasks = [t for t in tasks.values() if t['type'] == 'multi_class']
    
    tab1, tab2, tab3 = st.tabs(["Confusion Matrices", "Confidence Analysis", "Task Performance"])
    
    with tab1:
        # Binary classification confusion matrices
        if binary_tasks:
            st.markdown("### Binary Classification Tasks")
            n_binary = len(binary_tasks)
            n_cols = min(2, n_binary)
            
            for i in range(0, n_binary, n_cols):
                cols = st.columns(n_cols)
                for j in range(n_cols):
                    if i + j < n_binary:
                        task = binary_tasks[i + j]
                        task_key = task['key']
                        
                        with cols[j]:
                            confusion_data = pd.crosstab(
                                df[f'gt_{task_key}'], 
                                df[f'pred_{task_key}'], 
                                margins=True
                            )
                            st.markdown(f"**{task['name']} Confusion Matrix**")
                            st.dataframe(confusion_data, use_container_width=True)
        
        # Multi-class confusion matrices
        if multi_class_tasks:
            st.markdown("### Multi-class Classification Tasks")
            for task in multi_class_tasks:
                task_key = task['key']
                st.markdown(f"**{task['name']} Confusion Matrix**")
                confusion_data = pd.crosstab(
                    df[f'gt_{task_key}'], 
                    df[f'pred_{task_key}'], 
                    margins=True
                )
                st.dataframe(confusion_data, use_container_width=True)
    
    with tab2:
        # Confidence analysis plots
        fig1 = create_confidence_distribution_plot(df, task_config)
        st.plotly_chart(fig1, use_container_width=True)
    
    with tab3:
        # Task-wise performance
        st.markdown("**Performance by Task**")
        
        performance_data = []
        for task_key, task_info in tasks.items():
            accuracy = df[f'correct_{task_key}'].mean()
            confidence = df[f'conf_{task_key}'].mean()
            
            performance_data.append({
                'Task': task_info['name'],
                'Type': task_info['type'],
                'Accuracy': accuracy,
                'Avg Confidence': confidence
            })
        
        performance_df = pd.DataFrame(performance_data)
        st.dataframe(performance_df, use_container_width=True)
        
        # Performance visualization
        fig = px.scatter(performance_df, x='Avg Confidence', y='Accuracy', 
                        color='Type', text='Task',
                        title="Task Performance: Accuracy vs Confidence")
        fig.update_traces(textposition="top center")
        st.plotly_chart(fig, use_container_width=True)

def main():
    """Main Streamlit application"""
    # Header
    st.markdown('<div class="main-header">🔍 Pseudoable Classifier Evaluation Dashboard</div>', 
                unsafe_allow_html=True)
    
    # Load task configuration
    task_config = load_task_config()
    if not task_config:
        st.error("❌ Could not load task configuration. Please ensure task_config.json exists.")
        st.info("Expected location: ./task_config.json or ./checkpoints/task_config.json")
        return
    
    st.success(f"✅ Loaded task configuration with {len(task_config['tasks'])} tasks")
    
    # Display task information
    with st.expander("📋 Task Configuration", expanded=False):
        for task in task_config['tasks']:
            st.markdown(f"**{task['name']}** ({task['type']})")
            st.markdown(f"- *Description:* {task['description']}")
            st.markdown(f"- *Labels:* {', '.join(task['labels'])}")
            st.markdown("---")
    
    # Check and run validation if needed
    if not run_validation_if_needed():
        return
    
    # Load validation results
    validation_data = load_validation_results()
    if not validation_data:
        st.error("❌ Could not load validation results. Please check if validation_results.json exists.")
        return
    
    # Process data
    df = process_validation_data(validation_data, task_config)
    if df is None or len(df) == 0:
        st.error("❌ No validation data found or data processing failed.")
        return
    
    # Show basic info
    st.success(f"✅ Loaded {len(df)} validation samples successfully!")
    
    # Overview metrics
    create_overview_metrics(df, validation_data, task_config)
    
    # Apply filters
    filtered_df, ocr_filter, enable_ocr = apply_filters(df, task_config)
    
    # Display results
    if len(filtered_df) > 0:
        # Performance analysis
        create_confusion_matrices(filtered_df, task_config)
        
        # Sample images
        display_sample_images(filtered_df, task_config, ocr_filter, enable_ocr)
    else:
        st.warning("⚠️ No samples match the current filter criteria. Please adjust your filters.")
    
    # Footer
    st.markdown("---")
    st.markdown("**📝 Instructions:**")
    st.markdown("1. Use the sidebar to filter by task confidence and prediction classes")
    st.markdown("2. Filter images by text content using OCR (if screenshot detection task is configured)")
    st.markdown("3. Adjust the number of images to display and sorting order")
    st.markdown("4. View model performance metrics and confusion matrices")
    st.markdown("5. Browse sample images with predictions vs ground truth")
    st.markdown("6. Green badges indicate correct predictions, red badges indicate incorrect predictions")

if __name__ == "__main__":
    main()