mirror of
https://github.com/3b1b/manim.git
synced 2025-07-28 12:32:36 +08:00
Begin setting up Camera to work with shaders, not yet done
This commit is contained in:
Grant Sanderson
@ -1,41 +1,52 @@
|
||||
from functools import reduce
|
||||
import itertools as it
|
||||
import operator as op
|
||||
import time
|
||||
import copy
|
||||
import moderngl
|
||||
import re
|
||||
from colour import Color
|
||||
|
||||
from PIL import Image
|
||||
from scipy.spatial.distance import pdist
|
||||
import cairo
|
||||
import numpy as np
|
||||
import itertools as it
|
||||
|
||||
from manimlib.constants import *
|
||||
from manimlib.mobject.types.image_mobject import AbstractImageMobject
|
||||
from manimlib.mobject.mobject import Mobject
|
||||
from manimlib.mobject.types.point_cloud_mobject import PMobject
|
||||
from manimlib.mobject.types.vectorized_mobject import VMobject
|
||||
from manimlib.utils.color import color_to_int_rgba
|
||||
from manimlib.utils.config_ops import digest_config
|
||||
from manimlib.utils.images import get_full_raster_image_path
|
||||
from manimlib.utils.iterables import batch_by_property
|
||||
from manimlib.utils.iterables import list_difference_update
|
||||
from manimlib.utils.iterables import join_structured_arrays
|
||||
from manimlib.utils.family_ops import extract_mobject_family_members
|
||||
from manimlib.utils.simple_functions import fdiv
|
||||
from manimlib.utils.space_ops import angle_of_vector
|
||||
from manimlib.utils.space_ops import get_norm
|
||||
|
||||
|
||||
# TODO, think about how to incorporate perspective,
|
||||
# and change get_height, etc. to take orientation into account
|
||||
class CameraFrame(Mobject):
|
||||
CONFIG = {
|
||||
"width": FRAME_WIDTH,
|
||||
"height": FRAME_HEIGHT,
|
||||
"center": ORIGIN,
|
||||
}
|
||||
|
||||
def generate_points(self):
|
||||
self.points = np.array([UL, UR, DR, DL])
|
||||
self.set_width(self.width, stretch=True)
|
||||
self.set_height(self.height, stretch=True)
|
||||
self.move_to(self.center)
|
||||
|
||||
|
||||
class Camera(object):
|
||||
CONFIG = {
|
||||
"background_image": None,
|
||||
"frame_config": {
|
||||
"width": FRAME_WIDTH,
|
||||
"height": FRAME_HEIGHT,
|
||||
"center": ORIGIN,
|
||||
},
|
||||
"pixel_height": DEFAULT_PIXEL_HEIGHT,
|
||||
"pixel_width": DEFAULT_PIXEL_WIDTH,
|
||||
"frame_rate": DEFAULT_FRAME_RATE,
|
||||
# Note: frame height and width will be resized to match
|
||||
# the pixel aspect ratio
|
||||
"frame_height": FRAME_HEIGHT,
|
||||
"frame_width": FRAME_WIDTH,
|
||||
"frame_center": ORIGIN,
|
||||
"background_color": BLACK,
|
||||
"background_opacity": 1,
|
||||
# Points in vectorized mobjects with norm greater
|
||||
@ -44,57 +55,40 @@ class Camera(object):
|
||||
"image_mode": "RGBA",
|
||||
"n_channels": 4,
|
||||
"pixel_array_dtype": 'uint8',
|
||||
# z_buff_func is only used if the flag above is set to True.
|
||||
# round z coordinate to nearest hundredth when comparring
|
||||
"z_buff_func": lambda m: np.round(m.get_center()[2], 2),
|
||||
"cairo_line_width_multiple": 0.01,
|
||||
"line_width_multiple": 0.01,
|
||||
"background_fbo": None,
|
||||
}
|
||||
|
||||
def __init__(self, background=None, **kwargs):
|
||||
digest_config(self, kwargs, locals())
|
||||
self.rgb_max_val = np.iinfo(self.pixel_array_dtype).max
|
||||
self.pixel_array_to_cairo_context = {}
|
||||
self.init_background()
|
||||
self.resize_frame_shape()
|
||||
self.reset()
|
||||
self.init_frame()
|
||||
self.init_context()
|
||||
self.init_frame_buffer()
|
||||
self.init_shaders()
|
||||
|
||||
def __deepcopy__(self, memo):
|
||||
# This is to address a strange bug where deepcopying
|
||||
# will result in a segfault, which is somehow related
|
||||
# to the aggdraw library
|
||||
self.canvas = None
|
||||
return copy.copy(self)
|
||||
def init_frame(self):
|
||||
self.frame = CameraFrame(**self.frame_config)
|
||||
|
||||
def reset_pixel_shape(self, new_height, new_width):
|
||||
self.pixel_width = new_width
|
||||
self.pixel_height = new_height
|
||||
self.init_background()
|
||||
self.resize_frame_shape()
|
||||
self.reset()
|
||||
def init_context(self):
|
||||
# TODO, context with a window?
|
||||
ctx = moderngl.create_standalone_context()
|
||||
ctx.enable(moderngl.BLEND)
|
||||
ctx.blend_func = (
|
||||
moderngl.SRC_ALPHA, moderngl.ONE_MINUS_SRC_ALPHA,
|
||||
moderngl.ONE, moderngl.ONE
|
||||
)
|
||||
self.ctx = ctx
|
||||
|
||||
def get_pixel_height(self):
|
||||
return self.pixel_height
|
||||
# Methods associated with the frame buffer
|
||||
def init_frame_buffer(self):
|
||||
# TODO, account for live window
|
||||
self.fbo = self.get_fbo()
|
||||
self.fbo.use()
|
||||
self.clear()
|
||||
|
||||
def get_pixel_width(self):
|
||||
return self.pixel_width
|
||||
|
||||
def get_frame_height(self):
|
||||
return self.frame_height
|
||||
|
||||
def get_frame_width(self):
|
||||
return self.frame_width
|
||||
|
||||
def get_frame_center(self):
|
||||
return self.frame_center
|
||||
|
||||
def set_frame_height(self, frame_height):
|
||||
self.frame_height = frame_height
|
||||
|
||||
def set_frame_width(self, frame_width):
|
||||
self.frame_width = frame_width
|
||||
|
||||
def set_frame_center(self, frame_center):
|
||||
self.frame_center = frame_center
|
||||
def get_fbo(self):
|
||||
return self.ctx.simple_framebuffer(self.get_pixel_shape())
|
||||
|
||||
def resize_frame_shape(self, fixed_dimension=0):
|
||||
"""
|
||||
@ -115,98 +109,84 @@ class Camera(object):
|
||||
self.set_frame_height(frame_height)
|
||||
self.set_frame_width(frame_width)
|
||||
|
||||
def init_background(self):
|
||||
height = self.get_pixel_height()
|
||||
width = self.get_pixel_width()
|
||||
if self.background_image is not None:
|
||||
path = get_full_raster_image_path(self.background_image)
|
||||
image = Image.open(path).convert(self.image_mode)
|
||||
# TODO, how to gracefully handle backgrounds
|
||||
# with different sizes?
|
||||
self.background = np.array(image)[:height, :width]
|
||||
self.background = self.background.astype(self.pixel_array_dtype)
|
||||
def clear(self):
|
||||
if self.background_fbo:
|
||||
self.ctx.copy_framebuffer(self.fbo, self.background_fbo)
|
||||
else:
|
||||
background_rgba = color_to_int_rgba(
|
||||
self.background_color, self.background_opacity
|
||||
)
|
||||
self.background = np.zeros(
|
||||
(height, width, self.n_channels),
|
||||
dtype=self.pixel_array_dtype
|
||||
)
|
||||
self.background[:, :] = background_rgba
|
||||
rgba = (*Color(self.background_color).get_rgb(), self.background_opacity)
|
||||
self.fbo.clear(*rgba)
|
||||
|
||||
def lock_state_as_background(self):
|
||||
self.background_fbo = self.get_fbo()
|
||||
self.ctx.copy_framebuffer(self.background_fbo, self.fbo)
|
||||
|
||||
def unlock_background(self):
|
||||
self.background_fbo = None
|
||||
|
||||
def reset_pixel_shape(self, new_height, new_width):
|
||||
self.pixel_width = new_width
|
||||
self.pixel_height = new_height
|
||||
self.fbo.release()
|
||||
self.init_frame_buffer()
|
||||
|
||||
# Various ways to read from the fbo
|
||||
def get_raw_fbo_data(self, dtype='f1'):
|
||||
return self.fbo.read(components=self.n_channels, dtype=dtype)
|
||||
|
||||
def get_image(self, pixel_array=None):
|
||||
if pixel_array is None:
|
||||
pixel_array = self.pixel_array
|
||||
return Image.fromarray(
|
||||
pixel_array,
|
||||
mode=self.image_mode
|
||||
return Image.frombytes(
|
||||
'RGBA', self.fbo.size,
|
||||
self.get_raw_fbo_data(),
|
||||
'raw', 'RGBA', 0, -1
|
||||
)
|
||||
|
||||
def get_pixel_array(self):
|
||||
return self.pixel_array
|
||||
raw = self.get_raw_fbo_data(dtype='f4')
|
||||
flat_arr = np.frombuffer(raw, dtype='f4')
|
||||
arr = flat_arr.reshape([*self.fbo.size, self.n_channels])
|
||||
# Convert from float
|
||||
return (self.rgb_max_val * arr).astype(self.pixel_array_dtype)
|
||||
|
||||
def convert_pixel_array(self, pixel_array, convert_from_floats=False):
|
||||
retval = np.array(pixel_array)
|
||||
if convert_from_floats:
|
||||
retval = np.apply_along_axis(
|
||||
lambda f: (f * self.rgb_max_val).astype(self.pixel_array_dtype),
|
||||
2,
|
||||
retval
|
||||
)
|
||||
return retval
|
||||
|
||||
def set_pixel_array(self, pixel_array, convert_from_floats=False):
|
||||
converted_array = self.convert_pixel_array(
|
||||
pixel_array, convert_from_floats)
|
||||
if not (hasattr(self, "pixel_array") and self.pixel_array.shape == converted_array.shape):
|
||||
self.pixel_array = converted_array
|
||||
else:
|
||||
# Set in place
|
||||
self.pixel_array[:, :, :] = converted_array[:, :, :]
|
||||
|
||||
def set_background(self, pixel_array, convert_from_floats=False):
|
||||
self.background = self.convert_pixel_array(
|
||||
pixel_array, convert_from_floats)
|
||||
|
||||
# TODO, this should live in utils, not as a method of Camera
|
||||
# Also, this should be implement with a shader
|
||||
def make_background_from_func(self, coords_to_colors_func):
|
||||
"""
|
||||
Sets background by using coords_to_colors_func to determine each pixel's color. Each input
|
||||
to coords_to_colors_func is an (x, y) pair in space (in ordinary space coordinates; not
|
||||
pixel coordinates), and each output is expected to be an RGBA array of 4 floats.
|
||||
"""
|
||||
|
||||
print("Starting set_background; for reference, the current time is ", time.strftime("%H:%M:%S"))
|
||||
coords = self.get_coords_of_all_pixels()
|
||||
new_background = np.apply_along_axis(
|
||||
coords_to_colors_func,
|
||||
2,
|
||||
coords
|
||||
# Needed?
|
||||
def get_texture(self):
|
||||
texture = self.ctx.texture(
|
||||
size=self.fbo.size,
|
||||
components=4,
|
||||
data=self.get_raw_fbo_data(),
|
||||
dtype='f4'
|
||||
)
|
||||
print("Ending set_background; for reference, the current time is ", time.strftime("%H:%M:%S"))
|
||||
return texture
|
||||
|
||||
return self.convert_pixel_array(new_background, convert_from_floats=True)
|
||||
# Getting camera attributes
|
||||
def get_pixel_shape(self):
|
||||
return (self.pixel_width, self.pixel_height)
|
||||
|
||||
def set_background_from_func(self, coords_to_colors_func):
|
||||
self.set_background(
|
||||
self.make_background_from_func(coords_to_colors_func))
|
||||
def get_pixel_width(self):
|
||||
return self.get_pixel_shape()[0]
|
||||
|
||||
def reset(self):
|
||||
self.set_pixel_array(self.background)
|
||||
return self
|
||||
def get_pixel_height(self):
|
||||
return self.get_pixel_shape()[1]
|
||||
|
||||
###
|
||||
def get_mobjects_to_display(self, mobjects, excluded_mobjects=None):
|
||||
mobjects = extract_mobject_family_members(
|
||||
mobjects, only_those_with_points=True,
|
||||
)
|
||||
if excluded_mobjects:
|
||||
all_excluded = extract_mobject_family_members(excluded_mobjects)
|
||||
mobjects = list_difference_update(mobjects, all_excluded)
|
||||
return mobjects
|
||||
# TODO, make these work for a rotated frame
|
||||
def get_frame_height(self):
|
||||
return self.frame.get_height()
|
||||
|
||||
def get_frame_width(self):
|
||||
return self.frame.get_width()
|
||||
|
||||
def get_frame_center(self):
|
||||
return self.frame.get_center()
|
||||
|
||||
def set_frame_height(self, height):
|
||||
self.frame.set_height(height, stretch=True)
|
||||
|
||||
def set_frame_width(self, width):
|
||||
self.frame.set_width(width, stretch=True)
|
||||
|
||||
def set_frame_center(self, center):
|
||||
self.frame.move_to(center)
|
||||
|
||||
# TODO, account for 3d
|
||||
def is_in_frame(self, mobject):
|
||||
fc = self.get_frame_center()
|
||||
fh = self.get_frame_height()
|
||||
@ -218,476 +198,125 @@ class Camera(object):
|
||||
mobject.get_top()[1] < fc[1] - fh,
|
||||
])
|
||||
|
||||
# Rendering
|
||||
def get_mobjects_to_display(self, mobjects, excluded_mobjects=None):
|
||||
mobjects = extract_mobject_family_members(
|
||||
mobjects, only_those_with_points=True,
|
||||
)
|
||||
if excluded_mobjects:
|
||||
all_excluded = extract_mobject_family_members(excluded_mobjects)
|
||||
mobjects = list_difference_update(mobjects, all_excluded)
|
||||
return mobjects
|
||||
|
||||
def capture_mobject(self, mobject, **kwargs):
|
||||
return self.capture_mobjects([mobject], **kwargs)
|
||||
|
||||
def capture_mobjects(self, mobjects, **kwargs):
|
||||
mobjects = self.get_mobjects_to_display(mobjects, **kwargs)
|
||||
shader_infos = list(it.chain(*[
|
||||
mob.get_shader_info_list()
|
||||
for mob in mobjects
|
||||
]))
|
||||
# TODO, batching works well when the mobjects are already organized,
|
||||
# but can we somehow use z-buffering to better effect here?
|
||||
batches = batch_by_property(shader_infos, self.get_shader_id)
|
||||
for info_group, sid in batches:
|
||||
data = join_structured_arrays(*[info["data"] for info in info_group])
|
||||
shader = self.get_shader(sid)
|
||||
self.render_from_shader(shader, data)
|
||||
|
||||
# Organize this list into batches of the same type, and
|
||||
# apply corresponding function to those batches
|
||||
type_func_pairs = [
|
||||
(VMobject, self.display_multiple_vectorized_mobjects),
|
||||
(PMobject, self.display_multiple_point_cloud_mobjects),
|
||||
(AbstractImageMobject, self.display_multiple_image_mobjects),
|
||||
(Mobject, lambda batch, pa: batch), # Do nothing
|
||||
]
|
||||
# Shader stuff
|
||||
def init_shaders(self):
|
||||
self.id_to_shader = {}
|
||||
|
||||
def get_mobject_type(mobject):
|
||||
for mobject_type, func in type_func_pairs:
|
||||
if isinstance(mobject, mobject_type):
|
||||
return mobject_type
|
||||
raise Exception(
|
||||
"Trying to display something which is not of type Mobject"
|
||||
def get_shader_id(self, shader_info):
|
||||
# A unique id for a shader based on the names of the files holding its code
|
||||
return "|".join([
|
||||
shader_info.get(key, "")
|
||||
for key in ["vert", "geom", "frag"]
|
||||
])
|
||||
|
||||
def get_shader(self, sid):
|
||||
if sid not in self.id_to_shader:
|
||||
vert, geom, frag = sid.split("|")
|
||||
shader = self.ctx.program(
|
||||
vertex_shader=self.get_shader_code_from_file(vert),
|
||||
geometry_shader=self.get_shader_code_from_file(geom),
|
||||
fragment_shader=self.get_shader_code_from_file(frag),
|
||||
)
|
||||
batch_type_pairs = batch_by_property(mobjects, get_mobject_type)
|
||||
self.set_shader_uniforms(shader)
|
||||
self.id_to_shader[sid] = shader
|
||||
return self.id_to_shader[sid]
|
||||
|
||||
# Display in these batches
|
||||
for batch, batch_type in batch_type_pairs:
|
||||
# check what the type is, and call the appropriate function
|
||||
for mobject_type, func in type_func_pairs:
|
||||
if batch_type == mobject_type:
|
||||
func(batch, self.pixel_array)
|
||||
def get_shader_code_from_file(self, filename):
|
||||
if len(filename) == 0:
|
||||
return None
|
||||
|
||||
# Methods associated with svg rendering
|
||||
filepath = os.path.join(SHADER_DIR, filename)
|
||||
if not os.path.exists(filepath):
|
||||
warnings.warn(f"No file at {file_path}")
|
||||
return
|
||||
|
||||
def get_cached_cairo_context(self, pixel_array):
|
||||
return self.pixel_array_to_cairo_context.get(
|
||||
id(pixel_array), None
|
||||
)
|
||||
with open(filepath, "r") as f:
|
||||
result = f.read()
|
||||
|
||||
def cache_cairo_context(self, pixel_array, ctx):
|
||||
self.pixel_array_to_cairo_context[id(pixel_array)] = ctx
|
||||
# To share functionality between shaders, some functions are read in
|
||||
# from other files an inserted into the relevant strings before
|
||||
# passing to ctx.program for compiling
|
||||
# Replace "#INSERT " lines with relevant code
|
||||
insertions = re.findall(r"^#INSERT .*\.glsl$", result, flags=re.MULTILINE)
|
||||
for line in insertions:
|
||||
inserted_code = self.get_shader_code_from_file(line.replace("#INSERT ", ""))
|
||||
result = result.replace(line, inserted_code)
|
||||
return result
|
||||
|
||||
def get_cairo_context(self, pixel_array):
|
||||
cached_ctx = self.get_cached_cairo_context(pixel_array)
|
||||
if cached_ctx:
|
||||
return cached_ctx
|
||||
pw = self.get_pixel_width()
|
||||
ph = self.get_pixel_height()
|
||||
def set_shader_uniforms(self, shader):
|
||||
# TODO, think about how uniforms come from mobjects
|
||||
# as well.
|
||||
fw = self.get_frame_width()
|
||||
fh = self.get_frame_height()
|
||||
fc = self.get_frame_center()
|
||||
surface = cairo.ImageSurface.create_for_data(
|
||||
pixel_array,
|
||||
cairo.FORMAT_ARGB32,
|
||||
pw, ph
|
||||
)
|
||||
ctx = cairo.Context(surface)
|
||||
ctx.scale(pw, ph)
|
||||
ctx.set_matrix(cairo.Matrix(
|
||||
fdiv(pw, fw), 0,
|
||||
0, -fdiv(ph, fh),
|
||||
(pw / 2) - fc[0] * fdiv(pw, fw),
|
||||
(ph / 2) + fc[1] * fdiv(ph, fh),
|
||||
))
|
||||
self.cache_cairo_context(pixel_array, ctx)
|
||||
return ctx
|
||||
|
||||
def display_multiple_vectorized_mobjects(self, vmobjects, pixel_array):
|
||||
if len(vmobjects) == 0:
|
||||
return
|
||||
batch_file_pairs = batch_by_property(
|
||||
vmobjects,
|
||||
lambda vm: vm.get_background_image_file()
|
||||
)
|
||||
for batch, file_name in batch_file_pairs:
|
||||
if file_name:
|
||||
self.display_multiple_background_colored_vmobject(batch, pixel_array)
|
||||
else:
|
||||
self.display_multiple_non_background_colored_vmobjects(batch, pixel_array)
|
||||
shader['scale'].value = fh / 2
|
||||
shader['aspect_ratio'].value = fw / fh
|
||||
shader['anti_alias_width'].value = ANTI_ALIAS_WIDTH
|
||||
|
||||
def display_multiple_non_background_colored_vmobjects(self, vmobjects, pixel_array):
|
||||
ctx = self.get_cairo_context(pixel_array)
|
||||
for vmobject in vmobjects:
|
||||
self.display_vectorized(vmobject, ctx)
|
||||
|
||||
def display_vectorized(self, vmobject, ctx):
|
||||
self.set_cairo_context_path(ctx, vmobject)
|
||||
self.apply_stroke(ctx, vmobject, background=True)
|
||||
self.apply_fill(ctx, vmobject)
|
||||
self.apply_stroke(ctx, vmobject)
|
||||
return self
|
||||
|
||||
def set_cairo_context_path(self, ctx, vmobject):
|
||||
points = self.transform_points_pre_display(
|
||||
vmobject, vmobject.points
|
||||
)
|
||||
# TODO, shouldn't this be handled in transform_points_pre_display?
|
||||
# points = points - self.get_frame_center()
|
||||
if len(points) == 0:
|
||||
return
|
||||
|
||||
ctx.new_path()
|
||||
subpaths = vmobject.get_subpaths_from_points(points)
|
||||
for subpath in subpaths:
|
||||
quads = vmobject.get_cubic_bezier_tuples_from_points(subpath)
|
||||
ctx.new_sub_path()
|
||||
start = subpath[0]
|
||||
ctx.move_to(*start[:2])
|
||||
for p0, p1, p2, p3 in quads:
|
||||
ctx.curve_to(*p1[:2], *p2[:2], *p3[:2])
|
||||
if vmobject.consider_points_equals(subpath[0], subpath[-1]):
|
||||
ctx.close_path()
|
||||
return self
|
||||
|
||||
def set_cairo_context_color(self, ctx, rgbas, vmobject):
|
||||
if len(rgbas) == 1:
|
||||
# Use reversed rgb because cairo surface is
|
||||
# encodes it in reverse order
|
||||
ctx.set_source_rgba(
|
||||
*rgbas[0][2::-1], rgbas[0][3]
|
||||
)
|
||||
else:
|
||||
points = vmobject.get_gradient_start_and_end_points()
|
||||
points = self.transform_points_pre_display(
|
||||
vmobject, points
|
||||
)
|
||||
pat = cairo.LinearGradient(*it.chain(*[
|
||||
point[:2] for point in points
|
||||
]))
|
||||
step = 1.0 / (len(rgbas) - 1)
|
||||
offsets = np.arange(0, 1 + step, step)
|
||||
for rgba, offset in zip(rgbas, offsets):
|
||||
pat.add_color_stop_rgba(
|
||||
offset, *rgba[2::-1], rgba[3]
|
||||
)
|
||||
ctx.set_source(pat)
|
||||
return self
|
||||
|
||||
def apply_fill(self, ctx, vmobject):
|
||||
self.set_cairo_context_color(
|
||||
ctx, self.get_fill_rgbas(vmobject), vmobject
|
||||
)
|
||||
ctx.fill_preserve()
|
||||
return self
|
||||
|
||||
def apply_stroke(self, ctx, vmobject, background=False):
|
||||
width = vmobject.get_stroke_width(background)
|
||||
if width == 0:
|
||||
return self
|
||||
self.set_cairo_context_color(
|
||||
ctx,
|
||||
self.get_stroke_rgbas(vmobject, background=background),
|
||||
vmobject
|
||||
)
|
||||
ctx.set_line_width(
|
||||
width * self.cairo_line_width_multiple *
|
||||
# This ensures lines have constant width
|
||||
# as you zoom in on them.
|
||||
(self.get_frame_width() / FRAME_WIDTH)
|
||||
)
|
||||
ctx.stroke_preserve()
|
||||
return self
|
||||
|
||||
def get_stroke_rgbas(self, vmobject, background=False):
|
||||
return vmobject.get_stroke_rgbas(background)
|
||||
|
||||
def get_fill_rgbas(self, vmobject):
|
||||
return vmobject.get_fill_rgbas()
|
||||
|
||||
def get_background_colored_vmobject_displayer(self):
|
||||
# Quite wordy to type out a bunch
|
||||
bcvd = "background_colored_vmobject_displayer"
|
||||
if not hasattr(self, bcvd):
|
||||
setattr(self, bcvd, BackgroundColoredVMobjectDisplayer(self))
|
||||
return getattr(self, bcvd)
|
||||
|
||||
def display_multiple_background_colored_vmobject(self, cvmobjects, pixel_array):
|
||||
displayer = self.get_background_colored_vmobject_displayer()
|
||||
cvmobject_pixel_array = displayer.display(*cvmobjects)
|
||||
self.overlay_rgba_array(pixel_array, cvmobject_pixel_array)
|
||||
return self
|
||||
|
||||
# Methods for other rendering
|
||||
|
||||
def display_multiple_point_cloud_mobjects(self, pmobjects, pixel_array):
|
||||
for pmobject in pmobjects:
|
||||
self.display_point_cloud(
|
||||
pmobject,
|
||||
pmobject.points,
|
||||
pmobject.rgbas,
|
||||
self.adjusted_thickness(pmobject.stroke_width),
|
||||
pixel_array,
|
||||
)
|
||||
|
||||
def display_point_cloud(self, pmobject, points, rgbas, thickness, pixel_array):
|
||||
if len(points) == 0:
|
||||
return
|
||||
pixel_coords = self.points_to_pixel_coords(
|
||||
pmobject, points
|
||||
)
|
||||
pixel_coords = self.thickened_coordinates(
|
||||
pixel_coords, thickness
|
||||
)
|
||||
rgba_len = pixel_array.shape[2]
|
||||
|
||||
rgbas = (self.rgb_max_val * rgbas).astype(self.pixel_array_dtype)
|
||||
target_len = len(pixel_coords)
|
||||
factor = target_len // len(rgbas)
|
||||
rgbas = np.array([rgbas] * factor).reshape((target_len, rgba_len))
|
||||
|
||||
on_screen_indices = self.on_screen_pixels(pixel_coords)
|
||||
pixel_coords = pixel_coords[on_screen_indices]
|
||||
rgbas = rgbas[on_screen_indices]
|
||||
|
||||
ph = self.get_pixel_height()
|
||||
pw = self.get_pixel_width()
|
||||
|
||||
flattener = np.array([1, pw], dtype='int')
|
||||
flattener = flattener.reshape((2, 1))
|
||||
indices = np.dot(pixel_coords, flattener)[:, 0]
|
||||
indices = indices.astype('int')
|
||||
|
||||
new_pa = pixel_array.reshape((ph * pw, rgba_len))
|
||||
new_pa[indices] = rgbas
|
||||
pixel_array[:, :] = new_pa.reshape((ph, pw, rgba_len))
|
||||
|
||||
def display_multiple_image_mobjects(self, image_mobjects, pixel_array):
|
||||
for image_mobject in image_mobjects:
|
||||
self.display_image_mobject(image_mobject, pixel_array)
|
||||
|
||||
def display_image_mobject(self, image_mobject, pixel_array):
|
||||
corner_coords = self.points_to_pixel_coords(
|
||||
image_mobject, image_mobject.points
|
||||
)
|
||||
ul_coords, ur_coords, dl_coords = corner_coords
|
||||
right_vect = ur_coords - ul_coords
|
||||
down_vect = dl_coords - ul_coords
|
||||
center_coords = ul_coords + (right_vect + down_vect) / 2
|
||||
|
||||
sub_image = Image.fromarray(
|
||||
image_mobject.get_pixel_array(),
|
||||
mode="RGBA"
|
||||
)
|
||||
|
||||
# Reshape
|
||||
pixel_width = max(int(pdist([ul_coords, ur_coords])), 1)
|
||||
pixel_height = max(int(pdist([ul_coords, dl_coords])), 1)
|
||||
sub_image = sub_image.resize(
|
||||
(pixel_width, pixel_height), resample=Image.BICUBIC
|
||||
)
|
||||
|
||||
# Rotate
|
||||
angle = angle_of_vector(right_vect)
|
||||
adjusted_angle = -int(360 * angle / TAU)
|
||||
if adjusted_angle != 0:
|
||||
sub_image = sub_image.rotate(
|
||||
adjusted_angle, resample=Image.BICUBIC, expand=1
|
||||
)
|
||||
|
||||
# TODO, there is no accounting for a shear...
|
||||
|
||||
# Paste into an image as large as the camear's pixel array
|
||||
full_image = Image.fromarray(
|
||||
np.zeros((self.get_pixel_height(), self.get_pixel_width())),
|
||||
mode="RGBA"
|
||||
)
|
||||
new_ul_coords = center_coords - np.array(sub_image.size) / 2
|
||||
new_ul_coords = new_ul_coords.astype(int)
|
||||
full_image.paste(
|
||||
sub_image,
|
||||
box=(
|
||||
new_ul_coords[0],
|
||||
new_ul_coords[1],
|
||||
new_ul_coords[0] + sub_image.size[0],
|
||||
new_ul_coords[1] + sub_image.size[1],
|
||||
)
|
||||
)
|
||||
# Paint on top of existing pixel array
|
||||
self.overlay_PIL_image(pixel_array, full_image)
|
||||
|
||||
def overlay_rgba_array(self, pixel_array, new_array):
|
||||
self.overlay_PIL_image(
|
||||
pixel_array,
|
||||
self.get_image(new_array),
|
||||
)
|
||||
|
||||
def overlay_PIL_image(self, pixel_array, image):
|
||||
pixel_array[:, :] = np.array(
|
||||
Image.alpha_composite(
|
||||
self.get_image(pixel_array),
|
||||
image
|
||||
),
|
||||
dtype='uint8'
|
||||
)
|
||||
|
||||
def adjust_out_of_range_points(self, points):
|
||||
if not np.any(points > self.max_allowable_norm):
|
||||
return points
|
||||
norms = np.apply_along_axis(get_norm, 1, points)
|
||||
violator_indices = norms > self.max_allowable_norm
|
||||
violators = points[violator_indices, :]
|
||||
violator_norms = norms[violator_indices]
|
||||
reshaped_norms = np.repeat(
|
||||
violator_norms.reshape((len(violator_norms), 1)),
|
||||
points.shape[1], 1
|
||||
)
|
||||
rescaled = self.max_allowable_norm * violators / reshaped_norms
|
||||
points[violator_indices] = rescaled
|
||||
return points
|
||||
|
||||
def transform_points_pre_display(self, mobject, points):
|
||||
# Subclasses (like ThreeDCamera) may want to
|
||||
# adjust points futher before they're shown
|
||||
if np.any(np.isnan(points)) or np.any(points == np.inf):
|
||||
# TODO, print some kind of warning about
|
||||
# mobject having invalid points?
|
||||
points = np.zeros((1, 3))
|
||||
return points
|
||||
|
||||
def points_to_pixel_coords(self, mobject, points):
|
||||
points = self.transform_points_pre_display(
|
||||
mobject, points
|
||||
)
|
||||
shifted_points = points - self.get_frame_center()
|
||||
|
||||
result = np.zeros((len(points), 2))
|
||||
pixel_height = self.get_pixel_height()
|
||||
pixel_width = self.get_pixel_width()
|
||||
frame_height = self.get_frame_height()
|
||||
frame_width = self.get_frame_width()
|
||||
width_mult = pixel_width / frame_width
|
||||
width_add = pixel_width / 2
|
||||
height_mult = pixel_height / frame_height
|
||||
height_add = pixel_height / 2
|
||||
# Flip on y-axis as you go
|
||||
height_mult *= -1
|
||||
|
||||
result[:, 0] = shifted_points[:, 0] * width_mult + width_add
|
||||
result[:, 1] = shifted_points[:, 1] * height_mult + height_add
|
||||
return result.astype('int')
|
||||
|
||||
def on_screen_pixels(self, pixel_coords):
|
||||
return reduce(op.and_, [
|
||||
pixel_coords[:, 0] >= 0,
|
||||
pixel_coords[:, 0] < self.get_pixel_width(),
|
||||
pixel_coords[:, 1] >= 0,
|
||||
pixel_coords[:, 1] < self.get_pixel_height(),
|
||||
])
|
||||
|
||||
def adjusted_thickness(self, thickness):
|
||||
# TODO: This seems...unsystematic
|
||||
big_sum = op.add(
|
||||
PRODUCTION_QUALITY_CAMERA_CONFIG["pixel_height"],
|
||||
PRODUCTION_QUALITY_CAMERA_CONFIG["pixel_width"],
|
||||
)
|
||||
this_sum = op.add(
|
||||
self.get_pixel_height(),
|
||||
self.get_pixel_width(),
|
||||
)
|
||||
factor = fdiv(big_sum, this_sum)
|
||||
return 1 + (thickness - 1) / factor
|
||||
|
||||
def get_thickening_nudges(self, thickness):
|
||||
thickness = int(thickness)
|
||||
_range = list(range(-thickness // 2 + 1, thickness // 2 + 1))
|
||||
return np.array(list(it.product(_range, _range)))
|
||||
|
||||
def thickened_coordinates(self, pixel_coords, thickness):
|
||||
nudges = self.get_thickening_nudges(thickness)
|
||||
pixel_coords = np.array([
|
||||
pixel_coords + nudge
|
||||
for nudge in nudges
|
||||
])
|
||||
size = pixel_coords.size
|
||||
return pixel_coords.reshape((size // 2, 2))
|
||||
|
||||
# TODO, reimplement using cairo matrix
|
||||
def get_coords_of_all_pixels(self):
|
||||
# These are in x, y order, to help me keep things straight
|
||||
full_space_dims = np.array([
|
||||
self.get_frame_width(),
|
||||
self.get_frame_height()
|
||||
])
|
||||
full_pixel_dims = np.array([
|
||||
self.get_pixel_width(),
|
||||
self.get_pixel_height()
|
||||
])
|
||||
|
||||
# These are addressed in the same y, x order as in pixel_array, but the values in them
|
||||
# are listed in x, y order
|
||||
uncentered_pixel_coords = np.indices(
|
||||
[self.get_pixel_height(), self.get_pixel_width()]
|
||||
)[::-1].transpose(1, 2, 0)
|
||||
uncentered_space_coords = fdiv(
|
||||
uncentered_pixel_coords * full_space_dims,
|
||||
full_pixel_dims)
|
||||
# Could structure above line's computation slightly differently, but figured (without much
|
||||
# thought) multiplying by frame_shape first, THEN dividing by pixel_shape, is probably
|
||||
# better than the other order, for avoiding underflow quantization in the division (whereas
|
||||
# overflow is unlikely to be a problem)
|
||||
|
||||
centered_space_coords = (
|
||||
uncentered_space_coords - fdiv(full_space_dims, 2)
|
||||
)
|
||||
|
||||
# Have to also flip the y coordinates to account for pixel array being listed in
|
||||
# top-to-bottom order, opposite of screen coordinate convention
|
||||
centered_space_coords = centered_space_coords * (1, -1)
|
||||
|
||||
return centered_space_coords
|
||||
def render_from_shader(self, shader, data):
|
||||
vbo = shader.ctx.buffer(data.tobytes())
|
||||
vao = shader.ctx.simple_vertex_array(shader, vbo, *data.dtype.names)
|
||||
vao.render(moderngl.TRIANGLES) # TODO, allow different render types
|
||||
|
||||
|
||||
class BackgroundColoredVMobjectDisplayer(object):
|
||||
def __init__(self, camera):
|
||||
self.camera = camera
|
||||
self.file_name_to_pixel_array_map = {}
|
||||
self.pixel_array = np.array(camera.get_pixel_array())
|
||||
self.reset_pixel_array()
|
||||
def get_vmob_shader(ctx, type):
|
||||
vert_file = f"quadratic_bezier_{type}_vert.glsl"
|
||||
geom_file = f"quadratic_bezier_{type}_geom.glsl"
|
||||
frag_file = f"quadratic_bezier_{type}_frag.glsl"
|
||||
|
||||
def reset_pixel_array(self):
|
||||
self.pixel_array[:, :] = 0
|
||||
shader = ctx.program(
|
||||
vertex_shader=get_code_from_file(vert_file),
|
||||
geometry_shader=get_code_from_file(geom_file),
|
||||
fragment_shader=get_code_from_file(frag_file),
|
||||
)
|
||||
set_shader_uniforms(shader)
|
||||
return shader
|
||||
|
||||
def resize_background_array(
|
||||
self, background_array,
|
||||
new_width, new_height,
|
||||
mode="RGBA"
|
||||
):
|
||||
image = Image.fromarray(background_array)
|
||||
image = image.convert(mode)
|
||||
resized_image = image.resize((new_width, new_height))
|
||||
return np.array(resized_image)
|
||||
|
||||
def resize_background_array_to_match(self, background_array, pixel_array):
|
||||
height, width = pixel_array.shape[:2]
|
||||
mode = "RGBA" if pixel_array.shape[2] == 4 else "RGB"
|
||||
return self.resize_background_array(background_array, width, height, mode)
|
||||
def get_stroke_shader(ctx):
|
||||
return get_vmob_shader(ctx, "stroke")
|
||||
|
||||
def get_background_array(self, file_name):
|
||||
if file_name in self.file_name_to_pixel_array_map:
|
||||
return self.file_name_to_pixel_array_map[file_name]
|
||||
full_path = get_full_raster_image_path(file_name)
|
||||
image = Image.open(full_path)
|
||||
back_array = np.array(image)
|
||||
|
||||
pixel_array = self.pixel_array
|
||||
if not np.all(pixel_array.shape == back_array.shape):
|
||||
back_array = self.resize_background_array_to_match(
|
||||
back_array, pixel_array
|
||||
)
|
||||
def get_fill_shader(ctx):
|
||||
return get_vmob_shader(ctx, "fill")
|
||||
|
||||
self.file_name_to_pixel_array_map[file_name] = back_array
|
||||
return back_array
|
||||
|
||||
def display(self, *cvmobjects):
|
||||
batch_image_file_pairs = batch_by_property(
|
||||
cvmobjects, lambda cv: cv.get_background_image_file()
|
||||
)
|
||||
curr_array = None
|
||||
for batch, image_file in batch_image_file_pairs:
|
||||
background_array = self.get_background_array(image_file)
|
||||
pixel_array = self.pixel_array
|
||||
self.camera.display_multiple_non_background_colored_vmobjects(
|
||||
batch, pixel_array
|
||||
)
|
||||
new_array = np.array(
|
||||
(background_array * pixel_array.astype('float') / 255),
|
||||
dtype=self.camera.pixel_array_dtype
|
||||
)
|
||||
if curr_array is None:
|
||||
curr_array = new_array
|
||||
else:
|
||||
curr_array = np.maximum(curr_array, new_array)
|
||||
self.reset_pixel_array()
|
||||
return curr_array
|
||||
def render_vmob_stroke(shader, vmobs):
|
||||
assert(len(vmobs) > 0)
|
||||
data_arrays = [vmob.get_stroke_shader_data() for vmob in vmobs]
|
||||
data = join_arrays(*data_arrays)
|
||||
send_data_to_shader(shader, data)
|
||||
|
||||
|
||||
def render_vmob_fill(shader, vmobs):
|
||||
assert(len(vmobs) > 0)
|
||||
data_arrays = [vmob.get_fill_shader_data() for vmob in vmobs]
|
||||
data = join_arrays(*data_arrays)
|
||||
send_data_to_shader(shader, data)
|
||||
|
||||
Reference in New Issue
Block a user