# -*- coding: utf-8 -*- from __future__ import with_statement import os import logging import subprocess from os.path import abspath, join from PIL import Image from PIL import TiffImagePlugin from PIL import PngImagePlugin from PIL import GifImagePlugin from PIL import JpegImagePlugin class OCR(object): __name__ = "OCR" __type__ = "ocr" __version__ = "0.1" __description__ = """OCR base plugin""" __author_name__ = "pyLoad Team" __author_mail__ = "admin@pyload.org" def __init__(self): self.logger = logging.getLogger("log") def load_image(self, image): self.image = Image.open(image) self.pixels = self.image.load() self.result_captcha = '' def unload(self): """delete all tmp images""" pass def threshold(self, value): self.image = self.image.point(lambda a: a * value + 10) def run(self, command): """Run a command""" popen = subprocess.Popen(command, bufsize=-1, stdout=subprocess.PIPE, stderr=subprocess.PIPE) popen.wait() output = popen.stdout.read() + " | " + popen.stderr.read() popen.stdout.close() popen.stderr.close() self.logger.debug("Tesseract ReturnCode %s Output: %s" % (popen.returncode, output)) def run_tesser(self, subset=False, digits=True, lowercase=True, uppercase=True): #self.logger.debug("create tmp tif") #tmp = tempfile.NamedTemporaryFile(suffix=".tif") tmp = open(join("tmp", "tmpTif_%s.tif" % self.__name__), "wb") tmp.close() #self.logger.debug("create tmp txt") #tmpTxt = tempfile.NamedTemporaryFile(suffix=".txt") tmpTxt = open(join("tmp", "tmpTxt_%s.txt" % self.__name__), "wb") tmpTxt.close() self.logger.debug("save tiff") self.image.save(tmp.name, 'TIFF') if os.name == "nt": tessparams = [join(pypath, "tesseract", "tesseract.exe")] else: tessparams = ['tesseract'] tessparams.extend([abspath(tmp.name), abspath(tmpTxt.name).replace(".txt", "")]) if subset and (digits or lowercase or uppercase): #self.logger.debug("create temp subset config") #tmpSub = tempfile.NamedTemporaryFile(suffix=".subset") tmpSub = open(join("tmp", "tmpSub_%s.subset" % self.__name__), "wb") tmpSub.write("tessedit_char_whitelist ") if digits: tmpSub.write("0123456789") if lowercase: tmpSub.write("abcdefghijklmnopqrstuvwxyz") if uppercase: tmpSub.write("ABCDEFGHIJKLMNOPQRSTUVWXYZ") tmpSub.write("\n") tessparams.append("nobatch") tessparams.append(abspath(tmpSub.name)) tmpSub.close() self.logger.debug("run tesseract") self.run(tessparams) self.logger.debug("read txt") try: with open(tmpTxt.name, 'r') as f: self.result_captcha = f.read().replace("\n", "") except: self.result_captcha = "" self.logger.debug(self.result_captcha) try: os.remove(tmp.name) os.remove(tmpTxt.name) if subset and (digits or lowercase or uppercase): os.remove(tmpSub.name) except: pass def get_captcha(self, name): raise NotImplementedError def to_greyscale(self): if self.image.mode != 'L': self.image = self.image.convert('L') self.pixels = self.image.load() def eval_black_white(self, limit): self.pixels = self.image.load() w, h = self.image.size for x in xrange(w): for y in xrange(h): if self.pixels[x, y] > limit: self.pixels[x, y] = 255 else: self.pixels[x, y] = 0 def clean(self, allowed): pixels = self.pixels w, h = self.image.size for x in xrange(w): for y in xrange(h): if pixels[x, y] == 255: continue # No point in processing white pixels since we only want to remove black pixel count = 0 try: if pixels[x - 1, y - 1] != 255: count += 1 if pixels[x - 1, y] != 255: count += 1 if pixels[x - 1, y + 1] != 255: count += 1 if pixels[x, y + 1] != 255: count += 1 if pixels[x + 1, y + 1] != 255: count += 1 if pixels[x + 1, y] != 255: count += 1 if pixels[x + 1, y - 1] != 255: count += 1 if pixels[x, y - 1] != 255: count += 1 except: pass # not enough neighbors are dark pixels so mark this pixel # to be changed to white if count < allowed: pixels[x, y] = 1 # second pass: this time set all 1's to 255 (white) for x in xrange(w): for y in xrange(h): if pixels[x, y] == 1: pixels[x, y] = 255 self.pixels = pixels def derotate_by_average(self): """rotate by checking each angle and guess most suitable""" w, h = self.image.size pixels = self.pixels for x in xrange(w): for y in xrange(h): if pixels[x, y] == 0: pixels[x, y] = 155 highest = {} counts = {} for angle in xrange(-45, 45): tmpimage = self.image.rotate(angle) pixels = tmpimage.load() w, h = self.image.size for x in xrange(w): for y in xrange(h): if pixels[x, y] == 0: pixels[x, y] = 255 count = {} for x in xrange(w): count[x] = 0 for y in xrange(h): if pixels[x, y] == 155: count[x] += 1 sum = 0 cnt = 0 for x in count.values(): if x != 0: sum += x cnt += 1 avg = sum / cnt counts[angle] = cnt highest[angle] = 0 for x in count.values(): if x > highest[angle]: highest[angle] = x highest[angle] = highest[angle] - avg hkey = 0 hvalue = 0 for key, value in highest.iteritems(): if value > hvalue: hkey = key hvalue = value self.image = self.image.rotate(hkey) pixels = self.image.load() for x in xrange(w): for y in xrange(h): if pixels[x, y] == 0: pixels[x, y] = 255 if pixels[x, y] == 155: pixels[x, y] = 0 self.pixels = pixels def split_captcha_letters(self): captcha = self.image started = False letters = [] width, height = captcha.size bottomY, topY = 0, height pixels = captcha.load() for x in xrange(width): black_pixel_in_col = False for y in xrange(height): if pixels[x, y] != 255: if not started: started = True firstX = x lastX = x if y > bottomY: bottomY = y if y < topY: topY = y if x > lastX: lastX = x black_pixel_in_col = True if black_pixel_in_col is False and started is True: rect = (firstX, topY, lastX, bottomY) new_captcha = captcha.crop(rect) w, h = new_captcha.size if w > 5 and h > 5: letters.append(new_captcha) started = False bottomY, topY = 0, height return letters def correct(self, values, var=None): if var: result = var else: result = self.result_captcha for key, item in values.iteritems(): if key.__class__ == str: result = result.replace(key, item) else: for expr in key: result = result.replace(expr, item) if var: return result else: self.result_captcha = result