app.py

"""
Viral Images — Image Scoring System
Upload an image, score it, and improve it.

Deployed on HuggingFace Spaces: https://huggingface.co/spaces/Babajaan/viral-images
"""

import os
import sys
import time
import traceback
from pathlib import Path

# ── Compatibility shim: gradio 5.0.0 + huggingface_hub >= 0.27.0 ──
# HfFolder was removed in huggingface_hub 0.27+ but gradio 5.0.0 still imports it.
# We inject a minimal stub directly into huggingface_hub's module dict BEFORE gradio loads.
try:
    import sys
    import huggingface_hub
    hf_mod = sys.modules.get("huggingface_hub") or huggingface_hub
    if "HfFolder" not in hf_mod.__dict__:
        class _HfFolderCompat:
            @staticmethod
            def get_token():
                return os.environ.get("HF_TOKEN", None)
            @staticmethod
            def save_token(token):
                pass
            @staticmethod
            def delete_token():
                pass
        # Inject into the real module dict so `from huggingface_hub import HfFolder` works
        hf_mod.__dict__["HfFolder"] = _HfFolderCompat()
        if hasattr(hf_mod, "__all__") and "HfFolder" not in hf_mod.__all__:
            hf_mod.__all__.append("HfFolder")
        print("[ViralImages] HfFolder compatibility shim applied.")
    else:
        print("[ViralImages] HfFolder already present, no shim needed.")
except Exception as e:
    print(f"[ViralImages] HfFolder shim failed (non-fatal): {e}")

# Ensure our modules are importable
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import gradio as gr
import pandas as pd
from PIL import Image

from utils.preprocessing import preprocess_image, validate_image
from utils.formatting import (
    format_overall_score,
    format_sub_scores,
    format_strengths,
    format_weaknesses,
    format_suggestions,
    format_projected_improvement,
    format_comparison,
)
from scoring.engine import ScoringEngine, ScoreResponse
from recommendations.engine import RecommendationEngine, compute_improvement_potential
from compare.comparator import compare_scores
from models.loader import ModelLoader

# ───────────────────────────────────────────────────────────────
# Global state (models loaded once at startup)
# ───────────────────────────────────────────────────────────────

MODEL_LOADER = None
SCORING_ENGINE = None
REC_ENGINE = None


def _init_engines():
    """Initialize scoring engines on startup."""
    global MODEL_LOADER, SCORING_ENGINE, REC_ENGINE
    
    print("[ViralImages] Initializing engines...")
    t0 = time.time()
    
    MODEL_LOADER = ModelLoader()
    
    # Try to find config files
    config_paths = {
        "scoring": None,
        "rules": None,
        "model": None,
    }
    
    for base in ["configs", "viral-images/configs", os.path.join(os.path.dirname(__file__), "configs")]:
        for key, fname in [("scoring", "scoring_weights.yaml"), 
                           ("rules", "suggestion_rules.yaml"),
                           ("model", "model_config.yaml")]:
            if config_paths[key] is None:
                p = os.path.join(base, fname)
                if os.path.exists(p):
                    config_paths[key] = p
    
    SCORING_ENGINE = ScoringEngine(
        model_loader=MODEL_LOADER,
        config_path=config_paths["scoring"]
    )
    REC_ENGINE = RecommendationEngine(
        rules_path=config_paths["rules"]
    )
    
    # Pre-warm (load CLIP if available)
    try:
        SCORING_ENGINE.warmup()
        print(f"[ViralImages] Engines ready in {time.time()-t0:.1f}s")
    except Exception as e:
        print(f"[ViralImages] Warmup failed: {e}")
        print("[ViralImages] Running in fallback mode (no CLIP).")


# Initialize on module load
_init_engines()

# ───────────────────────────────────────────────────────────────
# Score Image Tab
# ───────────────────────────────────────────────────────────────

