image_learning.py

#!/usr/bin/env python3
"""
Image Generation Learning System
Learns from user feedback to generate better images over time
"""

import json
import time
from pathlib import Path
from typing import Dict, List, Optional, Tuple
import os

class ImageLearningSystem:
    """
    Self-learning image generation system
    Remembers what users liked and improves accordingly
    """
    
    def __init__(self, data_dir: str = "noahski_data"):
        self.data_dir = Path(data_dir)
        self.data_dir.mkdir(parents=True, exist_ok=True)
        self.learning_file = self.data_dir / "image_learning_data.json"
        self.feedback_file = self.data_dir / "image_feedback.json"
        self.learning_data = self._load_learning_data()
        self.feedback_history = self._load_feedback()
    
    def _load_learning_data(self) -> Dict:
        """Load learning data"""
        if self.learning_file.exists():
            try:
                with open(self.learning_file, 'r') as f:
                    return json.load(f)
            except:
                return self._default_learning_data()
        return self._default_learning_data()
    
    def _default_learning_data(self) -> Dict:
        """Default learning parameters"""
        return {
            'successful_prompts': {},
            'parameter_scores': {
                'shading_intensity': 0.8,
                'texture_density': 0.6,
                'color_saturation': 0.85,
                'detail_level': 0.7,
                'lighting_strength': 0.75
            },
            'learned_colors': {},
            'total_feedback': 0,
            'positive_count': 0,
            'improvement_rate': 1.0
        }
    
    def _load_feedback(self) -> List:
        """Load feedback history"""
        if self.feedback_file.exists():
            try:
                with open(self.feedback_file, 'r') as f:
                    return json.load(f)
            except:
                return []
        return []
    
    def record_feedback(self, prompt: str, feedback_type: str, params: Dict = None) -> None:
        """Record user feedback for learning"""
        feedback = {
            'timestamp': time.time(),
            'prompt': prompt,
            'type': feedback_type,  # 'positive' or 'negative'
            'params': params or {}
        }
        
        self.feedback_history.append(feedback)
        
        # Update learning data
        self.learning_data['total_feedback'] += 1
        if feedback_type == 'positive':
            self.learning_data['positive_count'] += 1
        
        # Update prompt success rate
        if prompt not in self.learning_data['successful_prompts']:
            self.learning_data['successful_prompts'][prompt] = {'positive': 0, 'total': 0}
        
        self.learning_data['successful_prompts'][prompt]['total'] += 1
        if feedback_type == 'positive':
            self.learning_data['successful_prompts'][prompt]['positive'] += 1
        
        # Calculate improvement metrics
        if self.learning_data['total_feedback'] > 0:
            success_rate = self.learning_data['positive_count'] / self.learning_data['total_feedback']
            self.learning_data['improvement_rate'] = min(1.2, 0.8 + success_rate * 0.4)
        
        self._save_learning_data()
        self._save_feedback()
    
    def get_improved_parameters(self, base_params: Dict) -> Dict:
        """Get improved generation parameters based on learning"""
        improved = base_params.copy()
        
        # Apply learned improvements
        for param, value in self.learning_data['parameter_scores'].items():
            if param in improved:
                # Blend base with learned value
                blend = 0.6  # 60% learned, 40% original
                improved[param] = base_params[param] * (1 - blend) + value * blend
        
        # Apply improvement rate
        rate = self.learning_data['improvement_rate']
        for param in ['shading_intensity', 'texture_density', 'color_saturation']:
            if param in improved:
                improved[param] = min(1.0, improved[param] * rate)
        
        return improved
    
    def get_success_rate(self, prompt: str = None) -> float:
        """Get success rate for a prompt or overall"""
        if prompt and prompt in self.learning_data['successful_prompts']:
            data = self.learning_data['successful_prompts'][prompt]
            if data['total'] > 0:
                return data['positive'] / data['total']
        
        if self.learning_data['total_feedback'] > 0:
            return self.learning_data['positive_count'] / self.learning_data['total_feedback']
        
        return 0.5  # Default neutral
    
    def _save_learning_data(self) -> None:
        """Save learning data"""
        try:
            with open(self.learning_file, 'w') as f:
                json.dump(self.learning_data, f, indent=2)
        except:
            pass
    
    def _save_feedback(self) -> None:
        """Save feedback history"""
        try:
            with open(self.feedback_file, 'w') as f:
                json.dump(self.feedback_history, f, indent=2)
        except:
            pass
    
    def get_stats(self) -> Dict:
        """Get learning stats"""
        return {
            'total_feedback': self.learning_data['total_feedback'],
            'positive_count': self.learning_data['positive_count'],
            'success_rate': self.get_success_rate(),
            'improvement_rate': self.learning_data['improvement_rate']
        }


# Global instance
_learning_system = None


def get_learning_system() -> ImageLearningSystem:
    """Get global learning system instance"""
    global _learning_system
    if _learning_system is None:
        _learning_system = ImageLearningSystem()
    return _learning_system


def get_improved_generation_params(base_params: Dict = None) -> Dict:
    """Get generation parameters improved by learning"""
    if base_params is None:
        base_params = {
            'shading_intensity': 0.8,
            'texture_density': 0.6,
            'color_saturation': 0.85,
            'detail_level': 0.7,
            'lighting_strength': 0.75
        }
    
    system = get_learning_system()
    return system.get_improved_parameters(base_params)


def record_image_feedback(prompt: str, feedback_type: str) -> None:
    """Record feedback for an image"""
    system = get_learning_system()
    system.record_feedback(prompt, feedback_type)