Skip to content

Commit

Permalink
Merge remote-tracking branch 'origin/master'
Browse files Browse the repository at this point in the history
  • Loading branch information
goncalopp committed Jun 13, 2018
2 parents b8f4883 + 8daaedc commit 8611fc9
Show file tree
Hide file tree
Showing 11 changed files with 177 additions and 115 deletions.
6 changes: 3 additions & 3 deletions example.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from simpleocr.files import open_img
from simpleocr.files import open_image
from simpleocr.segmentation import ContourSegmenter
from simpleocr.feature_extraction import SimpleFeatureExtractor
from simpleocr.classification import KNNClassifier
Expand All @@ -9,9 +9,9 @@
classifier = KNNClassifier()
ocr = OCR(segmenter, extractor, classifier)

ocr.train(open_img('digits1'))
ocr.train(open_image('digits1'))

test_image = open_img('digits2')
test_image = open_image('digits2')
test_chars, test_classes, test_segments = ocr.ocr(test_image, show_steps=True)

print("accuracy:", accuracy(test_image.ground.classes, test_classes))
Expand Down
6 changes: 3 additions & 3 deletions example_grounding.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
from simpleocr.files import open_img
from simpleocr.files import open_image
from simpleocr.grounding import UserGrounder
from simpleocr.segmentation import ContourSegmenter, draw_segments
from simpleocr.segmentation import ContourSegmenter

segmenter = ContourSegmenter(blur_y=5, blur_x=5, block_size=11, c=10)
new_image = open_img('digits1')
new_image = open_image('digits1')
segments = segmenter.process(new_image.image)

grounder = UserGrounder()
Expand Down
6 changes: 4 additions & 2 deletions simpleocr/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,14 +42,16 @@ def is_python_3():
# Classifiers
from simpleocr.classification import KNNClassifier
# Files
from simpleocr.files import ImageFile
from simpleocr.files import open_image, Image, ImageFile
# Grounders
from simpleocr.grounding import TerminalGrounder, TextGrounder, UserGrounder
# Improver functions
from simpleocr.improver import enhance_image, crop_image, imagefile_to_pillow
from simpleocr.improver import enhance_image, crop_image, image_to_pil
# OCR functions
from simpleocr.ocr import reconstruct_chars, show_differences, OCR
# Segmenters
from simpleocr.segmentation import RawContourSegmenter, ContourSegmenter
# Extraction
from simpleocr.feature_extraction import FeatureExtractor, SimpleFeatureExtractor
# Pillow functions
from simpleocr.pillow_utils import pil_to_image
159 changes: 108 additions & 51 deletions simpleocr/files.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import os
import functools
from pkg_resources import resource_filename
import cv2
from .tesseract_utils import read_boxfile, write_boxfile
Expand All @@ -10,80 +9,138 @@
GROUND_EXTENSIONS_DEFAULT = GROUND_EXTENSIONS[0]


def open_img(path_or_name):
"""
Fuzzy finds a image file given a absolute or relative path, or a name.
The name might have no extension, or be in the DATA_DIRECTORY
"""
try_img_ext = functools.partial(try_extensions, IMAGE_EXTENSIONS)
data_dir_path = os.path.join(DATA_DIRECTORY, path_or_name)
path = path_or_name
if not os.path.exists(path):
# proceed even when there's no result. ImageFile decides on the exception to raise
path = try_img_ext(path_or_name) or try_img_ext(data_dir_path) or path
return ImageFile(path)

def try_extensions(extensions, path):
"""checks if various extensions of a path exist"""
"""Checks for various extensions of a path exist if the extension is appended"""
for ext in [""] + extensions:
if os.path.exists(path + ext):
return path + ext
return None


class GroundFile(object):
"""A file with ground truth data about a image (i.e.: characters and their position)"""
def __init__(self, path):
def open_image(path):
return ImageFile(get_file_path(path))


def get_file_path(path, ground=False):
"""Get the absolute path for an image or ground file.
The path can be either absolute, relative to the CWD or relative to the
DATA_DIRECTORY. The file extension may be omitted.
:param path: image path (str)
:param ground: whether the file must be a ground file
:return: The absolute path to the file requested
"""
extensions = GROUND_EXTENSIONS if ground else IMAGE_EXTENSIONS
# If the path exists, return the path, but make sure it's an absolute path first
if os.path.exists(path):
return os.path.abspath(path)
# Try to find the file with the passed path with the various extensions
image_with_extension = try_extensions(extensions, os.path.splitext(path)[0])
if image_with_extension:
return os.path.abspath(image_with_extension)
# The file must be in the data directory if it has not yet been found
image_datadir = try_extensions(extensions, os.path.join(DATA_DIRECTORY, path))
if image_datadir:
return os.path.abspath(image_datadir)
raise IOError # file not found


