emoji_processor.py

"""
Core emoji processing logic for the Emoji Mashup application.
"""

from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import requests
from PIL import Image
from io import BytesIO

from config import CONFIG
from utils import logger, kitchen_txt_to_dict

class EmojiProcessor:
    def __init__(self, model_name=CONFIG["model_name"]):
        """Initialize the emoji processor with the specified model.
        
        Args:
            model_name: Name of the sentence transformer model to use
        """
        logger.info(f"Loading model: {model_name}")
        self.model = SentenceTransformer(model_name)
        self.emotion_dict = {}
        self.event_dict = {}
        self.emotion_embeddings = {}
        self.event_embeddings = {}

    def load_emoji_dictionaries(self, emotion_file=CONFIG["emotion_file"], item_file=CONFIG["item_file"]):
        """Load emoji dictionaries from text files.
        
        Args:
            emotion_file: Path to the emotion emoji file
            item_file: Path to the item emoji file
        """
        logger.info("Loading emoji dictionaries")
        self.emotion_dict = kitchen_txt_to_dict(emotion_file)
        self.event_dict = kitchen_txt_to_dict(item_file)
        
        # Precompute embeddings
        logger.info("Computing embeddings for emoji dictionaries")
        self.emotion_embeddings = {emoji: self.model.encode(desc) for emoji, desc in self.emotion_dict.items()}
        self.event_embeddings = {emoji: self.model.encode(desc) for emoji, desc in self.event_dict.items()}

    def find_top_emojis(self, embedding, emoji_embeddings, top_n=1):
        """Find top matching emojis based on cosine similarity.
        
        Args:
            embedding: Sentence embedding to compare
            emoji_embeddings: Dictionary of emoji embeddings
            top_n: Number of top emojis to return
            
        Returns:
            List of top matching emojis
        """
        similarities = [
            (emoji, cosine_similarity([embedding], [e_embed])[0][0])
            for emoji, e_embed in emoji_embeddings.items()
        ]
        similarities.sort(key=lambda x: x[1], reverse=True)
        return [emoji for emoji, _ in similarities[:top_n]]

    def get_emoji_mashup_url(self, emoji1, emoji2, size=CONFIG["default_size"]):
        """Generate URL for emoji mashup.
        
        Args:
            emoji1: First emoji character
            emoji2: Second emoji character
            size: Image size in pixels
            
        Returns:
            URL for the emoji mashup
        """
        return f"{CONFIG['emoji_kitchen_url'].format(emoji1=emoji1, emoji2=emoji2)}?size={size}"

    def fetch_mashup_image(self, url):
        """Fetch emoji mashup image from URL.
        
        Args:
            url: URL of the emoji mashup image
            
        Returns:
            PIL Image object or None if fetch failed
        """
        try:
            response = requests.get(url)
            if response.status_code == 200 and "image" in response.headers.get("Content-Type", ""):
                return Image.open(BytesIO(response.content))
            else:
                logger.warning(f"Failed to fetch image: Status code {response.status_code}")
                return None
        except Exception as e:
            logger.error(f"Error fetching image: {e}")
            return None

    def sentence_to_emojis(self, sentence):
        """Process sentence to find matching emojis and generate mashup.
        
        Args:
            sentence: User input text
            
        Returns:
            Tuple of (emotion emoji, event emoji, mashup image)
        """
        if not sentence.strip():
            return "❓", "❓", None
            
        try:
            embedding = self.model.encode(sentence)
            
            top_emotion = self.find_top_emojis(embedding, self.emotion_embeddings, top_n=1)[0]
            top_event = self.find_top_emojis(embedding, self.event_embeddings, top_n=1)[0]
            
            mashup_url = self.get_emoji_mashup_url(top_emotion, top_event)
            mashup_image = self.fetch_mashup_image(mashup_url)
            
            return top_emotion, top_event, mashup_image
        except Exception as e:
            logger.error(f"Error processing sentence: {e}")
            return "❌", "❌", None