def score_image(image_input, concept, audience, use_case, detail_level):
    """
    Score an uploaded image.
    
    Returns:
        (overall_html, scores_df, strengths_text, weaknesses_text,
         suggestions_text, projected_html)
    """
    if image_input is None:
        return (
            "<div style='padding:20px; color:#e74c3c;'><b>⚠️ Please upload an image first.</b></div>",
            pd.DataFrame(),
            "", "", "",
            "<div style='padding:12px; color:#999;'>No image to analyze.</div>"
        )
    
    try:
        img = preprocess_image(image_input)
        val_error = validate_image(img)
        if val_error:
            return (
                f"<div style='padding:20px; color:#e74c3c;'><b>⚠️ {val_error}</b></div>",
                pd.DataFrame(),
                "", "", "",
                "<div style='padding:12px; color:#999;'>Image validation failed.</div>"
            )
        
        # Score
        response = SCORING_ENGINE.score(img, concept or "", audience, use_case)
        
        # Generate suggestions
        suggestions = REC_ENGINE.generate(
            response.sub_scores,
            response.raw_features,
            use_case
        )
        response.suggestions = [
            {
                "priority": s.priority,
                "sub_score_target": s.sub_score_target,
                "message": s.message,
                "projected_gain": s.projected_gain,
                "trigger_feature": s.trigger_feature,
                "trigger_value": s.trigger_value,
            }
            for s in suggestions
        ]
        
        # Projected improvement
        projected = REC_ENGINE.estimate_improvement(
            response.sub_scores,
            suggestions,
            use_case
        )
        response.projected_improvement = projected
        
        # Format outputs
        overall_html = format_overall_score(response.overall_score)
        scores_df = format_sub_scores(response.sub_scores)
        strengths_html = format_strengths(response.strengths)
        weaknesses_html = format_weaknesses(response.weaknesses)
        suggestions_html = format_suggestions(response.suggestions)
        projected_html = format_projected_improvement(
            response.overall_score, projected
        )
        
        # Metadata footer
        meta = response.metadata
        meta_html = f"<div style='font-size:11px; color:#999; margin-top:8px;'>⏱ {meta.get('processing_time_ms', '?')}ms | Mode: {meta.get('neural_richness_mode', '?')} | Confidence: {response.confidence:.0%}</div>"
        
        return (
            overall_html + meta_html,
            scores_df,
            strengths_html,
            weaknesses_html,
            suggestions_html,
            projected_html,
        )
        
    except Exception as e:
        traceback.print_exc()
        return (
            f"<div style='padding:20px; color:#e74c3c;'><b>⚠️ Error during scoring:</b><br/>{str(e)}</div>",
            pd.DataFrame(),
            "", "", "",
            "<div style='padding:12px; color:#999;'>Scoring failed.</div>"
        )


# ───────────────────────────────────────────────────────────────
# Compare Mode Tab
# ───────────────────────────────────────────────────────────────

def compare_images(original_img, revised_img, concept, audience, use_case):
    """
    Compare before/after images.
    
    Returns:
        (original_scores_label, revised_scores_label, comparison_html, next_steps_text)
    """
    if original_img is None or revised_img is None:
        return (
            {}, {},
            "<div style='padding:20px; color:#e74c3c;'><b>⚠️ Please upload both images.</b></div>",
            "Upload both original and revised images to compare."
        )
    
    try:
        # Score both
        original = SCORING_ENGINE.score(original_img, concept or "", audience, use_case)
        revised = SCORING_ENGINE.score(revised_img, concept or "", audience, use_case)
        
        # Compare
        result = compare_scores(original, revised)
        
        # Format
        orig_label = {k: v / 100 for k, v in original.sub_scores.items()}
        rev_label = {k: v / 100 for k, v in revised.sub_scores.items()}
        
        comp_html = format_comparison_result(result)
        
        next_steps = "\n".join(
            f"• {rec['message']}" for rec in result["next_recommendations"]
        ) if result["next_recommendations"] else "No further improvements needed — your revised image scores well!"
        
        return (
            orig_label,
            rev_label,
            comp_html,
            next_steps,
        )
        
    except Exception as e:
        traceback.print_exc()
        return (
            {}, {},
            f"<div style='padding:20px; color:#e74c3c;'><b>⚠️ Error:</b><br/>{str(e)}</div>",
            "Comparison failed. Please try again."
        )