class Ground(object):
"""Data class that includes labeled characters of an Image and their positions"""
def __init__(self, segments, classes):
self.segments = segments
self.classes = classes


class GroundFile(Ground):
"""Ground with file support. This class can write the data
to a box file so it can be restored when the image file the ground data belongs
to is opened again.
"""
def __init__(self, path, segments, classes):
Ground.__init__(self, segments, classes)
self.path = path
self.segments = None
self.classes = None

def read(self):
"""Update the ground data stored by reading the box file from disk"""
self.classes, self.segments = read_boxfile(self.path)

def write(self):
"""Write a new box file to disk containing the stored ground data"""
write_boxfile(self.path, self.classes, self.segments)


class ImageFile(object):
"""
An OCR image file. Has an image and its file path, and optionally
a ground (ground segments and classes) and it's file path
"""
class Image(object):
"""An image stored in memory. It optionally contains a Ground"""
def __init__(self, array):
""":param array: array with image data, must be OpenCV compatible
"""
self._image = array
self._ground = None

def __init__(self, image_path):
if not os.path.exists(image_path):
raise IOError("Image file not found: " + image_path)
self.image_path = image_path
self.image = cv2.imread(self.image_path)
basepath = os.path.splitext(image_path)[0]
self.ground_path = try_extensions(GROUND_EXTENSIONS, basepath)
if self.ground_path:
self.ground = GroundFile(self.ground_path)
self.ground.read()
else:
self.ground_path = basepath + GROUND_EXTENSIONS_DEFAULT
self.ground = None
def set_ground(self, segments, classes):
"""Creates the ground data"""
self._ground = Ground(segments=segments, classes=classes)

def remove_ground(self):
"""Removes the grounding data for the Image"""
self._ground = None

# These properties prevent the user from altering the attributes stored within
# the object and thus emphasize the immutability of the object
@property
def image(self):
return self._image

@property
def is_grounded(self):
"""checks if this file is grounded"""
return not (self.ground is None)
return not (self._ground is None)

@property
def ground(self):
return self._ground


class ImageFile(Image):
"""
Complete class that contains functions for creation from file.
Also supports grounding in memory.
"""
def __init__(self, path):
"""
:param path: path to the image to read, must be valid and absolute
"""
if not os.path.isabs(path):
raise ValueError("path value is not absolute: {0}".format(path))
array = cv2.imread(path)
Image.__init__(self, array)
self._path = path
basepath = os.path.splitext(path)[0]
self._ground_path = try_extensions(GROUND_EXTENSIONS, basepath)
if self._ground_path:
self._ground = GroundFile(self._ground_path, None, None)
self._ground.read()
else:
self._ground_path = basepath + GROUND_EXTENSIONS_DEFAULT
self._ground = None

def set_ground(self, segments, classes, write_file=False):
"""creates the ground, saves it to a file"""
if self.is_grounded:
print("Warning: grounding already grounded file")
self.ground = GroundFile(self.ground_path)
self.ground.segments = segments
self.ground.classes = classes
"""Creates the ground, saves it to a file"""
self._ground = GroundFile(self._ground_path, segments=segments, classes=classes)
if write_file:
self.ground.write()

def remove_ground(self, remove_file=False):
"""removes ground, optionally deleting it's file"""
if not self.is_grounded:
print("Warning: ungrounding ungrounded file")
self.ground = None
"""Removes ground, optionally deleting it's file"""
self._ground = None
if remove_file:
os.remove(self.ground_path)
os.remove(self._ground_path)

@property
def path(self):
return self._path

@property
def ground_path(self):
return self._ground_path


33 changes: 6 additions & 27 deletions simpleocr/improver.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
from PIL import ImageEnhance, Image, ImageOps
import numpy
import cv2
from PIL import ImageEnhance, ImageOps
from .pillow_utils import image_to_pil, pil_to_cv_array

