Add JAX typography, xector primitives, deferred effect chains, and GPU streaming
All checks were successful
Build and Deploy / build-and-deploy (push) Successful in 1m28s
All checks were successful
Build and Deploy / build-and-deploy (push) Successful in 1m28s
- Add JAX text rendering with font atlas, styled text placement, and typography primitives - Add xector (element-wise/reduction) operations library and sexp effects - Add deferred effect chain fusion for JIT-compiled effect pipelines - Expand drawing primitives with font management, alignment, shadow, and outline - Add interpreter support for function-style define and require - Add GPU persistence mode and hardware decode support to streaming - Add new sexp effects: cell_pattern, halftone, mosaic, and derived definitions - Add path registry for asset resolution - Add integration, primitives, and xector tests Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -1,126 +1,680 @@
|
||||
"""
|
||||
Drawing Primitives Library
|
||||
|
||||
Draw shapes, text, and characters on images.
|
||||
Draw shapes, text, and characters on images with sophisticated text handling.
|
||||
|
||||
Text Features:
|
||||
- Font loading from files or system fonts
|
||||
- Text measurement and fitting
|
||||
- Alignment (left/center/right, top/middle/bottom)
|
||||
- Opacity for fade effects
|
||||
- Multi-line text support
|
||||
- Shadow and outline effects
|
||||
"""
|
||||
import numpy as np
|
||||
import cv2
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
import os
|
||||
import glob as glob_module
|
||||
from typing import Optional, Tuple, List, Union
|
||||
|
||||
|
||||
# Default font (will be loaded lazily)
|
||||
_default_font = None
|
||||
# =============================================================================
|
||||
# Font Management
|
||||
# =============================================================================
|
||||
|
||||
# Font cache: (path, size) -> font object
|
||||
_font_cache = {}
|
||||
|
||||
# Common system font directories
|
||||
FONT_DIRS = [
|
||||
"/usr/share/fonts",
|
||||
"/usr/local/share/fonts",
|
||||
"~/.fonts",
|
||||
"~/.local/share/fonts",
|
||||
"/System/Library/Fonts", # macOS
|
||||
"/Library/Fonts", # macOS
|
||||
"C:/Windows/Fonts", # Windows
|
||||
]
|
||||
|
||||
# Default fonts to try (in order of preference)
|
||||
DEFAULT_FONTS = [
|
||||
"DejaVuSans.ttf",
|
||||
"DejaVuSansMono.ttf",
|
||||
"Arial.ttf",
|
||||
"Helvetica.ttf",
|
||||
"FreeSans.ttf",
|
||||
"LiberationSans-Regular.ttf",
|
||||
]
|
||||
|
||||
|
||||
def _get_default_font(size=16):
|
||||
"""Get default font, creating if needed."""
|
||||
global _default_font
|
||||
if _default_font is None or _default_font.size != size:
|
||||
def _find_font_file(name: str) -> Optional[str]:
|
||||
"""Find a font file by name in system directories."""
|
||||
# If it's already a full path
|
||||
if os.path.isfile(name):
|
||||
return name
|
||||
|
||||
# Expand user paths
|
||||
expanded = os.path.expanduser(name)
|
||||
if os.path.isfile(expanded):
|
||||
return expanded
|
||||
|
||||
# Search in font directories
|
||||
for font_dir in FONT_DIRS:
|
||||
font_dir = os.path.expanduser(font_dir)
|
||||
if not os.path.isdir(font_dir):
|
||||
continue
|
||||
|
||||
# Direct match
|
||||
direct = os.path.join(font_dir, name)
|
||||
if os.path.isfile(direct):
|
||||
return direct
|
||||
|
||||
# Recursive search
|
||||
for root, dirs, files in os.walk(font_dir):
|
||||
for f in files:
|
||||
if f.lower() == name.lower():
|
||||
return os.path.join(root, f)
|
||||
# Also match without extension
|
||||
base = os.path.splitext(f)[0]
|
||||
if base.lower() == name.lower():
|
||||
return os.path.join(root, f)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _get_default_font(size: int = 24) -> ImageFont.FreeTypeFont:
|
||||
"""Get a default font at the given size."""
|
||||
for font_name in DEFAULT_FONTS:
|
||||
path = _find_font_file(font_name)
|
||||
if path:
|
||||
try:
|
||||
return ImageFont.truetype(path, size)
|
||||
except:
|
||||
continue
|
||||
|
||||
# Last resort: PIL default
|
||||
return ImageFont.load_default()
|
||||
|
||||
|
||||
def prim_make_font(name_or_path: str, size: int = 24) -> ImageFont.FreeTypeFont:
|
||||
"""
|
||||
Load a font by name or path.
|
||||
|
||||
(make-font "Arial" 32) ; system font by name
|
||||
(make-font "/path/to/font.ttf" 24) ; font file path
|
||||
(make-font "DejaVuSans" 48) ; searches common locations
|
||||
|
||||
Returns a font object for use with text primitives.
|
||||
"""
|
||||
size = int(size)
|
||||
|
||||
# Check cache
|
||||
cache_key = (name_or_path, size)
|
||||
if cache_key in _font_cache:
|
||||
return _font_cache[cache_key]
|
||||
|
||||
# Find the font file
|
||||
path = _find_font_file(name_or_path)
|
||||
if not path:
|
||||
raise FileNotFoundError(f"Font not found: {name_or_path}")
|
||||
|
||||
# Load and cache
|
||||
font = ImageFont.truetype(path, size)
|
||||
_font_cache[cache_key] = font
|
||||
return font
|
||||
|
||||
|
||||
def prim_list_fonts() -> List[str]:
|
||||
"""
|
||||
List available system fonts.
|
||||
|
||||
(list-fonts) ; -> ("Arial.ttf" "DejaVuSans.ttf" ...)
|
||||
|
||||
Returns list of font filenames found in system directories.
|
||||
"""
|
||||
fonts = set()
|
||||
|
||||
for font_dir in FONT_DIRS:
|
||||
font_dir = os.path.expanduser(font_dir)
|
||||
if not os.path.isdir(font_dir):
|
||||
continue
|
||||
|
||||
for root, dirs, files in os.walk(font_dir):
|
||||
for f in files:
|
||||
if f.lower().endswith(('.ttf', '.otf', '.ttc')):
|
||||
fonts.add(f)
|
||||
|
||||
return sorted(fonts)
|
||||
|
||||
|
||||
def prim_font_size(font: ImageFont.FreeTypeFont) -> int:
|
||||
"""
|
||||
Get the size of a font.
|
||||
|
||||
(font-size my-font) ; -> 24
|
||||
"""
|
||||
return font.size
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Text Measurement
|
||||
# =============================================================================
|
||||
|
||||
def prim_text_size(text: str, font=None, font_size: int = 24) -> Tuple[int, int]:
|
||||
"""
|
||||
Measure text dimensions.
|
||||
|
||||
(text-size "Hello" my-font) ; -> (width height)
|
||||
(text-size "Hello" :font-size 32) ; -> (width height) with default font
|
||||
|
||||
For multi-line text, returns total bounding box.
|
||||
"""
|
||||
if font is None:
|
||||
font = _get_default_font(int(font_size))
|
||||
elif isinstance(font, (int, float)):
|
||||
font = _get_default_font(int(font))
|
||||
|
||||
# Create temporary image for measurement
|
||||
img = Image.new('RGB', (1, 1))
|
||||
draw = ImageDraw.Draw(img)
|
||||
|
||||
bbox = draw.textbbox((0, 0), str(text), font=font)
|
||||
width = bbox[2] - bbox[0]
|
||||
height = bbox[3] - bbox[1]
|
||||
|
||||
return (width, height)
|
||||
|
||||
|
||||
def prim_text_metrics(font=None, font_size: int = 24) -> dict:
|
||||
"""
|
||||
Get font metrics.
|
||||
|
||||
(text-metrics my-font) ; -> {ascent: 20, descent: 5, height: 25}
|
||||
|
||||
Useful for precise text layout.
|
||||
"""
|
||||
if font is None:
|
||||
font = _get_default_font(int(font_size))
|
||||
elif isinstance(font, (int, float)):
|
||||
font = _get_default_font(int(font))
|
||||
|
||||
ascent, descent = font.getmetrics()
|
||||
return {
|
||||
'ascent': ascent,
|
||||
'descent': descent,
|
||||
'height': ascent + descent,
|
||||
'size': font.size,
|
||||
}
|
||||
|
||||
|
||||
def prim_fit_text_size(text: str, max_width: int, max_height: int,
|
||||
font_name: str = None, min_size: int = 8,
|
||||
max_size: int = 500) -> int:
|
||||
"""
|
||||
Calculate font size to fit text within bounds.
|
||||
|
||||
(fit-text-size "Hello World" 400 100) ; -> 48
|
||||
(fit-text-size "Title" 800 200 :font-name "Arial")
|
||||
|
||||
Returns the largest font size that fits within max_width x max_height.
|
||||
"""
|
||||
max_width = int(max_width)
|
||||
max_height = int(max_height)
|
||||
min_size = int(min_size)
|
||||
max_size = int(max_size)
|
||||
text = str(text)
|
||||
|
||||
# Binary search for optimal size
|
||||
best_size = min_size
|
||||
low, high = min_size, max_size
|
||||
|
||||
while low <= high:
|
||||
mid = (low + high) // 2
|
||||
|
||||
if font_name:
|
||||
try:
|
||||
font = prim_make_font(font_name, mid)
|
||||
except:
|
||||
font = _get_default_font(mid)
|
||||
else:
|
||||
font = _get_default_font(mid)
|
||||
|
||||
w, h = prim_text_size(text, font)
|
||||
|
||||
if w <= max_width and h <= max_height:
|
||||
best_size = mid
|
||||
low = mid + 1
|
||||
else:
|
||||
high = mid - 1
|
||||
|
||||
return best_size
|
||||
|
||||
|
||||
def prim_fit_font(text: str, max_width: int, max_height: int,
|
||||
font_name: str = None, min_size: int = 8,
|
||||
max_size: int = 500) -> ImageFont.FreeTypeFont:
|
||||
"""
|
||||
Create a font sized to fit text within bounds.
|
||||
|
||||
(fit-font "Hello World" 400 100) ; -> font object
|
||||
(fit-font "Title" 800 200 :font-name "Arial")
|
||||
|
||||
Returns a font object at the optimal size.
|
||||
"""
|
||||
size = prim_fit_text_size(text, max_width, max_height,
|
||||
font_name, min_size, max_size)
|
||||
|
||||
if font_name:
|
||||
try:
|
||||
_default_font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf", size)
|
||||
return prim_make_font(font_name, size)
|
||||
except:
|
||||
_default_font = ImageFont.load_default()
|
||||
return _default_font
|
||||
pass
|
||||
|
||||
return _get_default_font(size)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Text Drawing
|
||||
# =============================================================================
|
||||
|
||||
def prim_text(img: np.ndarray, text: str,
|
||||
x: int = None, y: int = None,
|
||||
width: int = None, height: int = None,
|
||||
font=None, font_size: int = 24, font_name: str = None,
|
||||
color=None, opacity: float = 1.0,
|
||||
align: str = "left", valign: str = "top",
|
||||
fit: bool = False,
|
||||
shadow: bool = False, shadow_color=None, shadow_offset: int = 2,
|
||||
outline: bool = False, outline_color=None, outline_width: int = 1,
|
||||
line_spacing: float = 1.2) -> np.ndarray:
|
||||
"""
|
||||
Draw text with alignment, opacity, and effects.
|
||||
|
||||
Basic usage:
|
||||
(text frame "Hello" :x 100 :y 50)
|
||||
|
||||
Centered in frame:
|
||||
(text frame "Title" :align "center" :valign "middle")
|
||||
|
||||
Fit to box:
|
||||
(text frame "Big Text" :x 50 :y 50 :width 400 :height 100 :fit true)
|
||||
|
||||
With fade (for animations):
|
||||
(text frame "Fading" :x 100 :y 100 :opacity 0.5)
|
||||
|
||||
With effects:
|
||||
(text frame "Shadow" :x 100 :y 100 :shadow true)
|
||||
(text frame "Outline" :x 100 :y 100 :outline true :outline-color (0 0 0))
|
||||
|
||||
Args:
|
||||
img: Input frame
|
||||
text: Text to draw
|
||||
x, y: Position (if not specified, uses alignment in full frame)
|
||||
width, height: Bounding box (for fit and alignment within box)
|
||||
font: Font object from make-font
|
||||
font_size: Size if no font specified
|
||||
font_name: Font name to load
|
||||
color: RGB tuple (default white)
|
||||
opacity: 0.0 (invisible) to 1.0 (opaque) for fading
|
||||
align: "left", "center", "right"
|
||||
valign: "top", "middle", "bottom"
|
||||
fit: If true, auto-size font to fit in box
|
||||
shadow: Draw drop shadow
|
||||
shadow_color: Shadow color (default black)
|
||||
shadow_offset: Shadow offset in pixels
|
||||
outline: Draw text outline
|
||||
outline_color: Outline color (default black)
|
||||
outline_width: Outline thickness
|
||||
line_spacing: Multiplier for line height (for multi-line)
|
||||
|
||||
Returns:
|
||||
Frame with text drawn
|
||||
"""
|
||||
h, w = img.shape[:2]
|
||||
text = str(text)
|
||||
|
||||
# Default colors
|
||||
if color is None:
|
||||
color = (255, 255, 255)
|
||||
else:
|
||||
color = tuple(int(c) for c in color)
|
||||
|
||||
if shadow_color is None:
|
||||
shadow_color = (0, 0, 0)
|
||||
else:
|
||||
shadow_color = tuple(int(c) for c in shadow_color)
|
||||
|
||||
if outline_color is None:
|
||||
outline_color = (0, 0, 0)
|
||||
else:
|
||||
outline_color = tuple(int(c) for c in outline_color)
|
||||
|
||||
# Determine bounding box
|
||||
if x is None:
|
||||
x = 0
|
||||
if width is None:
|
||||
width = w
|
||||
if y is None:
|
||||
y = 0
|
||||
if height is None:
|
||||
height = h
|
||||
|
||||
x, y = int(x), int(y)
|
||||
box_width = int(width) if width else w - x
|
||||
box_height = int(height) if height else h - y
|
||||
|
||||
# Get or create font
|
||||
if font is None:
|
||||
if fit:
|
||||
font = prim_fit_font(text, box_width, box_height, font_name)
|
||||
elif font_name:
|
||||
try:
|
||||
font = prim_make_font(font_name, int(font_size))
|
||||
except:
|
||||
font = _get_default_font(int(font_size))
|
||||
else:
|
||||
font = _get_default_font(int(font_size))
|
||||
|
||||
# Measure text
|
||||
text_w, text_h = prim_text_size(text, font)
|
||||
|
||||
# Calculate position based on alignment
|
||||
if align == "center":
|
||||
draw_x = x + (box_width - text_w) // 2
|
||||
elif align == "right":
|
||||
draw_x = x + box_width - text_w
|
||||
else: # left
|
||||
draw_x = x
|
||||
|
||||
if valign == "middle":
|
||||
draw_y = y + (box_height - text_h) // 2
|
||||
elif valign == "bottom":
|
||||
draw_y = y + box_height - text_h
|
||||
else: # top
|
||||
draw_y = y
|
||||
|
||||
# Create RGBA image for compositing with opacity
|
||||
pil_img = Image.fromarray(img).convert('RGBA')
|
||||
|
||||
# Create text layer with transparency
|
||||
text_layer = Image.new('RGBA', (w, h), (0, 0, 0, 0))
|
||||
draw = ImageDraw.Draw(text_layer)
|
||||
|
||||
# Draw shadow first (if enabled)
|
||||
if shadow:
|
||||
shadow_x = draw_x + shadow_offset
|
||||
shadow_y = draw_y + shadow_offset
|
||||
shadow_rgba = shadow_color + (int(255 * opacity * 0.5),)
|
||||
draw.text((shadow_x, shadow_y), text, fill=shadow_rgba, font=font)
|
||||
|
||||
# Draw outline (if enabled)
|
||||
if outline:
|
||||
outline_rgba = outline_color + (int(255 * opacity),)
|
||||
ow = int(outline_width)
|
||||
for dx in range(-ow, ow + 1):
|
||||
for dy in range(-ow, ow + 1):
|
||||
if dx != 0 or dy != 0:
|
||||
draw.text((draw_x + dx, draw_y + dy), text,
|
||||
fill=outline_rgba, font=font)
|
||||
|
||||
# Draw main text
|
||||
text_rgba = color + (int(255 * opacity),)
|
||||
draw.text((draw_x, draw_y), text, fill=text_rgba, font=font)
|
||||
|
||||
# Composite
|
||||
result = Image.alpha_composite(pil_img, text_layer)
|
||||
return np.array(result.convert('RGB'))
|
||||
|
||||
|
||||
def prim_text_box(img: np.ndarray, text: str,
|
||||
x: int, y: int, width: int, height: int,
|
||||
font=None, font_size: int = 24, font_name: str = None,
|
||||
color=None, opacity: float = 1.0,
|
||||
align: str = "center", valign: str = "middle",
|
||||
fit: bool = True,
|
||||
padding: int = 0,
|
||||
background=None, background_opacity: float = 0.5,
|
||||
**kwargs) -> np.ndarray:
|
||||
"""
|
||||
Draw text fitted within a box, optionally with background.
|
||||
|
||||
(text-box frame "Title" 50 50 400 100)
|
||||
(text-box frame "Subtitle" 50 160 400 50
|
||||
:background (0 0 0) :background-opacity 0.7)
|
||||
|
||||
Convenience wrapper around text() for common box-with-text pattern.
|
||||
"""
|
||||
x, y = int(x), int(y)
|
||||
width, height = int(width), int(height)
|
||||
padding = int(padding)
|
||||
|
||||
result = img.copy()
|
||||
|
||||
# Draw background if specified
|
||||
if background is not None:
|
||||
bg_color = tuple(int(c) for c in background)
|
||||
|
||||
# Create background with opacity
|
||||
pil_img = Image.fromarray(result).convert('RGBA')
|
||||
bg_layer = Image.new('RGBA', (pil_img.width, pil_img.height), (0, 0, 0, 0))
|
||||
bg_draw = ImageDraw.Draw(bg_layer)
|
||||
bg_rgba = bg_color + (int(255 * background_opacity),)
|
||||
bg_draw.rectangle([x, y, x + width, y + height], fill=bg_rgba)
|
||||
result = np.array(Image.alpha_composite(pil_img, bg_layer).convert('RGB'))
|
||||
|
||||
# Draw text within padded box
|
||||
return prim_text(result, text,
|
||||
x=x + padding, y=y + padding,
|
||||
width=width - 2 * padding, height=height - 2 * padding,
|
||||
font=font, font_size=font_size, font_name=font_name,
|
||||
color=color, opacity=opacity,
|
||||
align=align, valign=valign, fit=fit,
|
||||
**kwargs)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Legacy text functions (keep for compatibility)
|
||||
# =============================================================================
|
||||
|
||||
def prim_draw_char(img, char, x, y, font_size=16, color=None):
|
||||
"""Draw a single character at (x, y)."""
|
||||
if color is None:
|
||||
color = [255, 255, 255]
|
||||
|
||||
pil_img = Image.fromarray(img)
|
||||
draw = ImageDraw.Draw(pil_img)
|
||||
font = _get_default_font(font_size)
|
||||
draw.text((x, y), char, fill=tuple(color), font=font)
|
||||
return np.array(pil_img)
|
||||
"""Draw a single character at (x, y). Legacy function."""
|
||||
return prim_text(img, str(char), x=int(x), y=int(y),
|
||||
font_size=int(font_size), color=color)
|
||||
|
||||
|
||||
def prim_draw_text(img, text, x, y, font_size=16, color=None):
|
||||
"""Draw text string at (x, y)."""
|
||||
"""Draw text string at (x, y). Legacy function."""
|
||||
return prim_text(img, str(text), x=int(x), y=int(y),
|
||||
font_size=int(font_size), color=color)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Shape Drawing
|
||||
# =============================================================================
|
||||
|
||||
def prim_fill_rect(img, x, y, w, h, color=None, opacity: float = 1.0):
|
||||
"""
|
||||
Fill a rectangle with color.
|
||||
|
||||
(fill-rect frame 10 10 100 50 (255 0 0))
|
||||
(fill-rect frame 10 10 100 50 (255 0 0) :opacity 0.5)
|
||||
"""
|
||||
if color is None:
|
||||
color = [255, 255, 255]
|
||||
|
||||
pil_img = Image.fromarray(img)
|
||||
draw = ImageDraw.Draw(pil_img)
|
||||
font = _get_default_font(font_size)
|
||||
draw.text((x, y), text, fill=tuple(color), font=font)
|
||||
return np.array(pil_img)
|
||||
|
||||
|
||||
def prim_fill_rect(img, x, y, w, h, color=None):
|
||||
"""Fill a rectangle with color."""
|
||||
if color is None:
|
||||
color = [255, 255, 255]
|
||||
|
||||
result = img.copy()
|
||||
x, y, w, h = int(x), int(y), int(w), int(h)
|
||||
result[y:y+h, x:x+w] = color
|
||||
return result
|
||||
|
||||
if opacity >= 1.0:
|
||||
result = img.copy()
|
||||
result[y:y+h, x:x+w] = color
|
||||
return result
|
||||
|
||||
# With opacity, use alpha compositing
|
||||
pil_img = Image.fromarray(img).convert('RGBA')
|
||||
layer = Image.new('RGBA', (pil_img.width, pil_img.height), (0, 0, 0, 0))
|
||||
draw = ImageDraw.Draw(layer)
|
||||
fill_rgba = tuple(int(c) for c in color) + (int(255 * opacity),)
|
||||
draw.rectangle([x, y, x + w, y + h], fill=fill_rgba)
|
||||
result = Image.alpha_composite(pil_img, layer)
|
||||
return np.array(result.convert('RGB'))
|
||||
|
||||
|
||||
def prim_draw_rect(img, x, y, w, h, color=None, thickness=1):
|
||||
def prim_draw_rect(img, x, y, w, h, color=None, thickness=1, opacity: float = 1.0):
|
||||
"""Draw rectangle outline."""
|
||||
if color is None:
|
||||
color = [255, 255, 255]
|
||||
|
||||
result = img.copy()
|
||||
cv2.rectangle(result, (int(x), int(y)), (int(x+w), int(y+h)),
|
||||
tuple(color), thickness)
|
||||
return result
|
||||
if opacity >= 1.0:
|
||||
result = img.copy()
|
||||
cv2.rectangle(result, (int(x), int(y)), (int(x+w), int(y+h)),
|
||||
tuple(int(c) for c in color), int(thickness))
|
||||
return result
|
||||
|
||||
# With opacity
|
||||
pil_img = Image.fromarray(img).convert('RGBA')
|
||||
layer = Image.new('RGBA', (pil_img.width, pil_img.height), (0, 0, 0, 0))
|
||||
draw = ImageDraw.Draw(layer)
|
||||
outline_rgba = tuple(int(c) for c in color) + (int(255 * opacity),)
|
||||
draw.rectangle([int(x), int(y), int(x+w), int(y+h)],
|
||||
outline=outline_rgba, width=int(thickness))
|
||||
result = Image.alpha_composite(pil_img, layer)
|
||||
return np.array(result.convert('RGB'))
|
||||
|
||||
|
||||
def prim_draw_line(img, x1, y1, x2, y2, color=None, thickness=1):
|
||||
def prim_draw_line(img, x1, y1, x2, y2, color=None, thickness=1, opacity: float = 1.0):
|
||||
"""Draw a line from (x1, y1) to (x2, y2)."""
|
||||
if color is None:
|
||||
color = [255, 255, 255]
|
||||
|
||||
result = img.copy()
|
||||
cv2.line(result, (int(x1), int(y1)), (int(x2), int(y2)),
|
||||
tuple(color), thickness)
|
||||
return result
|
||||
if opacity >= 1.0:
|
||||
result = img.copy()
|
||||
cv2.line(result, (int(x1), int(y1)), (int(x2), int(y2)),
|
||||
tuple(int(c) for c in color), int(thickness))
|
||||
return result
|
||||
|
||||
# With opacity
|
||||
pil_img = Image.fromarray(img).convert('RGBA')
|
||||
layer = Image.new('RGBA', (pil_img.width, pil_img.height), (0, 0, 0, 0))
|
||||
draw = ImageDraw.Draw(layer)
|
||||
line_rgba = tuple(int(c) for c in color) + (int(255 * opacity),)
|
||||
draw.line([(int(x1), int(y1)), (int(x2), int(y2))],
|
||||
fill=line_rgba, width=int(thickness))
|
||||
result = Image.alpha_composite(pil_img, layer)
|
||||
return np.array(result.convert('RGB'))
|
||||
|
||||
|
||||
def prim_draw_circle(img, cx, cy, radius, color=None, thickness=1, fill=False):
|
||||
def prim_draw_circle(img, cx, cy, radius, color=None, thickness=1,
|
||||
fill=False, opacity: float = 1.0):
|
||||
"""Draw a circle."""
|
||||
if color is None:
|
||||
color = [255, 255, 255]
|
||||
|
||||
result = img.copy()
|
||||
t = -1 if fill else thickness
|
||||
cv2.circle(result, (int(cx), int(cy)), int(radius), tuple(color), t)
|
||||
return result
|
||||
if opacity >= 1.0:
|
||||
result = img.copy()
|
||||
t = -1 if fill else int(thickness)
|
||||
cv2.circle(result, (int(cx), int(cy)), int(radius),
|
||||
tuple(int(c) for c in color), t)
|
||||
return result
|
||||
|
||||
# With opacity
|
||||
pil_img = Image.fromarray(img).convert('RGBA')
|
||||
layer = Image.new('RGBA', (pil_img.width, pil_img.height), (0, 0, 0, 0))
|
||||
draw = ImageDraw.Draw(layer)
|
||||
cx, cy, r = int(cx), int(cy), int(radius)
|
||||
rgba = tuple(int(c) for c in color) + (int(255 * opacity),)
|
||||
|
||||
if fill:
|
||||
draw.ellipse([cx - r, cy - r, cx + r, cy + r], fill=rgba)
|
||||
else:
|
||||
draw.ellipse([cx - r, cy - r, cx + r, cy + r],
|
||||
outline=rgba, width=int(thickness))
|
||||
|
||||
result = Image.alpha_composite(pil_img, layer)
|
||||
return np.array(result.convert('RGB'))
|
||||
|
||||
|
||||
def prim_draw_ellipse(img, cx, cy, rx, ry, angle=0, color=None, thickness=1, fill=False):
|
||||
def prim_draw_ellipse(img, cx, cy, rx, ry, angle=0, color=None,
|
||||
thickness=1, fill=False, opacity: float = 1.0):
|
||||
"""Draw an ellipse."""
|
||||
if color is None:
|
||||
color = [255, 255, 255]
|
||||
|
||||
result = img.copy()
|
||||
t = -1 if fill else thickness
|
||||
cv2.ellipse(result, (int(cx), int(cy)), (int(rx), int(ry)),
|
||||
angle, 0, 360, tuple(color), t)
|
||||
return result
|
||||
if opacity >= 1.0:
|
||||
result = img.copy()
|
||||
t = -1 if fill else int(thickness)
|
||||
cv2.ellipse(result, (int(cx), int(cy)), (int(rx), int(ry)),
|
||||
float(angle), 0, 360, tuple(int(c) for c in color), t)
|
||||
return result
|
||||
|
||||
# With opacity (note: PIL doesn't support rotated ellipses easily)
|
||||
# Fall back to cv2 on a separate layer
|
||||
layer = np.zeros((img.shape[0], img.shape[1], 4), dtype=np.uint8)
|
||||
t = -1 if fill else int(thickness)
|
||||
rgba = tuple(int(c) for c in color) + (int(255 * opacity),)
|
||||
cv2.ellipse(layer, (int(cx), int(cy)), (int(rx), int(ry)),
|
||||
float(angle), 0, 360, rgba, t)
|
||||
|
||||
pil_img = Image.fromarray(img).convert('RGBA')
|
||||
pil_layer = Image.fromarray(layer)
|
||||
result = Image.alpha_composite(pil_img, pil_layer)
|
||||
return np.array(result.convert('RGB'))
|
||||
|
||||
|
||||
def prim_draw_polygon(img, points, color=None, thickness=1, fill=False):
|
||||
def prim_draw_polygon(img, points, color=None, thickness=1,
|
||||
fill=False, opacity: float = 1.0):
|
||||
"""Draw a polygon from list of [x, y] points."""
|
||||
if color is None:
|
||||
color = [255, 255, 255]
|
||||
|
||||
result = img.copy()
|
||||
pts = np.array(points, dtype=np.int32).reshape((-1, 1, 2))
|
||||
if opacity >= 1.0:
|
||||
result = img.copy()
|
||||
pts = np.array(points, dtype=np.int32).reshape((-1, 1, 2))
|
||||
if fill:
|
||||
cv2.fillPoly(result, [pts], tuple(int(c) for c in color))
|
||||
else:
|
||||
cv2.polylines(result, [pts], True,
|
||||
tuple(int(c) for c in color), int(thickness))
|
||||
return result
|
||||
|
||||
# With opacity
|
||||
pil_img = Image.fromarray(img).convert('RGBA')
|
||||
layer = Image.new('RGBA', (pil_img.width, pil_img.height), (0, 0, 0, 0))
|
||||
draw = ImageDraw.Draw(layer)
|
||||
|
||||
pts_flat = [(int(p[0]), int(p[1])) for p in points]
|
||||
rgba = tuple(int(c) for c in color) + (int(255 * opacity),)
|
||||
|
||||
if fill:
|
||||
cv2.fillPoly(result, [pts], tuple(color))
|
||||
draw.polygon(pts_flat, fill=rgba)
|
||||
else:
|
||||
cv2.polylines(result, [pts], True, tuple(color), thickness)
|
||||
draw.polygon(pts_flat, outline=rgba, width=int(thickness))
|
||||
|
||||
return result
|
||||
result = Image.alpha_composite(pil_img, layer)
|
||||
return np.array(result.convert('RGB'))
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# PRIMITIVES Export
|
||||
# =============================================================================
|
||||
|
||||
PRIMITIVES = {
|
||||
# Text
|
||||
# Font management
|
||||
'make-font': prim_make_font,
|
||||
'list-fonts': prim_list_fonts,
|
||||
'font-size': prim_font_size,
|
||||
|
||||
# Text measurement
|
||||
'text-size': prim_text_size,
|
||||
'text-metrics': prim_text_metrics,
|
||||
'fit-text-size': prim_fit_text_size,
|
||||
'fit-font': prim_fit_font,
|
||||
|
||||
# Text drawing
|
||||
'text': prim_text,
|
||||
'text-box': prim_text_box,
|
||||
|
||||
# Legacy text (compatibility)
|
||||
'draw-char': prim_draw_char,
|
||||
'draw-text': prim_draw_text,
|
||||
|
||||
|
||||
@@ -8,12 +8,18 @@ GPU Acceleration:
|
||||
- Set STREAMING_GPU_PERSIST=1 to output CuPy arrays (frames stay on GPU)
|
||||
- Hardware video decoding (NVDEC) is used when available
|
||||
- Dramatically improves performance on GPU nodes
|
||||
|
||||
Async Prefetching:
|
||||
- Set STREAMING_PREFETCH=1 to enable background frame prefetching
|
||||
- Decodes upcoming frames while current frame is being processed
|
||||
"""
|
||||
|
||||
import os
|
||||
import numpy as np
|
||||
import subprocess
|
||||
import json
|
||||
import threading
|
||||
from collections import deque
|
||||
from pathlib import Path
|
||||
|
||||
# Try to import CuPy for GPU acceleration
|
||||
@@ -28,6 +34,10 @@ except ImportError:
|
||||
# Disabled by default until all primitives support GPU frames
|
||||
GPU_PERSIST = os.environ.get("STREAMING_GPU_PERSIST", "0") == "1" and CUPY_AVAILABLE
|
||||
|
||||
# Async prefetch mode - decode frames in background thread
|
||||
PREFETCH_ENABLED = os.environ.get("STREAMING_PREFETCH", "1") == "1"
|
||||
PREFETCH_BUFFER_SIZE = int(os.environ.get("STREAMING_PREFETCH_SIZE", "10"))
|
||||
|
||||
# Check for hardware decode support (cached)
|
||||
_HWDEC_AVAILABLE = None
|
||||
|
||||
@@ -283,6 +293,122 @@ class VideoSource:
|
||||
self._proc = None
|
||||
|
||||
|
||||
class PrefetchingVideoSource:
|
||||
"""
|
||||
Video source with background prefetching for improved performance.
|
||||
|
||||
Wraps VideoSource and adds a background thread that pre-decodes
|
||||
upcoming frames while the main thread processes the current frame.
|
||||
"""
|
||||
|
||||
def __init__(self, path: str, fps: float = 30, buffer_size: int = None):
|
||||
self._source = VideoSource(path, fps)
|
||||
self._buffer_size = buffer_size or PREFETCH_BUFFER_SIZE
|
||||
self._buffer = {} # time -> frame
|
||||
self._buffer_lock = threading.Lock()
|
||||
self._prefetch_time = 0.0
|
||||
self._frame_time = 1.0 / fps
|
||||
self._stop_event = threading.Event()
|
||||
self._request_event = threading.Event()
|
||||
self._target_time = 0.0
|
||||
|
||||
# Start prefetch thread
|
||||
self._thread = threading.Thread(target=self._prefetch_loop, daemon=True)
|
||||
self._thread.start()
|
||||
|
||||
import sys
|
||||
print(f"PrefetchingVideoSource: {path} buffer_size={self._buffer_size}", file=sys.stderr)
|
||||
|
||||
def _prefetch_loop(self):
|
||||
"""Background thread that pre-reads frames."""
|
||||
while not self._stop_event.is_set():
|
||||
# Wait for work or timeout
|
||||
self._request_event.wait(timeout=0.01)
|
||||
self._request_event.clear()
|
||||
|
||||
if self._stop_event.is_set():
|
||||
break
|
||||
|
||||
# Prefetch frames ahead of target time
|
||||
target = self._target_time
|
||||
with self._buffer_lock:
|
||||
# Clean old frames (more than 1 second behind)
|
||||
old_times = [t for t in self._buffer.keys() if t < target - 1.0]
|
||||
for t in old_times:
|
||||
del self._buffer[t]
|
||||
|
||||
# Count how many frames we have buffered ahead
|
||||
buffered_ahead = sum(1 for t in self._buffer.keys() if t >= target)
|
||||
|
||||
# Prefetch if buffer not full
|
||||
if buffered_ahead < self._buffer_size:
|
||||
# Find next time to prefetch
|
||||
prefetch_t = target
|
||||
with self._buffer_lock:
|
||||
existing_times = set(self._buffer.keys())
|
||||
for _ in range(self._buffer_size):
|
||||
if prefetch_t not in existing_times:
|
||||
break
|
||||
prefetch_t += self._frame_time
|
||||
|
||||
# Read the frame (this is the slow part)
|
||||
try:
|
||||
frame = self._source.read_at(prefetch_t)
|
||||
with self._buffer_lock:
|
||||
self._buffer[prefetch_t] = frame
|
||||
except Exception as e:
|
||||
import sys
|
||||
print(f"Prefetch error at t={prefetch_t}: {e}", file=sys.stderr)
|
||||
|
||||
def read_at(self, t: float) -> np.ndarray:
|
||||
"""Read frame at specific time, using prefetch buffer if available."""
|
||||
self._target_time = t
|
||||
self._request_event.set() # Wake up prefetch thread
|
||||
|
||||
# Round to frame time for buffer lookup
|
||||
t_key = round(t / self._frame_time) * self._frame_time
|
||||
|
||||
# Check buffer first
|
||||
with self._buffer_lock:
|
||||
if t_key in self._buffer:
|
||||
return self._buffer[t_key]
|
||||
# Also check for close matches (within half frame time)
|
||||
for buf_t, frame in self._buffer.items():
|
||||
if abs(buf_t - t) < self._frame_time * 0.5:
|
||||
return frame
|
||||
|
||||
# Not in buffer - read directly (blocking)
|
||||
frame = self._source.read_at(t)
|
||||
|
||||
# Store in buffer
|
||||
with self._buffer_lock:
|
||||
self._buffer[t_key] = frame
|
||||
|
||||
return frame
|
||||
|
||||
def read(self) -> np.ndarray:
|
||||
"""Read frame (uses last cached or t=0)."""
|
||||
return self.read_at(0)
|
||||
|
||||
def skip(self):
|
||||
"""No-op for seek-based reading."""
|
||||
pass
|
||||
|
||||
@property
|
||||
def size(self):
|
||||
return self._source.size
|
||||
|
||||
@property
|
||||
def path(self):
|
||||
return self._source.path
|
||||
|
||||
def close(self):
|
||||
self._stop_event.set()
|
||||
self._request_event.set() # Wake up thread to exit
|
||||
self._thread.join(timeout=1.0)
|
||||
self._source.close()
|
||||
|
||||
|
||||
class AudioAnalyzer:
|
||||
"""Audio analyzer for energy and beat detection."""
|
||||
|
||||
@@ -394,7 +520,12 @@ class AudioAnalyzer:
|
||||
# === Primitives ===
|
||||
|
||||
def prim_make_video_source(path: str, fps: float = 30):
|
||||
"""Create a video source from a file path."""
|
||||
"""Create a video source from a file path.
|
||||
|
||||
Uses PrefetchingVideoSource if STREAMING_PREFETCH=1 (default).
|
||||
"""
|
||||
if PREFETCH_ENABLED:
|
||||
return PrefetchingVideoSource(path, fps)
|
||||
return VideoSource(path, fps)
|
||||
|
||||
|
||||
|
||||
1382
sexp_effects/primitive_libs/xector.py
Normal file
1382
sexp_effects/primitive_libs/xector.py
Normal file
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user