import gradio as gr
import torch
from transformers import ViTFeatureExtractor, ViTForImageClassification, pipeline
from PIL import Image
import easyocr
from aesthetic_score import AestheticScorer

# Initialize models
# Vision Transformer for Emotion Detection
feature_extractor = ViTFeatureExtractor.from_pretrained("nateraw/vit-base-beauty")
model = ViTForImageClassification.from_pretrained("nateraw/vit-base-beauty")
emotion_detector = pipeline("image-classification", model=model, feature_extractor=feature_extractor)

# OCR for Text Detection
reader = easyocr.Reader(['en'])

# Aesthetic Scoring Model (Placeholder)
aesthetic_scorer = AestheticScorer()

# Function to analyze a single thumbnail
def analyze_single_thumbnail(image):
    # Text detection using OCR
    ocr_result = reader.readtext(image)
    text_detected = " ".join([item[1] for item in ocr_result]) if ocr_result else "No text found"

    # Emotion detection for overall appeal
    emotions = emotion_detector(image)
    main_emotion = max(emotions, key=lambda x: x['score'])['label'] if emotions else "Unknown"

    # Aesthetic scoring
    aesthetic_score = aesthetic_scorer.score(image)

    return {
        "Detected Text": text_detected,
        "Emotion Detected": main_emotion,
        "Aesthetic Score": aesthetic_score,
    }

# Function to compare two thumbnails
def compare_thumbnails(image1, image2):
    result1 = analyze_single_thumbnail(image1)
    result2 = analyze_single_thumbnail(image2)

    # Determine which thumbnail might perform better based on scores
    if result1["Aesthetic Score"] > result2["Aesthetic Score"]:
        better_thumbnail = "Thumbnail 1 is likely better."
    elif result2["Aesthetic Score"] > result1["Aesthetic Score"]:
        better_thumbnail = "Thumbnail 2 is likely better."
    else:
        better_thumbnail = "Both thumbnails have similar appeal."

    return result1, result2, better_thumbnail

# Gradio Interface
iface = gr.Interface(
    fn=compare_thumbnails,
    inputs=[gr.Image(type="pil"), gr.Image(type="pil")],
    outputs=[
        gr.JSON(label="Thumbnail 1 Analysis"),
        gr.JSON(label="Thumbnail 2 Analysis"),
        gr.Textbox(label="Comparison Result"),
    ],
    title="YouTube Thumbnail Comparator",
    description="Upload two thumbnails to compare their effectiveness based on detected text, emotions, and aesthetic score.",
)

iface.launch()