animator.py

import streamlit as st
import os
import numpy as np
from PIL import Image, ImageEnhance, ImageFilter, ImageDraw
import time
from concurrent.futures import ThreadPoolExecutor
from functools import partial

class Animator:
    def __init__(self):
        self.frame_cache = {}
        self.aspect_ratio = "1:1"  # Default aspect ratio
        self.frames_per_animation = 15  # Default number of frames per animation for smoother transitions
        
    def set_aspect_ratio(self, aspect_ratio):
        """Set the aspect ratio for animations"""
        self.aspect_ratio = aspect_ratio
    
    def set_frames_per_animation(self, frames):
        """Set the number of frames per animation"""
        self.frames_per_animation = max(10, min(frames, 20))  # Keep between 10-20 frames for balance
    
    def apply_cinematic_effects(self, image):
        """Apply cinematic effects to enhance the frame quality"""
        try:
            # Convert to PIL Image if it's a path
            if isinstance(image, str):
                img = Image.open(image)
            else:
                img = image
                
            # Enhance contrast slightly
            enhancer = ImageEnhance.Contrast(img)
            img = enhancer.enhance(1.2)
            
            # Enhance color saturation slightly
            enhancer = ImageEnhance.Color(img)
            img = enhancer.enhance(1.1)
            
            # Add subtle vignette effect
            # Create a radial gradient mask
            mask = Image.new('L', img.size, 255)
            draw = ImageDraw.Draw(mask)
            
            width, height = img.size
            center_x, center_y = width // 2, height // 2
            max_radius = min(width, height) // 2
            
            for y in range(height):
                for x in range(width):
                    # Calculate distance from center
                    distance = np.sqrt((x - center_x)**2 + (y - center_y)**2)
                    # Create vignette effect (darker at edges)
                    intensity = int(255 * (1 - 0.3 * (distance / max_radius)**2))
                    mask.putpixel((x, y), intensity)
            
            # Apply the mask
            img = Image.composite(img, Image.new('RGB', img.size, (0, 0, 0)), mask)
            
            # Add subtle film grain
            grain = Image.effect_noise((img.width, img.height), 10)
            grain = grain.convert('L')
            grain = grain.filter(ImageFilter.GaussianBlur(radius=1))
            img = Image.blend(img, Image.composite(img, Image.new('RGB', img.size, (128, 128, 128)), grain), 0.05)
            
            return img
        except Exception as e:
            # If effects fail, return original image
            if isinstance(image, str):
                return Image.open(image)
            return image
    
    def add_zoom_animation(self, image_path, num_frames=None, zoom_factor=1.05, output_dir="temp"):
        """Add a simple zoom animation to an image with cinematic effects"""
        if num_frames is None:
            num_frames = self.frames_per_animation
            
        # Check cache first
        cache_key = f"zoom_{image_path}_{num_frames}_{zoom_factor}_{self.aspect_ratio}"
        if cache_key in self.frame_cache:
            return self.frame_cache[cache_key]
            
        # Ensure output directory exists
        os.makedirs(output_dir, exist_ok=True)
        
        # Load the image
        img = Image.open(image_path)
        
        # Create a sequence of slightly modified images for animation
        frames = []
        for scale in np.linspace(1.0, zoom_factor, num_frames):  # Subtle zoom
            size = (int(img.width * scale), int(img.height * scale))
            scaled_img = img.resize(size, Image.LANCZOS)
            
            # Center the scaled image
            new_img = Image.new("RGB", (img.width, img.height))
            left = (img.width - scaled_img.width) // 2
            top = (img.height - scaled_img.height) // 2
            new_img.paste(scaled_img, (left, top))
            
            # Apply cinematic effects
            new_img = self.apply_cinematic_effects(new_img)
            
            # Save the frame
            frame_path = f"{output_dir}/frame_{os.path.basename(image_path)}_{len(frames)}.png"
            new_img.save(frame_path)
            frames.append(frame_path)
        
        # Cache the result
        self.frame_cache[cache_key] = frames
        return frames
    
    def add_pan_animation(self, image_path, num_frames=None, direction="right", output_dir="temp"):
        """Add a simple panning animation to an image with cinematic effects"""
        if num_frames is None:
            num_frames = self.frames_per_animation
            
        # Check cache first
        cache_key = f"pan_{image_path}_{num_frames}_{direction}_{self.aspect_ratio}"
        if cache_key in self.frame_cache:
            return self.frame_cache[cache_key]
            
        # Ensure output directory exists
        os.makedirs(output_dir, exist_ok=True)
        
        # Load the image
        img = Image.open(image_path)
        
        # Create a sequence of panned images
        frames = []
        
        # Calculate pan parameters based on aspect ratio
        # For portrait (9:16), horizontal panning should be more subtle
        # For landscape (16:9), vertical panning should be more subtle
        pan_factor = 0.1  # Default pan factor
        
        if self.aspect_ratio == "9:16" and (direction == "left" or direction == "right"):
            pan_factor = 0.05  # Reduce horizontal pan for portrait
        elif self.aspect_ratio == "16:9" and (direction == "up" or direction == "down"):
            pan_factor = 0.05  # Reduce vertical pan for landscape
        
        # Calculate pan parameters
        if direction == "right":
            x_shifts = np.linspace(0, img.width * pan_factor, num_frames)
            y_shifts = np.zeros(num_frames)
        elif direction == "left":
            x_shifts = np.linspace(0, -img.width * pan_factor, num_frames)
            y_shifts = np.zeros(num_frames)
        elif direction == "down":
            x_shifts = np.zeros(num_frames)
            y_shifts = np.linspace(0, img.height * pan_factor, num_frames)
        elif direction == "up":
            x_shifts = np.zeros(num_frames)
            y_shifts = np.linspace(0, -img.height * pan_factor, num_frames)
        else:
            # Default to right
            x_shifts = np.linspace(0, img.width * pan_factor, num_frames)
            y_shifts = np.zeros(num_frames)
        
        for i in range(num_frames):
            # Create a new image with the same size
            new_img = Image.new("RGB", (img.width, img.height))
            
            # Paste the original image with shift
            new_img.paste(img, (int(x_shifts[i]), int(y_shifts[i])))
            
            # Apply cinematic effects
            new_img = self.apply_cinematic_effects(new_img)
            
            # Save the frame
            frame_path = f"{output_dir}/frame_{os.path.basename(image_path)}_{i}.png"
            new_img.save(frame_path)
            frames.append(frame_path)
        
        # Cache the result
        self.frame_cache[cache_key] = frames
        return frames
    
    def add_fade_animation(self, image_path, num_frames=None, fade_type="in", output_dir="temp"):
        """Add a fade in/out animation to an image with cinematic effects"""
        if num_frames is None:
            num_frames = self.frames_per_animation
            
        # Check cache first
        cache_key = f"fade_{image_path}_{num_frames}_{fade_type}_{self.aspect_ratio}"
        if cache_key in self.frame_cache:
            return self.frame_cache[cache_key]
            
        # Ensure output directory exists
        os.makedirs(output_dir, exist_ok=True)
        
        # Load the image
        img = Image.open(image_path)
        
        # Create a sequence of images with changing opacity
        frames = []
        
        if fade_type == "in":
            alphas = np.linspace(0.3, 1.0, num_frames)
        elif fade_type == "out":
            alphas = np.linspace(1.0, 0.3, num_frames)
        else:
            # Default to fade in
            alphas = np.linspace(0.3, 1.0, num_frames)
        
        for i, alpha in enumerate(alphas):
            # Create a new image with adjusted brightness
            enhancer = Image.new("RGBA", img.size, (0, 0, 0, 0))
            new_img = Image.blend(enhancer, img.convert("RGBA"), alpha)
            new_img = new_img.convert("RGB")
            
            # Apply cinematic effects
            new_img = self.apply_cinematic_effects(new_img)
            
            # Save the frame
            frame_path = f"{output_dir}/frame_{os.path.basename(image_path)}_{i}.png"
            new_img.save(frame_path)
            frames.append(frame_path)
        
        # Cache the result
        self.frame_cache[cache_key] = frames
        return frames
    
    def add_ken_burns_effect(self, image_path, num_frames=None, output_dir="temp"):
        """Add a Ken Burns effect (combination of pan and zoom) with cinematic effects"""
        if num_frames is None:
            num_frames = self.frames_per_animation
            
        # Check cache first
        cache_key = f"kenburns_{image_path}_{num_frames}_{self.aspect_ratio}"
        if cache_key in self.frame_cache:
            return self.frame_cache[cache_key]
            
        # Ensure output directory exists
        os.makedirs(output_dir, exist_ok=True)
        
        # Load the image
        img = Image.open(image_path)
        
        # Create a sequence of images with Ken Burns effect
        frames = []
        
        # Determine direction based on aspect ratio and image content
        import random
        if self.aspect_ratio == "16:9":
            # For landscape, prefer horizontal movement
            direction = random.choice(["right", "left"])
        elif self.aspect_ratio == "9:16":
            # For portrait, prefer vertical movement
            direction = random.choice(["up", "down"])
        else:
            # For square, random direction
            direction = random.choice(["right", "left", "up", "down"])
        
        # Calculate pan parameters
        if direction == "right":
            x_shifts = np.linspace(0, img.width * 0.05, num_frames)
            y_shifts = np.zeros(num_frames)
        elif direction == "left":
            x_shifts = np.linspace(0, -img.width * 0.05, num_frames)
            y_shifts = np.zeros(num_frames)
        elif direction == "down":
            x_shifts = np.zeros(num_frames)
            y_shifts = np.linspace(0, img.height * 0.05, num_frames)
        elif direction == "up":
            x_shifts = np.zeros(num_frames)
            y_shifts = np.linspace(0, -img.height * 0.05, num_frames)
        
        # Calculate zoom factors
        zoom_factors = np.linspace(1.0, 1.05, num_frames)
        
        for i in range(num_frames):
            # Apply zoom
            size = (int(img.width * zoom_factors[i]), int(img.height * zoom_factors[i]))
            zoomed_img = img.resize(size, Image.LANCZOS)
            
            # Create a new image with the same size as original
            new_img = Image.new("RGB", (img.width, img.height))
            
            # Calculate position with both zoom and pan
            left = (img.width - zoomed_img.width) // 2 + int(x_shifts[i])
            top = (img.height - zoomed_img.height) // 2 + int(y_shifts[i])
            
            # Paste the zoomed image with shift
            new_img.paste(zoomed_img, (left, top))
            
            # Apply cinematic effects
            new_img = self.apply_cinematic_effects(new_img)
            
            # Save the frame
            frame_path = f"{output_dir}/frame_{os.path.basename(image_path)}_{i}.png"
            new_img.save(frame_path)
            frames.append(frame_path)
        
        # Cache the result
        self.frame_cache[cache_key] = frames
        return frames
    
    def animate_single_image(self, img_path, animation_type="random", output_dir="temp", num_frames=None):
        """Animate a single image with cinematic effects"""
        if num_frames is None:
            num_frames = self.frames_per_animation
            
        # Choose animation type
        animation_types = ["zoom", "pan_right", "pan_left", "fade_in", "ken_burns"]
        
        # For different aspect ratios, prioritize certain animations
        if self.aspect_ratio == "16:9":
            # For landscape, prioritize horizontal panning
            animation_types = ["zoom", "pan_left", "pan_right", "ken_burns", "fade_in"]
        elif self.aspect_ratio == "9:16":
            # For portrait, prioritize vertical panning
            animation_types = ["zoom", "ken_burns", "fade_in", "pan_up", "pan_down"]
        
        if animation_type == "random":
            # Use hash of image path to deterministically select animation type
            import random
            random.seed(hash(img_path))
            chosen_type = random.choice(animation_types)
        else:
            chosen_type = animation_type
        
        # Apply the chosen animation
        if chosen_type == "ken_burns":
            frames = self.add_ken_burns_effect(img_path, num_frames=num_frames, output_dir=output_dir)
        elif chosen_type.startswith("pan"):
            direction = chosen_type.split("_")[1] if "_" in chosen_type else "right"
            frames = self.add_pan_animation(img_path, num_frames=num_frames, direction=direction, output_dir=output_dir)
        elif chosen_type.startswith("fade"):
            fade_type = chosen_type.split("_")[1] if "_" in chosen_type else "in"
            frames = self.add_fade_animation(img_path, num_frames=num_frames, fade_type=fade_type, output_dir=output_dir)
        else:  # Default to zoom
            frames = self.add_zoom_animation(img_path, num_frames=num_frames, output_dir=output_dir)
        
        return frames
    
    def animate_images(self, image_paths, animation_type="random", output_dir="temp", 
                      progress_callback=None, parallel=False, max_workers=4, batch_size=2, num_frames=None):
        """Add animations to a list of images with parallel processing and batching"""
        if num_frames is None:
            num_frames = self.frames_per_animation
            
        all_animated_frames = []
        
        if parallel and len(image_paths) > 1:
            # Process in parallel using ThreadPoolExecutor
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                # Create a partial function with fixed parameters
                animate_func = partial(self.animate_single_image, 
                                      animation_type=animation_type, 
                                      output_dir=output_dir,
                                      num_frames=num_frames)
                
                # Process images in parallel
                if progress_callback:
                    progress_callback("Animating images in parallel...")
                
                # Map and collect results
                all_animated_frames = list(executor.map(animate_func, image_paths))
        else:
            # Process in batches
            for i in range(0, len(image_paths), batch_size):
                batch = image_paths[i:i+batch_size]
                
                if progress_callback:
                    progress_callback(f"Animating batch {i//batch_size + 1}/{(len(image_paths) + batch_size - 1)//batch_size}...")
                
                batch_frames = []
                for img_path in batch:
                    frames = self.animate_single_image(img_path, animation_type, output_dir, num_frames)
                    batch_frames.append(frames)
                
                all_animated_frames.extend(batch_frames)
        
        return all_animated_frames
    
    def clear_cache(self):
        """Clear the animation frame cache"""
        self.frame_cache = {}
        return True