mirror of
https://github.com/OpenVGLab/OmniLottie.git
synced 2026-09-16 23:26:26 +00:00
249 lines
9.1 KiB
Python
249 lines
9.1 KiB
Python
# NOTE: requires pillow, pypotrace>=0.2, numpy, scipy to be installed
|
|
from PIL import Image
|
|
import potrace
|
|
import numpy
|
|
import enum
|
|
from scipy.cluster.vq import kmeans
|
|
from .. import objects
|
|
from ..nvector import NVector
|
|
from .pixel import _vectorizing_func
|
|
|
|
|
|
class QuanzationMode(enum.Enum):
|
|
Nearest = 1
|
|
Exact = 2
|
|
|
|
|
|
class RasterImage:
|
|
def __init__(self, data):
|
|
self.data = data
|
|
|
|
@classmethod
|
|
def from_pil(cls, image):
|
|
return cls(numpy.array(image))
|
|
|
|
#@classmethod
|
|
#def open(cls, filename):
|
|
#return cls.from_pil(Image.open(filename))
|
|
|
|
def k_means(self, n_colors):
|
|
"""!
|
|
Returns a list of centroids
|
|
"""
|
|
colors = []
|
|
for row in range(self.data.shape[0]):
|
|
for column in range(self.data.shape[1]):
|
|
if self.get_alpha(row, column) == 255:
|
|
colors.append(self.data[row][column])
|
|
|
|
colors = numpy.array(colors, numpy.float)
|
|
return kmeans(colors, n_colors+1)[0]
|
|
|
|
def get_alpha(self, row, column):
|
|
if self.data.shape[2] >= 4:
|
|
return self.data[row][column][3]
|
|
return 255
|
|
|
|
def quantize(self, codebook, quantization_mode=QuanzationMode.Nearest):
|
|
"""!
|
|
Returns a list of tuple [color, data] where for each color in codebook
|
|
data is a bit mask for the image
|
|
|
|
You can get codebook from k_means
|
|
"""
|
|
if codebook is None or len(codebook) == 0:
|
|
return [(numpy.array([0., 0., 0., 255.]), self.mono())]
|
|
|
|
mono_data = []
|
|
for c in codebook:
|
|
mono_data.append((c, numpy.zeros(self.data.shape[:2])))
|
|
|
|
for row in range(self.data.shape[0]):
|
|
for column in range(self.data.shape[1]):
|
|
if self.get_alpha(row, column) == 255:
|
|
if quantization_mode == QuanzationMode.Nearest:
|
|
min_norm = 511 # (norm of [255, 255, 255, 255]) + 1
|
|
best = None
|
|
for color, bitmap in mono_data:
|
|
norm = numpy.linalg.norm(self.data[row][column] - color)
|
|
if norm < min_norm:
|
|
min_norm = norm
|
|
best = bitmap
|
|
if norm == 0:
|
|
break
|
|
best[row][column] = 1
|
|
else:
|
|
for color, bitmap in mono_data:
|
|
if numpy.array_equal(color, self.data[row][column]):
|
|
bitmap[row][column] = 1
|
|
break
|
|
|
|
return mono_data
|
|
|
|
def mono(self):
|
|
"""!
|
|
Returns a bit mask of opaque pixels
|
|
"""
|
|
mono_data = numpy.zeros(self.data.shape[:2])
|
|
for row in range(self.data.shape[0]):
|
|
for column in range(self.data.shape[1]):
|
|
mono_data[row][column] = int(self.data[row][column][3] == 255)
|
|
return mono_data
|
|
|
|
|
|
class Vectorizer:
|
|
def __init__(self):
|
|
self.palette = None
|
|
self.layers = {}
|
|
|
|
def _create_layer(self, animation, layer_name):
|
|
layer = animation.add_layer(objects.ShapeLayer())
|
|
if layer_name:
|
|
self.layers[layer_name] = layer
|
|
layer.name = layer_name
|
|
return layer
|
|
|
|
def prepare_layer(self, animation, layer_name=None):
|
|
layer = self._create_layer(animation, layer_name)
|
|
layer._max_verts = {}
|
|
if self.palette is None:
|
|
group = layer.add_shape(objects.Group())
|
|
group.name = "bitmap"
|
|
layer._max_verts[group.name] = 0
|
|
group.add_shape(objects.Path())
|
|
group.add_shape(objects.Fill(NVector(0, 0, 0)))
|
|
else:
|
|
for color in self.palette:
|
|
group = layer.add_shape(objects.Group())
|
|
group.name = "color_%s" % "".join("%02x" % int(c) for c in color)
|
|
layer._max_verts[group.name] = 0
|
|
fcol = color/255
|
|
fill = group.add_shape(objects.Fill(NVector(*fcol)))
|
|
if len(fcol) > 3 and fcol[3] < 1:
|
|
fill.opacity.value = fcol[3] * 100
|
|
return layer
|
|
|
|
def raster_to_layer(self, animation, raster, layer_name=None, mode=QuanzationMode.Nearest):
|
|
layer = self.prepare_layer(animation, layer_name)
|
|
mono_data = raster.quantize(self.palette, mode)
|
|
for (color, bitmap), group in zip(mono_data, layer.shapes):
|
|
self.raster_to_shapes(group, bitmap)
|
|
return layer
|
|
|
|
def raster_to_shapes(self, group, mono_data):
|
|
shapes = []
|
|
for bezier in self.raster_to_bezier(mono_data):
|
|
shape = group.insert_shape(0, objects.Path())
|
|
shapes.append(shape)
|
|
shape.shape.value = bezier
|
|
return shapes
|
|
|
|
def raster_to_bezier(self, mono_data):
|
|
bmp = potrace.Bitmap(mono_data)
|
|
path = bmp.trace()
|
|
shapes = []
|
|
for curve in path:
|
|
bezier = objects.Bezier()
|
|
shapes.append(bezier)
|
|
bezier.add_point(NVector(*curve.start_point))
|
|
for segment in curve:
|
|
if segment.is_corner:
|
|
bezier.add_point(NVector(*segment.c))
|
|
bezier.add_point(NVector(*segment.end_point))
|
|
else:
|
|
sp = NVector(*bezier.vertices[-1])
|
|
ep = NVector(*segment.end_point)
|
|
c1 = NVector(*segment.c1) - sp
|
|
c2 = NVector(*segment.c2) - ep
|
|
bezier.out_tangents[-1] = c1
|
|
bezier.add_point(ep, c2)
|
|
return shapes
|
|
|
|
def _frame_keyframe(self, layer, group, time, shapes, beziers):
|
|
if shapes:
|
|
# TODO handle multiple shapes
|
|
nverts = len(beziers[0].vertices)
|
|
if nverts > layer._max_verts[group.name]:
|
|
layer._max_verts[group.name] = nverts
|
|
for shape, bezier in zip(shapes, beziers):
|
|
shape.shape.add_keyframe(time, bezier)
|
|
|
|
def raster_to_frame(self, animation, raster, layer_name, time, mode=QuanzationMode.Nearest):
|
|
mono_data = raster.quantize(self.palette, mode)
|
|
|
|
if layer_name not in self.layers:
|
|
layer = self.prepare_layer(animation, layer_name)
|
|
for (color, bitmap), group in zip(mono_data, layer.shapes):
|
|
shapes = self.raster_to_shapes(group, bitmap)
|
|
beziers = [s.shape.value for s in shapes]
|
|
self._frame_keyframe(layer, group, time, shapes, beziers)
|
|
else:
|
|
layer = self.layers[layer_name]
|
|
|
|
for (color, bitmap), group in zip(mono_data, layer.shapes):
|
|
shapes = [s for s in group.shapes if isinstance(s, objects.Path)]
|
|
beziers = self.raster_to_bezier(bitmap)
|
|
self._frame_keyframe(layer, group, time, shapes, beziers)
|
|
|
|
def adjust_missing_vertices(self, layer_name):
|
|
layer = self.layers[layer_name]
|
|
for group in layer.shapes:
|
|
# TODO handle multiple shapes
|
|
shape = group.shapes[0]
|
|
nverts = layer._max_verts[group.name]
|
|
if shape.shape.animated:
|
|
for kf in shape.shape.keyframes:
|
|
bezier = kf.start
|
|
count = nverts - len(bezier.vertices)
|
|
bezier.vertices += [bezier.vertices[-1]] * count
|
|
bezier.in_tangents += [NVector(0, 0)] * count
|
|
bezier.out_tangents += [NVector(0, 0)] * count
|
|
|
|
def duplicate_start_frame(self, layer_name, time):
|
|
layer = self.layers[layer_name]
|
|
for group in layer.shapes:
|
|
shape = group.shapes[0]
|
|
bezier = shape.shape.keyframes[0].start
|
|
group.shapes[0].shape.add_keyframe(time, bezier)
|
|
|
|
|
|
def color2numpy(vcolor):
|
|
l = (vcolor * 255).components
|
|
if len(l) == 3:
|
|
l.append(255)
|
|
return numpy.array(l, numpy.uint8)
|
|
|
|
|
|
def raster_to_animation(filenames, n_colors=1, frame_delay=1,
|
|
looping=True, framerate=60, palette=[],
|
|
mode=QuanzationMode.Nearest):
|
|
|
|
vc = Vectorizer()
|
|
|
|
def callback(animation, raster, frame):
|
|
raster = RasterImage.from_pil(raster)
|
|
if vc.palette is None:
|
|
if palette:
|
|
vc.palette = [color2numpy(c) for c in palette]
|
|
elif n_colors > 1:
|
|
vc.palette = raster.k_means(n_colors)
|
|
#vc.raster_to_frame(animation, raster, "anim", frame * frame_delay, mode)
|
|
layer = vc.raster_to_layer(animation, raster, "frame_%s" % frame, mode)
|
|
layer.in_point = frame * frame_delay
|
|
layer.out_point = (frame + 1) * frame_delay
|
|
|
|
animation = _vectorizing_func(filenames, frame_delay, framerate, callback)
|
|
|
|
#vc.adjust_missing_vertices("anim")
|
|
#if looping and animation._nframes > 1:
|
|
#animation.out_point += frame_delay
|
|
#vc.duplicate_start_frame("anim", animation.out_point)
|
|
#elif animation._nframes == 1:
|
|
#for g in animation.find("anim").shapes:
|
|
#for shape in g.find_all(objects.Path):
|
|
#shape.shape.clear_animation(shape.shape.get_value(0))
|
|
|
|
#animation.find("anim").out_point = animation.out_point
|
|
|
|
return animation
|