model/vision_ai.py

"""
🖼️ EKALAVYA Vision AI - Image Generation & Analysis
🎨 Generate Images • 👁️ Understand Images • 📊 Analyze Visual Data
"""

import numpy as np
from typing import Dict, List, Optional, Tuple
from PIL import Image, ImageDraw, ImageFilter
import io
import base64

class VisionAI:
    """🖼️ Advanced image generation and analysis"""
    
    def __init__(self):
        self.generation_history = []
        self.analysis_cache = {}
    
    # 🎨 IMAGE GENERATION
    
    def generate_image(self, prompt: str, style: str = "realistic", size: Tuple[int, int] = (512, 512)) -> Dict:
        """🎨 Generate image from text prompt"""
        print(f"🎨 Generating image: {prompt}")
        
        # Create image based on prompt keywords
        img = Image.new('RGB', size, color='white')
        draw = ImageDraw.Draw(img)
        
        # Parse prompt for visual elements
        elements = self._parse_prompt(prompt)
        
        # Generate visual elements
        if 'landscape' in elements or 'nature' in elements:
            self._draw_landscape(draw, size)
        elif 'portrait' in elements or 'person' in elements:
            self._draw_portrait(draw, size)
        elif 'abstract' in elements or 'art' in elements:
            self._draw_abstract(draw, size)
        elif 'geometric' in elements or 'shape' in elements:
            self._draw_geometric(draw, size)
        else:
            self._draw_generic(draw, size, prompt)
        
        # Convert to base64
        buffer = io.BytesIO()
        img.save(buffer, format='PNG')
        img_base64 = base64.b64encode(buffer.getvalue()).decode()
        
        return {
            "status": "✅ success",
            "prompt": prompt,
            "style": style,
            "size": size,
            "image_base64": img_base64,
            "emoji": "🎨",
            "message": "🎨 Image generated successfully!"
        }
    
    def _parse_prompt(self, prompt: str) -> List[str]:
        """🔍 Parse prompt for visual elements"""
        elements = []
        prompt_lower = prompt.lower()
        
        keywords = {
            'landscape': ['mountain', 'forest', 'ocean', 'sky', 'sunset', 'sunrise'],
            'portrait': ['person', 'face', 'portrait', 'human', 'people'],
            'abstract': ['abstract', 'art', 'creative', 'creative', 'artistic'],
            'geometric': ['geometric', 'shape', 'circle', 'square', 'triangle', 'pattern']
        }
        
        for category, words in keywords.items():
            if any(word in prompt_lower for word in words):
                elements.append(category)
        
        return elements if elements else ['generic']
    
    def _draw_landscape(self, draw: ImageDraw.Draw, size: Tuple[int, int]):
        """🏔️ Draw landscape"""
        width, height = size
        
        # Sky gradient
        for y in range(height // 2):
            color = (135, 206, 235, int(255 * (1 - y / (height // 2))))
            draw.line([(0, y), (width, y)], fill=(135, 206, 235))
        
        # Mountains
        points = [(0, height // 2)]
        for x in range(0, width, 50):
            y = height // 2 + np.random.randint(-50, 50)
            points.append((x, y))
        points.append((width, height // 2))
        points.append((width, height))
        points.append((0, height))
        draw.polygon(points, fill=(34, 139, 34))
        
        # Sun
        draw.ellipse([width - 100, 50, width - 50, 100], fill=(255, 215, 0))
    
    def _draw_portrait(self, draw: ImageDraw.Draw, size: Tuple[int, int]):
        """👤 Draw portrait"""
        width, height = size
        center_x, center_y = width // 2, height // 2
        
        # Head
        draw.ellipse([center_x - 80, center_y - 100, center_x + 80, center_y + 100], 
                    fill=(255, 218, 185))
        
        # Eyes
        draw.ellipse([center_x - 40, center_y - 20, center_x - 20, center_y], fill=(0, 0, 0))
        draw.ellipse([center_x + 20, center_y - 20, center_x + 40, center_y], fill=(0, 0, 0))
        
        # Mouth
        draw.arc([center_x - 30, center_y + 30, center_x + 30, center_y + 60], 
                0, 180, fill=(0, 0, 0), width=2)
    
    def _draw_abstract(self, draw: ImageDraw.Draw, size: Tuple[int, int]):
        """🎨 Draw abstract art"""
        width, height = size
        
        # Random colorful shapes
        for _ in range(20):
            x1 = np.random.randint(0, width)
            y1 = np.random.randint(0, height)
            x2 = x1 + np.random.randint(50, 150)
            y2 = y1 + np.random.randint(50, 150)
            color = tuple(np.random.randint(0, 255, 3))
            
            if np.random.rand() > 0.5:
                draw.ellipse([x1, y1, x2, y2], fill=color)
            else:
                draw.rectangle([x1, y1, x2, y2], fill=color)
    
    def _draw_geometric(self, draw: ImageDraw.Draw, size: Tuple[int, int]):
        """🔷 Draw geometric patterns"""
        width, height = size
        
        # Grid pattern
        for x in range(0, width, 50):
            draw.line([(x, 0), (x, height)], fill=(100, 100, 100), width=2)
        for y in range(0, height, 50):
            draw.line([(0, y), (width, y)], fill=(100, 100, 100), width=2)
        
        # Shapes
        for i in range(5):
            x = (i + 1) * 100
            y = (i + 1) * 80
            draw.polygon([(x, y), (x + 50, y + 50), (x, y + 50)], fill=(255, 0, 0))
    
    def _draw_generic(self, draw: ImageDraw.Draw, size: Tuple[int, int], prompt: str):
        """🖼️ Draw generic image"""
        width, height = size
        
        # Background
        draw.rectangle([0, 0, width, height], fill=(240, 240, 240))
        
        # Text
        draw.text((50, height // 2), prompt[:50], fill=(0, 0, 0))
    
    # 👁️ IMAGE ANALYSIS
    
    def analyze_image(self, image_data: str) -> Dict:
        """👁️ Analyze image content"""
        print("👁️ Analyzing image...")
        
        # Decode base64 image
        try:
            if image_data.startswith('data:image'):
                image_data = image_data.split(',')[1]
            image_bytes = base64.b64decode(image_data)
            img = Image.open(io.BytesIO(image_bytes))
        except Exception as e:
            return {"error": f"❌ Cannot decode image: {e}"}
        
        # Analyze image properties
        width, height = img.size
        mode = img.mode
        
        # Convert to RGB if needed
        if mode != 'RGB':
            img = img.convert('RGB')
        
        # Convert to numpy array for analysis
        img_array = np.array(img)
        
        # Calculate image statistics
        analysis = {
            "dimensions": {"width": width, "height": height},
            "mode": mode,
            "statistics": {
                "mean_brightness": float(np.mean(img_array)),
                "std_brightness": float(np.std(img_array)),
                "min_brightness": float(np.min(img_array)),
                "max_brightness": float(np.max(img_array))
            },
            "color_analysis": self._analyze_colors(img_array),
            "detected_objects": self._detect_objects(img_array),
            "scene_type": self._classify_scene(img_array),
            "quality_metrics": self._assess_quality(img_array)
        }
        
        return {
            "status": "✅ success",
            "analysis": analysis,
            "emoji": "👁️",
            "message": "👁️ Image analysis complete!"
        }
    
    def _analyze_colors(self, img_array: np.ndarray) -> Dict:
        """🎨 Analyze color distribution"""
        # Flatten image to 2D array (pixels x channels)
        pixels = img_array.reshape(-1, 3)
        
        # Calculate dominant colors
        # Simple color quantization
        r_mean, g_mean, b_mean = np.mean(pixels, axis=0)
        
        # Determine dominant color family
        if r_mean > g_mean and r_mean > b_mean:
            dominant = "🔴 Red/Warm"
        elif g_mean > r_mean and g_mean > b_mean:
            dominant = "🟢 Green/Nature"
        elif b_mean > r_mean and b_mean > g_mean:
            dominant = "🔵 Blue/Cool"
        else:
            dominant = "⚪ Neutral"
        
        return {
            "dominant_color": dominant,
            "rgb_mean": [float(r_mean), float(g_mean), float(b_mean)],
            "color_variance": float(np.std(pixels))
        }
    
    def _detect_objects(self, img_array: np.ndarray) -> List[str]:
        """🔍 Detect objects in image (simplified)"""
        objects = []
        
        # Simple heuristic-based detection
        # In production, use actual object detection model
        
        # Check for faces (skin color detection)
        skin_mask = (img_array[:,:,0] > 95) & (img_array[:,:,1] > 40) & (img_array[:,:,2] > 20)
        if np.sum(skin_mask) > 1000:
            objects.append("👤 Person/Face")
        
        # Check for text (high contrast regions)
        gray = np.mean(img_array, axis=2)
        if np.std(gray) > 50:
            objects.append("📝 Text/Documents")
        
        # Check for nature (green regions)
        green_mask = (img_array[:,:,1] > img_array[:,:,0]) & (img_array[:,:,1] > img_array[:,:,2])
        if np.sum(green_mask) > 10000:
            objects.append("🌿 Nature/Plants")
        
        # Check for sky (blue regions)
        blue_mask = (img_array[:,:,2] > img_array[:,:,0]) & (img_array[:,:,2] > img_array[:,:,1])
        if np.sum(blue_mask) > 20000:
            objects.append("☁️ Sky")
        
        return objects if objects else ["🖼️ General Image"]
    
    def _classify_scene(self, img_array: np.ndarray) -> str:
        """🏞️ Classify scene type"""
        # Calculate brightness and color statistics
        brightness = np.mean(img_array)
        color_variance = np.std(img_array)
        
        if brightness > 200:
            return "☀️ Bright/Daylight"
        elif brightness < 50:
            return "🌙 Dark/Night"
        elif color_variance > 80:
            return "🎨 Colorful/Vibrant"
        elif color_variance < 30:
            return "⚪ Monotone/Minimal"
        else:
            return "📸 Standard Photo"
    
    def _assess_quality(self, img_array: np.ndarray) -> Dict:
        """📊 Assess image quality"""
        # Calculate sharpness (edge detection)
        from PIL import ImageFilter
        img = Image.fromarray(img_array)
        edges = img.filter(ImageFilter.FIND_EDGES)
        edge_array = np.array(edges)
        sharpness = np.mean(edge_array)
        
        # Calculate noise (variance in uniform regions)
        noise = np.std(img_array)
        
        # Overall quality score
        quality_score = min(100, sharpness * 2 + (100 - noise * 0.5))
        
        return {
            "sharpness": float(sharpness),
            "noise_level": float(noise),
            "quality_score": float(quality_score),
            "rating": "🌟 Excellent" if quality_score > 80 else "✅ Good" if quality_score > 60 else "⚠️ Fair"
        }
    
    # 🔄 IMAGE PROCESSING
    
    def process_image(self, image_data: str, operation: str = "enhance") -> Dict:
        """🔧 Process image with various operations"""
        print(f"🔧 Processing image: {operation}")
        
        # Decode image
        try:
            if image_data.startswith('data:image'):
                image_data = image_data.split(',')[1]
            image_bytes = base64.b64decode(image_data)
            img = Image.open(io.BytesIO(image_bytes))
        except Exception as e:
            return {"error": f"❌ Cannot decode image: {e}"}
        
        # Apply operation
        if operation == "enhance":
            processed = self._enhance_image(img)
        elif operation == "resize":
            processed = img.resize((256, 256))
        elif operation == "grayscale":
            processed = img.convert('L').convert('RGB')
        elif operation == "blur":
            processed = img.filter(ImageFilter.BLUR)
        elif operation == "sharpen":
            processed = img.filter(ImageFilter.SHARPEN)
        elif operation == "edge_detect":
            processed = img.filter(ImageFilter.FIND_EDGES).convert('RGB')
        else:
            processed = img
        
        # Convert to base64
        buffer = io.BytesIO()
        processed.save(buffer, format='PNG')
        processed_base64 = base64.b64encode(buffer.getvalue()).decode()
        
        return {
            "status": "✅ success",
            "operation": operation,
            "processed_image_base64": processed_base64,
            "emoji": "🔧",
            "message": f"🔧 Image {operation} complete!"
        }
    
    def _enhance_image(self, img: Image.Image) -> Image.Image:
        """✨ Enhance image quality"""
        # Enhance contrast
        from PIL import ImageEnhance
        enhancer = ImageEnhance.Contrast(img)
        img = enhancer.enhance(1.5)
        
        # Enhance sharpness
        enhancer = ImageEnhance.Sharpness(img)
        img = enhancer.enhance(1.3)
        
        # Enhance color
        enhancer = ImageEnhance.Color(img)
        img = enhancer.enhance(1.2)
        
        return img

# Export class
__all__ = ['VisionAI']