def format_comparison_result(result: dict) -> str:
    """Format comparison as HTML."""
    orig = result["original_overall"]
    rev = result["revised_overall"]
    delta = result["delta_overall"]
    
    color = "#27ae60" if delta > 0 else "#e74c3c" if delta < 0 else "#999"
    emoji = "📈" if delta > 0 else "📉" if delta < 0 else "➡️"
    
    html = f"""
    <div style="padding:16px; background:#f8f9fa; border-radius:8px;">
        <div style="display:flex; justify-content:space-between; align-items:center; margin-bottom:16px;">
            <div style="text-align:center; flex:1;">
                <div style="font-size:12px; color:#666;">Original</div>
                <div style="font-size:36px; font-weight:bold; color:#999;">{orig:.0f}</div>
            </div>
            <div style="font-size:28px; color:#ccc; padding:0 16px;">→</div>
            <div style="text-align:center; flex:1;">
                <div style="font-size:12px; color:#666;">Revised</div>
                <div style="font-size:36px; font-weight:bold; color:{color};">{rev:.0f}</div>
            </div>
        </div>
        <div style="text-align:center; color:{color}; font-weight:bold; font-size:16px; margin-bottom:12px;">
            {emoji} Overall: {abs(delta):.0f} point{'s' if abs(delta) != 1 else ''} {'gained' if delta > 0 else 'lost' if delta < 0 else 'unchanged'}
        </div>
    """
    
    if result["improvements"]:
        html += "<div style='font-weight:bold; color:#27ae60; margin-top:8px;'>✅ Improvements:</div><ul style='margin:4px 0; padding-left:20px;'>"
        for name, val in result["improvements"]:
            html += f"<li><b>{name.replace('_', ' ').title()}:</b> +{val:.0f} points</li>"
        html += "</ul>"
    
    if result["regressions"]:
        html += "<div style='font-weight:bold; color:#e74c3c; margin-top:8px;'>⚠️ Regressions:</div><ul style='margin:4px 0; padding-left:20px;'>"
        for name, val in result["regressions"]:
            html += f"<li><b>{name.replace('_', ' ').title()}:</b> {val:.0f} points</li>"
        html += "</ul>"
    
    if result["next_recommendations"]:
        html += "<div style='font-weight:bold; color:#2980b9; margin-top:8px;'>📋 Next Steps:</div><div style='margin:4px 0;'>"
        for rec in result["next_recommendations"]:
            html += f"<div style='margin:4px 0; padding:6px 10px; background:#e3f2fd; border-radius:4px; font-size:13px;'>{rec['message']}</div>"
        html += "</div>"
    
    html += "</div>"
    return html


# ───────────────────────────────────────────────────────────────
# Gradio App
# ───────────────────────────────────────────────────────────────

ABOUT_TEXT = """
## 🧠 About Viral Images

**Viral Images** scores any image across 8 dimensions and suggests concrete improvements.

### What It Measures
| Dimension | Description |
|-----------|-------------|
| **Concept Match** | Does the image convey your intended topic? |
| **Visual Focus** | Is there a clear focal point that draws the eye? |
| **Readability** | Is text legible and well-sized? |
| **Complexity Balance** | Is detail level appropriate — not empty, not cluttered? |
| **Communication Clarity** | Is the visual hierarchy clear and well-organized? |
| **Predicted Neural Richness** | How visually engaging is the image? |
| **Memorability** | Will viewers remember it after one viewing? |
| **Improvement Potential** | How much room for improvement remains? |

### How to Use
1. **Upload an image** (JPG, PNG, or any common format)
2. **Describe the concept** — what should this image communicate?
3. **Select audience and use case**
4. Click **"Score Image"**
5. Review scores, strengths, weaknesses, and suggestions

### Compare Mode
Upload an **original** and a **revised** version to see what improved and what to focus on next.

### ⚠️ Important Disclaimer
- Scores are **AI-predicted proxies**, not measurements of real brain activity.
- Scores do **not guarantee engagement, virality, or aesthetic quality**.
- Use as **directional feedback** to complement your own judgment.
- Predicted Neural Richness is estimated via computational proxy models, not from actual fMRI data.

### Technical Details
Models used: CLIP ViT-B/32 (concept matching), heuristic analysis (OpenCV), 
saliency heuristics, OCR heuristics, and aesthetic proxies.
Neural richness is estimated via proxy, not measured.
"""


