Create utils.py

utils.py

@@ -0,0 +1,158 @@
+import numpy as np
+import cv2
+import numexpr
+import re
+from PIL import Image, ImageOps
+
+# --- Math & Schedule Parsing ---
+
+def get_inbetweens(key_frames, max_frames, integer=False):
+    """Interpolates values between keyframes (simple linear for now, but robust)."""
+    key_frames = dict(sorted(key_frames.items()))
+    keys = list(key_frames.keys())
+    vals = list(key_frames.values())
+    
+    # Fill array
+    series = np.linspace(vals[0], vals[0], max_frames)
+    for i in range(len(keys)-1):
+        idx_start, idx_end = keys[i], keys[i+1]
+        val_start, val_end = vals[i], vals[i+1]
+        if idx_end > max_frames: idx_end = max_frames
+        
+        # Linear interpolation
+        range_len = idx_end - idx_start
+        if range_len > 0:
+            segment = np.linspace(val_start, val_end, range_len, endpoint=False)
+            series[idx_start:idx_end] = segment
+            
+    # Fill tail
+    if keys[-1] < max_frames:
+        series[keys[-1]:] = vals[-1]
+        
+    return series.astype(int) if integer else series
+
+def parse_weight_string(string, max_frames):
+    """
+    Parses '0:(0.5), 10:(1.0)' including Math like '0:(sin(t/10))'.
+    Returns a numpy array of float values for every frame.
+    """
+    # Clean string
+    string = re.sub(r'\s+', '', string)
+    keyframes = {}
+    
+    # Split by comma, respecting parentheses might be needed in complex regex, 
+    # but simple split usually works for Deforum format
+    parts = string.split(',')
+    
+    for part in parts:
+        try:
+            frame_str, val_str = part.split(':')
+            frame = int(frame_str)
+            # Remove parentheses
+            val_str = val_str.strip('()')
+            keyframes[frame] = val_str
+        except:
+            continue
+            
+    if 0 not in keyframes:
+        keyframes[0] = "0"
+        
+    sorted_frames = sorted(keyframes.keys())
+    series = np.zeros(max_frames)
+    
+    # Evaluate math for every frame
+    for i in range(len(sorted_frames)):
+        f_start = sorted_frames[i]
+        f_end = sorted_frames[i+1] if i < len(sorted_frames)-1 else max_frames
+        
+        formula = keyframes[f_start]
+        
+        for f in range(f_start, f_end):
+            t = f # Deforum standard variable
+            try:
+                # Safe evaluation environment
+                val = numexpr.evaluate(formula, local_dict={'t': t, 's': f_start, 'pi': np.pi, 'sin': np.sin, 'cos': np.cos, 'tan': np.tan})
+                series[f] = float(val)
+            except Exception as e:
+                # If static value or fail
+                try:
+                    series[f] = float(formula)
+                except:
+                    series[f] = series[f-1] if f > 0 else 0.0
+                    
+    return series
+
+# --- Image Processing ---
+
+def maintain_colors(prev_img, color_match_sample):
+    """
+    Matches the coloring of the previous frame (or frame 0) to prevent color drift.
+    Uses LAB color space transfer.
+    """
+    prev_img_np = np.array(prev_img).astype(np.uint8)
+    sample_np = np.array(color_match_sample).astype(np.uint8)
+    
+    prev_lab = cv2.cvtColor(prev_img_np, cv2.COLOR_RGB2LAB)
+    sample_lab = cv2.cvtColor(sample_np, cv2.COLOR_RGB2LAB)
+    
+    l_avg_p, a_avg_p, b_avg_p = np.mean(prev_lab[:,:,0]), np.mean(prev_lab[:,:,1]), np.mean(prev_lab[:,:,2])
+    l_avg_s, a_avg_s, b_avg_s = np.mean(sample_lab[:,:,0]), np.mean(sample_lab[:,:,1]), np.mean(sample_lab[:,:,2])
+    
+    l, a, b = cv2.split(prev_lab)
+    
+    # Shift current image logic towards sample mean
+    # Note: Deforum usually matches the NEW generation to the OLD image.
+    # Here we are adjusting the image we just warped (prev) to match the original anchor? 
+    # Actually, standard Deforum 'Match Frame 0' means we force the init image to look like Frame 0 colors.
+    
+    l = l - l_avg_p + l_avg_s
+    a = a - a_avg_p + a_avg_s
+    b = b - b_avg_p + b_avg_s
+    
+    l = np.clip(l, 0, 255)
+    a = np.clip(a, 0, 255)
+    b = np.clip(b, 0, 255)
+    
+    matched_lab = cv2.merge([l.astype(np.uint8), a.astype(np.uint8), b.astype(np.uint8)])
+    return Image.fromarray(cv2.cvtColor(matched_lab, cv2.COLOR_LAB2RGB))
+
+def add_noise(img, noise_amt):
+    """Adds uniform noise to the image to give the diffusion model texture to latch onto."""
+    if noise_amt <= 0: return img
+    
+    img_np = np.array(img).astype(np.float32)
+    noise = np.random.normal(0, noise_amt * 255, img_np.shape).astype(np.float32)
+    noisy_img = np.clip(img_np + noise, 0, 255).astype(np.uint8)
+    return Image.fromarray(noisy_img)
+
+def anim_frame_warp_2d(prev_img_pil, args_dict):
+    """
+    Standard Deforum 2D Warping.
+    args_dict must contain: angle, zoom, translation_x, translation_y
+    """
+    cv2_img = np.array(prev_img_pil)
+    height, width = cv2_img.shape[:2]
+    
+    center = (width // 2, height // 2)
+    
+    # Rotation & Zoom
+    angle = args_dict.get('angle', 0)
+    zoom = args_dict.get('zoom', 1.0)
+    
+    trans_mat = cv2.getRotationMatrix2D(center, angle, zoom)
+    
+    # Translation
+    tx = args_dict.get('translation_x', 0)
+    ty = args_dict.get('translation_y', 0)
+    trans_mat[0, 2] += tx
+    trans_mat[1, 2] += ty
+    
+    # Warp with Reflection to handle edges naturally
+    warped = cv2.warpAffine(
+        cv2_img, 
+        trans_mat, 
+        (width, height), 
+        borderMode=cv2.BORDER_REFLECT_101
+    )
+    
+    return Image.fromarray(warped)