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/>.


from io import BytesIO

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

from .utils import cached_property


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):
        self.parent = self
        self.data = data
        self.y1 = 0
        self.y2 = self.width
        self.x1 = 0
        self.x2 = self.height

    def __getitem__(self, key):
        """Return a SubImage 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

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

        y1 = self.y1 + ystart
        y2 = self.y1 + yend
        x1 = self.x1 + xstart
        x2 = self.x1 + xend
        return SubImage(self.parent, y1, y2, x1, x2)

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

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

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

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

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

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

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

        0 = white (blank) pixel
        1 = black (letter) pixel
        """

        grayscale = rgb2gray(self.data)
        return (grayscale < 1).astype('b')

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

    @property
    def key(self):
        """Return a byte string uniquely representing the image."""
        return self.tostring()

    def tostring(self):
        """Serialize the image as a string."""
        height, width, *_ = self.shape
        shape = '{}x{}'.format(height, width).encode('latin1')
        bitmap = np.packbits(self.bitmap).tostring()
        return shape + b':' + bitmap

    @classmethod
    def fromstring(cls, data):
        """Deserialize an image from a string."""
        # TODO
        raise NotImplementedError

    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, :]


class SubImage(Image):
    def __init__(self, parent, y1, y2, x1, x2):
        self.parent = parent
        self.y1 = y1
        self.y2 = y2
        self.x1 = x1
        self.x2 = x2

    @property
    def data(self):
        return self.parent.data[self.y1:self.y2, self.x1:self.x2]

    @property
    def bitmap(self):
        return self.parent.bitmap[self.y1:self.y2, self.x1:self.x2]

    @property
    def T(self):
        return type(self)(self.parent.T, self.x1, self.x2, self.y1, self.y2)

    def _iter_lines(self, min_space, T=False):
        line_start = None
        prev_line_end = 0

        for i, row in enumerate(self.bitmap):
            if _is_nonblank(row):
                if line_start is None:
                    line_start = i
                    height = line_start - prev_line_end
                    if height >= min_space:
                        yield self[prev_line_end:line_start]
            else:
                if line_start is not None:
                    yield self[line_start:i,:]
                    line_start = None  
                    prev_line_end = i