def build_app():
    """Build and return the Gradio app."""
    
    with gr.Blocks() as demo:
        gr.Markdown(
            """
            # 🧠 Viral Images
            ### *Upload an image, score it, and improve it.*
            """
        )
        
        # ── TAB 1: Score Image ──
        with gr.Tab("📊 Score Image"):
            with gr.Row():
                # Input column
                with gr.Column(scale=2):
                    image_input = gr.Image(
                        type="pil",
                        label="📷 Upload Image",
                        sources=["upload", "clipboard"],
                        height=300,
                    )
                    
                    concept_input = gr.Textbox(
                        label="💡 Concept / Theme",
                        placeholder="e.g., photosynthesis process, product launch, data visualization...",
                        lines=1,
                    )
                    
                    audience_input = gr.Dropdown(
                        choices=[
                            "General",
                            "Children (K-8)",
                            "High School Students",
                            "College Students",
                            "Professionals",
                            "Researchers",
                            "General Public",
                            "Social Media Audience",
                        ],
                        label="👥 Target Audience",
                        value="General",
                    )
                    
                    usecase_input = gr.Dropdown(
                        choices=[
                            ("Social Media Post", "social_media"),
                            ("Thumbnail", "thumbnail"),
                            ("Educational", "educational"),
                            ("Scientific Figure", "scientific_figure"),
                            ("Infographic", "infographic"),
                            ("Presentation Slide", "presentation_slide"),
                            ("Marketing", "marketing"),
                            ("General", "default"),
                        ],
                        label="🎯 Use Case",
                        value="social_media",
                    )
                    
                    score_btn = gr.Button(
                        "🔍 Score Image",
                        variant="primary",
                        size="lg",
                    )
                    
                    with gr.Accordion("⚙️ Advanced Options", open=False):
                        detail_slider = gr.Slider(
                            1, 5, value=3, step=1,
                            label="Detail Level",
                        )
                
                # Output column
                with gr.Column(scale=3):
                    overall_score = gr.HTML(
                        label="Overall Score",
                        value="<div style='padding:40px; text-align:center; color:#999;'>Upload an image and click 'Score Image' to see results.</div>"
                    )
                    
                    with gr.Row():
                        strengths_box = gr.HTML(
                            label="✅ Strengths",
                            value="",
                        )
                        weaknesses_box = gr.HTML(
                            label="⚠️ Weaknesses",
                            value="",
                        )
                    
                    score_chart = gr.BarPlot(
                        x="Sub-Score",
                        y="Score",
                        color="Sub-Score",
                        title="Score Breakdown",
                        y_lim=[0, 100],
                        height=280,
                        value=pd.DataFrame({"Sub-Score": [], "Score": []}),
                    )
                    
                    suggestions_box = gr.HTML(
                        label="💡 Improvement Suggestions",
                        value="",
                    )
                    
                    projected_html = gr.HTML(
                        label="Projected Improvement",
                        value="",
                    )
            
            score_btn.click(
                fn=score_image,
                inputs=[image_input, concept_input, audience_input, usecase_input, detail_slider],
                outputs=[overall_score, score_chart, strengths_box, weaknesses_box, suggestions_box, projected_html],
                concurrency_id="inference",
                concurrency_limit=1,
            )
        
        # ── TAB 2: Compare Mode ──
        with gr.Tab("🔄 Compare Mode"):
            with gr.Row():
                before_img = gr.Image(
                    type="pil",
                    label="📷 Original Image",
                    sources=["upload"],
                    height=250,
                )
                after_img = gr.Image(
                    type="pil",
                    label="📷 Revised Image",
                    sources=["upload"],
                    height=250,
                )
            
            with gr.Row():
                compare_concept = gr.Textbox(
                    label="💡 Concept",
                    placeholder="Same concept as original...",
                    lines=1,
                )
                compare_audience = gr.Dropdown(
                    choices=[
                        "General", "Children (K-8)", "High School Students",
                        "College Students", "Professionals", "Researchers",
                        "General Public", "Social Media Audience",
                    ],
                    label="👥 Audience",
                    value="General",
                )
                compare_usecase = gr.Dropdown(
                    choices=[
                        ("Social Media Post", "social_media"),
                        ("Thumbnail", "thumbnail"),
                        ("Educational", "educational"),
                        ("Scientific Figure", "scientific_figure"),
                        ("Infographic", "infographic"),
                        ("Presentation Slide", "presentation_slide"),
                        ("Marketing", "marketing"),
                        ("General", "default"),
                    ],
                    label="🎯 Use Case",
                    value="social_media",
                )
            
            compare_btn = gr.Button(
                "🔄 Compare",
                variant="primary",
                size="lg",
            )
            
            with gr.Row():
                before_scores = gr.Label(
                    label="Original Scores",
                    value={},
                )
                after_scores = gr.Label(
                    label="Revised Scores",
                    value={},
                )
            
            comparison_html = gr.HTML(
                label="Comparison Results",
                value="<div style='padding:20px; text-align:center; color:#999;'>Upload both images to compare.</div>"
            )
            
            next_steps = gr.HTML(
                label="📋 Next Recommended Edits",
                value="",
            )
            
            compare_btn.click(
                fn=compare_images,
                inputs=[before_img, after_img, compare_concept, compare_audience, compare_usecase],
                outputs=[before_scores, after_scores, comparison_html, next_steps],
                concurrency_id="inference",
                concurrency_limit=1,
            )
        
        # ── TAB 3: About ──
        with gr.Tab("ℹ️ About"):
            gr.Markdown(ABOUT_TEXT)
        
        # Footer
        gr.Markdown(
            """
            <div style="margin-top:20px; padding:12px; background:#f8f9fa; border-radius:8px; font-size:12px; color:#666;">
            ⚠️ <b>Disclaimer:</b> Scores are AI-predicted proxies, not measurements of brain activity or guaranteed engagement. 
            Use as directional feedback, not absolute truth. 
            <a href="https://github.com/facebookresearch/algonauts-2025" target="_blank">TRIBE</a> is used only as inspiration for the neural richness proxy.
            </div>
            """
        )
    
    return demo


# ───────────────────────────────────────────────────────────────
# Launch
# ───────────────────────────────────────────────────────────────

if __name__ == "__main__":
    demo = build_app()
    demo.queue(max_size=10, api_open=False).launch(
        show_api=False,
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
    )