File indexing completed on 2025-01-19 03:59:54

 # 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