"""
These functions are not suitable for use on images to be grounded and then trained, as the file on disk is not actually
Expand All @@ -22,7 +21,7 @@ def enhance_image(imagefile, color=None, brightness=None, contrast=None, sharpne
:param invert: Invert the colors of the image, bool
:return: modified ImageFile object, with no changes written to the actual file
"""
image = imagefile_to_pillow(imagefile)
image = image_to_pil(imagefile)
if color is not None:
image = ImageEnhance.Color(image).enhance(color)
if brightness is not None:
Expand All @@ -33,7 +32,7 @@ def enhance_image(imagefile, color=None, brightness=None, contrast=None, sharpne
image = ImageEnhance.Sharpness(image).enhance(sharpness)
if invert:
image = ImageOps.invert(image)
imagefile.image = pillow_to_numpy(image)
imagefile.image = pil_to_cv_array(image)
return imagefile


Expand All @@ -49,27 +48,7 @@ def crop_image(imagefile, box):
raise ValueError("The box parameter is not a tuple")
if not len(box) == 4:
raise ValueError("The box parameter does not have length 4")
image = imagefile_to_pillow(imagefile)
image = image_to_pil(imagefile)
image.crop(box)
imagefile.image = pillow_to_numpy(image)
imagefile.image = pil_to_cv_array(image)
return imagefile


def imagefile_to_pillow(imagefile):
"""
Convert an ImageFile object to a Pillow Image object
:param imagefile: ImageFile object
:return: Image object
"""
pillow = cv2.cvtColor(imagefile.image, cv2.COLOR_BGR2RGB)
return Image.fromarray(pillow)


def pillow_to_numpy(pillow):
"""
Convert a Pillow Image object to an ImageFile object
:param pillow: Image object
:return: cv2 compatible array that fits into ImageFile.image
"""
imagefile = numpy.array(pillow)
return imagefile[:, :, ::-1].copy()
14 changes: 7 additions & 7 deletions simpleocr/ocr.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from . import classification as classifiers
from . import feature_extraction as extractors
from . import grounding as grounders
from .files import ImageFile
from .files import open_image, Image
from six import unichr

SEGMENTERS = {
Expand Down Expand Up @@ -57,17 +57,17 @@ def __init__(self, segmenter=None, extractor=None, classifier=None, grounder=Non

def train(self, image_file):
"""feeds the training data to the OCR"""
if not isinstance(image_file, ImageFile):
image_file = ImageFile(image_file)
if not isinstance(image_file, Image):
image_file = open_image(image_file)
if not image_file.is_grounded:
raise Exception("The provided file is not grounded")
features = self.extractor.extract(image_file.image, image_file.ground.segments)
self.classifier.train(features, image_file.ground.classes)

def ocr(self, image_file, show_steps=False):
"""performs ocr used trained classifier"""
if not isinstance(image_file, ImageFile):
image_file = ImageFile(image_file)
if not isinstance(image_file, Image):
image_file = open_image(image_file)
segments = self.segmenter.process(image_file.image)
if show_steps:
self.segmenter.display()
Expand All @@ -83,8 +83,8 @@ def ground(self, image_file, text=None):
:param text: The text, if self.grounder is a TextGrounder (defaults to None)
:return:
"""
if not isinstance(image_file, ImageFile):
image_file = ImageFile(image_file)
if not isinstance(image_file, Image):
image_file = open_image(image_file)
segments = self.segmenter.process(image_file.image)
if isinstance(self.grounder, grounders.TextGrounder):
if not text:
Expand Down
24 changes: 24 additions & 0 deletions simpleocr/pillow_utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
from .files import Image
from PIL import Image
import numpy
import cv2


def image_to_pil(imagefile):
"""Convert an ImageFile or ImageBuffer object to a Pillow Image object
:param imagefile: ImageFile object
:return: Image object
"""
pillow = cv2.cvtColor(imagefile.image, cv2.COLOR_BGR2RGB)
return Image.fromarray(pillow)


def pil_to_image(pillow):
"""Convert a Pillow Image object to an ImageBuffer object"""
return Image.fromarray(pil_to_cv_array(pillow))


def pil_to_cv_array(pillow):
"""Convert a Pillow Image object to a cv compatible array"""
imagefile = numpy.array(pillow)
return imagefile[:, :, ::-1].copy()
Loading

0 comments on commit 8611fc9

Please sign in to comment.