app.py

# app.py - Enhanced with Text Quality Assessment
import gradio as gr
from core import IQA, TextAwareIQA
import numpy as np

# Initialize both simple and advanced models
simple_metric = IQA("qualiclip+")  
advanced_metric = TextAwareIQA("qualiclip+", text_weight=0.3)

def simple_assessment(image):
    """Original simple interface - backward compatible"""
    score = simple_metric(image)
    print(f"Simple score: {score}")  # Debug output
    return f"Quality Score: {score:.2f}/100" if score is not None else "Error: Invalid image"

def detailed_assessment(image, text_penalty_mode="balanced"):
    """Enhanced assessment with detailed breakdown"""
    try:
        # Get detailed analysis
        details = advanced_metric.evaluate(image, text_penalty_mode=text_penalty_mode)
        
        if details is None or 'error' in details:
            return "Error: Could not analyze image", "", ""
        
        # Format main result
        main_result = f"**Combined Quality Score: {details['combined_score']:.2f}/100**"
        
        # Format breakdown
        breakdown = f"""
## Score Breakdown:
- **Traditional IQA ({details['model_used']}): {details['traditional_score']:.2f}/100**
- **Text Quality: {details['text_analysis']['text_quality_score']:.2f}/100**
- **Text Weight in Final Score: {details['text_weight']*100:.0f}%**
- **Penalty Mode: {details.get('penalty_mode', 'balanced')}**
"""
        
        # Format text analysis details
        text_analysis = details['text_analysis']
        if text_analysis['text_detected']:
            text_details = f"""
## Text Analysis:
- **Text Regions Detected: {text_analysis['text_regions']}**
- **Low Quality Text Regions: {text_analysis['low_quality_regions']}**
- **Average OCR Confidence: {text_analysis['avg_confidence']:.3f}**
- **Assessment: {text_analysis['details']}**

### Text Quality Interpretation:
- **90-100**: Excellent text clarity, no distortions
- **70-89**: Good text quality, minor artifacts
- **50-69**: Readable but noticeable distortions
- **30-49**: Poor text quality, significant distortions
- **0-29**: Severely distorted or unreadable text
"""
        else:
            text_details = """
## Text Analysis:
- **No text detected in image**
- **Using traditional IQA score only**
"""
        
        return main_result, breakdown, text_details
        
    except Exception as e:
        error_msg = f"Error: {str(e)}"
        return error_msg, "", ""

def batch_comparison(image):
    """Compare different penalty modes side by side"""
    try:
        modes = ['lenient', 'balanced', 'strict']
        results = []
        
        for mode in modes:
            details = advanced_metric.evaluate(image, text_penalty_mode=mode)
            if details and 'error' not in details:
                results.append({
                    'mode': mode.title(),
                    'score': details['combined_score'],
                    'traditional': details['traditional_score'],
                    'text': details['text_analysis']['text_quality_score'],
                    'text_detected': details['text_analysis']['text_detected']
                })
        
        if not results:
            return "Error: Could not analyze image"
        
        # Format comparison
        comparison = "## Penalty Mode Comparison:\n\n"
        for result in results:
            comparison += f"**{result['mode']} Mode:**\n"
            comparison += f"- Combined Score: {result['score']:.2f}/100\n"
            comparison += f"- Traditional IQA: {result['traditional']:.2f}/100\n"
            comparison += f"- Text Quality: {result['text']:.2f}/100\n\n"
        
        if results[0]['text_detected']:
            comparison += """
### Mode Explanations:
- **Lenient**: Minimal penalty for text issues (10% weight)
- **Balanced**: Moderate penalty for text issues (30% weight)  
- **Strict**: Heavy penalty for text issues (50% weight)
"""
        else:
            comparison += "\n*No text detected - all modes return same score*"
        
        return comparison
        
    except Exception as e:
        return f"Error in batch comparison: {str(e)}"

# Create Gradio interface with tabs
with gr.Blocks(title="Enhanced Image Quality Assessment", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🔍 Enhanced Image Quality Assessment
    
    This tool evaluates image quality using both traditional perceptual metrics **AND** text-specific quality assessment.
    Perfect for detecting letter distortions in AI-generated images that traditional IQA tools miss.
    """)
    
    with gr.Tabs():
        # Simple tab - backward compatible
        with gr.TabItem("🚀 Quick Assessment"):
            gr.Markdown("### Simple quality score (backward compatible with your original setup)")
            with gr.Row():
                with gr.Column():
                    input_image_simple = gr.Image(
                        label="Upload Image", 
                        type="pil", 
                        format="png"
                    )
                    assess_btn_simple = gr.Button("Assess Quality", variant="primary")
                with gr.Column():
                    output_simple = gr.Textbox(
                        label="Quality Score", 
                        lines=2
                    )
            
            assess_btn_simple.click(
                simple_assessment, 
                inputs=input_image_simple, 
                outputs=output_simple
            )
        
        # Detailed tab - new enhanced features
        with gr.TabItem("🔬 Detailed Analysis"):
            gr.Markdown("### Comprehensive quality assessment with text-specific evaluation")
            with gr.Row():
                with gr.Column():
                    input_image_detailed = gr.Image(
                        label="Upload Image", 
                        type="pil", 
                        format="png"
                    )
                    penalty_mode = gr.Radio(
                        choices=["lenient", "balanced", "strict"],
                        value="balanced",
                        label="Text Penalty Mode",
                        info="How harshly to penalize text quality issues"
                    )
                    assess_btn_detailed = gr.Button("Detailed Analysis", variant="primary")
                
                with gr.Column():
                    output_main = gr.Markdown(label="Final Score")
                    output_breakdown = gr.Markdown(label="Score Breakdown")
                    output_text_details = gr.Markdown(label="Text Analysis")
            
            assess_btn_detailed.click(
                detailed_assessment,
                inputs=[input_image_detailed, penalty_mode],
                outputs=[output_main, output_breakdown, output_text_details]
            )
        
        # Comparison tab
        with gr.TabItem("⚖️ Mode Comparison"):
            gr.Markdown("### Compare how different penalty modes affect the final score")
            with gr.Row():
                with gr.Column():
                    input_image_comparison = gr.Image(
                        label="Upload Image", 
                        type="pil", 
                        format="png"
                    )
                    compare_btn = gr.Button("Compare All Modes", variant="primary")
                
                with gr.Column():
                    output_comparison = gr.Markdown(label="Comparison Results")
            
            compare_btn.click(
                batch_comparison,
                inputs=input_image_comparison,
                outputs=output_comparison
            )
    
    gr.Markdown("""
    ---
    ### 💡 How It Works:
    - **Traditional IQA**: Uses qalign (or your chosen model) for overall perceptual quality
    - **Text Quality**: Uses OCR confidence scores to detect letter distortions and rendering issues
    - **Combined Score**: Weighted combination that penalizes text quality problems
    
    ### 🎯 Perfect For:
    - Detecting letter distortions in AI-generated images
    - Evaluating text rendering quality in synthetic images  
    - Getting more accurate quality assessments for images containing text
    
    ### ⚙️ Penalty Modes:
    - **Lenient**: Light penalty for text issues (good for artistic images)
    - **Balanced**: Moderate penalty (recommended for most use cases)
    - **Strict**: Heavy penalty for any text problems (best for document quality)
    """)

if __name__ == "__main__":
    demo.launch(share=True)