pixelocr/pixelocr/image.py

# Copyright (C) 2014  Andrey Golovizin
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.


import itertools
from io import BytesIO

import numpy as np
from skimage.io import imread, imsave

from .utils import cached_property, pairwise


def bbox(*images):
    return (
        min(image.left for image in images),
        min(image.top for image in images),
        max(image.right for image in images),
        max(image.bottom for image in images),
    )


def overlap(*images):
    for image1, image2 in itertools.combinations(images, 2):
        xoverlap = (image1.right > image2.left and image2.right > image1.left)
        yoverlap = (image1.bottom > image2.top and image2.bottom > image1.top)
        if xoverlap and yoverlap:
            return True
    return False


def combine(*images):
    # TODO: properly combine overlapping glyphs
    # if overlap(*images):
    #     raise NotImplementedError

    left, top, right, bottom = bbox(*images)
    width = right - left
    height = bottom - top

    data = np.zeros((height, width, images[0].data.shape[2]), images[0].data.dtype)
    data.fill(255)

    for image in images:
        xoffset = image.left - left
        yoffset = image.top - top
        data[yoffset:yoffset + image.height, xoffset:xoffset + image.width] = image.data
    return Image(data, left, top)


def is_nonblank(bitmap):
    """Return True if bitmap contains at least one black (=1) pixel."""
    return bitmap.any()


class Image(object):
    """Basic image class."""

    def __init__(self, data, x=0, y=0):
        self.data = data
        self.x = x
        self.y = y

    def __getitem__(self, key):
        """Return an Image for the specified region."""

        def indices(sliceobj, length):
            """Decode a slice object and return a pair of end:start indices."""
            if sliceobj is None:
                return 0, length
            elif isinstance(sliceobj, int):
                return sliceobj, sliceobj + 1
            elif isinstance(sliceobj, slice):
                start, end, stride = sliceobj.indices(length)
                if stride != 1:
                    raise NotImplementedError
                return start, end
            else:
                raise NotImplementedError(sliceobj)

        if not isinstance(key, tuple):
            yslice = key
            xslice = None
        else:
            yslice, xslice = key

        xstart, xend = indices(xslice, self.width)
        ystart, yend = indices(yslice, self.height)

        x = self.x + xstart
        y = self.y + ystart
        return Image(self.data[key], x, y)

    def space(self, x, y, width, height):
        data = np.zeros((height, width, self.shape[2]), dtype=self.data.dtype)
        data.fill(255)
        return Image(data, x, y)

    def _repr_png_(self):
        buf = BytesIO()
        imsave(buf, self.data)
        return buf.getvalue()

    @classmethod
    def fromfile(cls, filename):
        return cls(imread(filename))

    @property
    def left(self):
        return self.x

    @property
    def right(self):
        return self.x + self.width

    @property
    def top(self):
        return self.y

    @property
    def bottom(self):
        return self.y + self.height

    @property
    def shape(self):
        return self.data.shape

    @property
    def height(self):
        return self.data.shape[0]

    @property
    def width(self):
        return self.data.shape[1]

    @cached_property
    def T(self):
        return type(self)(self.data.swapaxes(0, 1), x=self.y, y=self.x)

    @cached_property
    def bitmap(self):
        """Return a two-color version of the image.

        0 = white (blank) pixel
        1 = black (glyph) pixel
        """
        return (self.data != 255).any(axis=2)

    @cached_property
    def isspace(self):
        return not is_nonblank(self.bitmap)

    @cached_property
    def color(self):
        """Return the most frequent foreground color."""
        if self.isspace:
            return None
        mask3 = np.dstack([~self.bitmap] * 3)
        colors = np.ma.MaskedArray(self.data, mask3).reshape(-1, 3)
        r, g, b = colors.mean(axis=0)
        return int(r), int(g), int(b)

    def serialize(self):
        """Serialize the image as some hashable object."""
        bitmap = self.data.astype(np.uint8).tostring()
        return self.shape, bitmap

    @classmethod
    def deserialize(cls, obj):
        """Deserialize an image."""
        shape, data = obj
        array = np.fromstring(data, dtype=np.uint8).reshape(shape)
        return cls(array)

    @cached_property
    def qimage(self):
        from PyQt4.QtGui import QImage
        return QImage(
            self.data.astype(np.uint8).data,
            self.data.shape[1], self.data.shape[0],
            self.data.shape[1] * 3,
            QImage.Format_RGB888,
        )

    def unframe(self, width=2):
        return Image(self.data[width:-width,width:-width])

    def strip(self):
        """Strip top and bottom blank space.
        
        All-whitespace images are not stripped.
        """

        if self.isspace:
            return self

        def _get_margin_height(rows):
            for i, row in enumerate(rows):
                if is_nonblank(row):
                    return i
            return 0

        top_margin = _get_margin_height(self.bitmap)
        bottom_margin = _get_margin_height(reversed(self.bitmap))
        return self[top_margin:self.height - bottom_margin, :]

    def mask(self, mask):
        mask3 = np.dstack([mask] * 3)
        data = np.ma.masked_array(self.data, mask3, fill_value=255).filled()
        return Image(data, self.x, self.y)

    def fits(self, left, top, right, bottom):
        """Return True if the image fits into the given bounding box."""

        return (
            self.top >= top
            and self.left >= left
            and self.right <= right
            and self.bottom <= bottom
        )