refactor

Author: aflag

break_captcha

@@ -22,16 +22,10 @@
 
 import sys
 import Image
-from sklearn.externals import joblib
-from image_processing import DigitSeparator
 import urllib2
 import StringIO
 
-def detect_number(model, image):
-    digit_separator = DigitSeparator(image)
-    digits = digit_separator.get_digits()
-    labels = model.predict(digits)
-    print ''.join(map(lambda x: str(int(x)), labels))
+from sklearn.externals import joblib
 
 def main():
     model = joblib.load(sys.argv[1])
@@ -47,7 +41,7 @@ def main():
                 else:
                     with open(img_path) as f_image:
                         image = Image.open(f_image).convert('L')
-                detect_number(model, image)
+                print model.decode_image(image)
         except EOFError:
             pass
 

dataset.py

@@ -0,0 +1,38 @@
+# Copyright (C) 2012 Rafael Cunha de Almeida <[email protected]>
+#
+# Permission is hereby granted, free of charge, to any person obtaining
+# a copy of this software and associated documentation files (the
+# "Software"), to deal in the Software without restriction, including
+# without limitation the rights to use, copy, modify, merge, publish,
+# distribute, sublicense, and/or sell copies of the Software, and to
+# permit persons to whom the Software is furnished to do so, subject to
+# the following conditions:
+#
+# The above copyright notice and this permission notice shall be
+# included in all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+# IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
+# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+# OTHER DEALINGS IN THE SOFTWARE.
+
+import os
+import re
+import Image
+
+def _get_files(base_dir):
+    return map(lambda x: os.path.join(base_dir, x), os.listdir(base_dir))
+
+def load_captcha_dataset(base_dir):
+    files = _get_files(base_dir)
+    dataset = []
+    for file_path in files:
+        file_name = os.path.basename(file_path)
+        label = re.findall(r'^([0-9]+)-[0-9]+\..*$', file_name)[0]
+        with open(file_path) as f:
+            captcha = Image.open(f).convert('L')
+        dataset.append((captcha, label))
+    return zip(*dataset)  # unzip

features.py

@@ -21,7 +21,6 @@
 import ImageOps
 import ImageFilter
 import numpy
-from sklearn.feature_extraction import DictVectorizer
 
 class compose_extractors(object):
     def __init__(self, extractors):
@@ -39,24 +38,6 @@ def __call__(self, arg):
             extractor(image, image_features)
         return image_features
 
-class FeatureHandler(object):
-    def __init__(self, extractor, dataset):
-        self.extractor = extractor
-        self.vectorizer = DictVectorizer()
-        digits = self.__extract_features(dataset[0])
-        self.train_digits = self.vectorizer.fit_transform(digits).toarray()
-        self.labels = dataset[1]
-
-    def __extract_features(self, values):
-        return map(self.extractor, values)
-
-    def sklearn_format_train(self):
-        return self.train_digits,self.labels
-
-    def sklearn_format_test(self, items):
-        features = self.__extract_features(items)
-        return self.vectorizer.transform(features).toarray()
-
 def border_detection(digit):
     digit.image = digit.image.filter(ImageFilter.FIND_EDGES)
     digit.pix = digit.image.load()

models.py

@@ -12,33 +12,23 @@
 #
 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE X
-# CONSORTIUM BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
 # ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
 # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 from features import *
+from image_processing import DigitSeparator
 from functools import partial
-from sklearn import naive_bayes
-from sklearn import tree
-from sklearn import linear_model
-from sklearn import svm
-from sklearn import ensemble
-from sklearn.neighbors.nearest_centroid import NearestCentroid
-from sklearn import decomposition
+
 import time
 
-class ScikitWrapper(object):
-    def __init__(self, engine, extractors, dataset):
-        self.feature_handler = FeatureHandler(
-                compose_extractors(extractors),
-                dataset)
-        self.engine = engine
-        vector, labels = self.feature_handler.sklearn_format_train()
-        self.engine.fit(vector, labels)
+from sklearn import svm
+from sklearn import ensemble
+from sklearn.feature_extraction import DictVectorizer
 
-    def predict(self, items):
-        return self.engine.predict(self.feature_handler.sklearn_format_test(items))
+class ModelUnavailable(Exception):
+    pass
 
 ALL_EXTRACTORS = [
         x_histogram,
@@ -55,20 +45,48 @@ def predict(self, items):
         horizontal_symmetry,
 ]
 
-def NaiveBayes(dataset):
-    return ScikitWrapper(naive_bayes.MultinomialNB(), [positions], dataset)
+SVM_EXTRACTORS = [positions]
+def svm_engine():
+    return svm.SVC(kernel='poly', degree=2)
+
+FOREST_EXTRACTORS = ALL_EXTRACTORS
+def forest_engine():
+    return ensemble.RandomForestClassifier(n_estimators=50, n_jobs=2)
+
+class CaptchaDecoder(object):
+    def __init__(self, x, y):
+        self.engine = svm_engine()
+        self.feature_extractor = compose_extractors(SVM_EXTRACTORS)
+        self.fit(x,y)
 
-def DecisionTree(dataset):
-    return ScikitWrapper(tree.DecisionTreeRegressor(), [positions, reversed_horizontal_silhouette, horizontal_silhouette], dataset)
+    def fit(self, x, y):
+        digits = []
+        labels = []
+        for image,param_labels in zip(x,y):
+            separator = DigitSeparator(image)
+            digits.extend(map(self.feature_extractor, separator.get_digits()))
+            labels.extend(param_labels)
+        self.vectorizer = DictVectorizer()
+        train_array = self.vectorizer.fit_transform(digits).toarray()
+        self.engine.fit(train_array, labels)
 
-def SGD(dataset):
-    return ScikitWrapper(linear_model.SGDClassifier(loss="hinge", penalty="l2"), [positions, reversed_horizontal_silhouette, horizontal_silhouette], dataset)
+    def predict(self, x):
+        prediction = []
+        for image in x:
+           separator = DigitSeparator(image) 
+           features = map(self.feature_extractor, separator.get_digits())
+           digits = self.vectorizer.transform(features).toarray()
+           labels = self.engine.predict(digits)
+           prediction.append(''.join(map(lambda x: '%d'%x, labels)))
+        return prediction
 
-def SVM(dataset):
-    return ScikitWrapper(svm.SVC(kernel='poly', degree=2), [positions], dataset)
+    def score(self, data, labels):
+        pred_labels = self.predict(data)
+        matches = sum(map(lambda (x,y): x==y, zip(labels, pred_labels)))
+        return float(matches)/len(labels)
 
-def NN(dataset):
-    return ScikitWrapper(NearestCentroid(), [positions, reversed_horizontal_silhouette, horizontal_silhouette], dataset)
+    def decode_image(self, image):
+        return self.predict([image])[0]
 
-def RandomForest(dataset):
-    return ScikitWrapper(ensemble.RandomForestClassifier(n_estimators=50, n_jobs=2), ALL_EXTRACTORS, dataset)
+    def get_params(self, *args, **kwargs):
+        return self.engine(*args, **kwargs)

profiler

@@ -20,39 +20,21 @@
 # ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 # OTHER DEALINGS IN THE SOFTWARE.
 
-import os
 import sys
-import re
-import Image
 import time
-from sklearn.externals import joblib
-from image_processing import DigitSeparator
 import cProfile
 
-model = joblib.load(sys.argv[1])
+from sklearn.externals import joblib
 
-def make_test_dataset(files):
-    dataset = []
-    for file_path in files:
-        file_name = os.path.basename(file_path)
-        label = re.findall(r'^([0-9]+)-[0-9]+\..*$', file_name)[0]
-        with open(file_path) as f:
-            digits = DigitSeparator(Image.open(f).convert("L")).get_digits()
-        dataset.append((label, digits))
-    return dataset
+from dataset import load_captcha_dataset
+
+model = joblib.load(sys.argv[1])
 
 def prof():
-    base_dir = sys.argv[2]
-    files = map(lambda x: os.path.join(base_dir, x), os.listdir(base_dir))
-    dataset = make_test_dataset(files)
+    dataset = load_captcha_dataset(sys.argv[2])
     t0 = time.time()
-    matches = 0
-    for labels,digits in dataset:
-        pred_labels = model.predict(digits)
-        if labels == ''.join(map(lambda x: str(int(x)), pred_labels)):
-            matches += 1
+    print 'Matches:', model.score(dataset[0], dataset[1])
     spent_time = time.time() - t0
     print 'Spent time:', spent_time, 'avg per predict:', spent_time/len(dataset)
-    print 'Matches:', float(matches)/len(dataset)
 
 cProfile.run('prof()', filename='/tmp/profile')

train.py

@@ -18,73 +18,41 @@
 # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 import sys
-import re
 import random
-import os
-import pickle
+
 from models import *
 from sklearn.externals import joblib
 from image_processing import DigitSeparator
 
-def make_train_dataset(files):
-    dataset = []
-    for file_path in files:
-        file_name = os.path.basename(file_path)
-        labels = re.findall(r'^([0-9]+)-[0-9]+\..*$', file_name)[0]
-        with open(file_path) as f:
-            digits = DigitSeparator(Image.open(f).convert("L")).get_digits()
-        for i,digit in enumerate(digits):
-            dataset.append((digit, labels[i]))
-    return zip(*dataset)  # unzip
-
-def make_test_dataset(files):
-    dataset = []
-    for file_path in files:
-        file_name = os.path.basename(file_path)
-        label = re.findall(r'^([0-9]+)-[0-9]+\..*$', file_name)[0]
-        with open(file_path) as f:
-            digits = DigitSeparator(Image.open(f).convert("L")).get_digits()
-        dataset.append((label, digits))
-    return dataset
-
-def get_files(base_dir):
-    return map(lambda x: os.path.join(base_dir, x), os.listdir(base_dir))
+from dataset import load_captcha_dataset
 
 def generate_datasets(base_dir):
-    files = get_files(base_dir)
-    random.shuffle(files)
-    train_size = int(0.4*len(files))
-    train = files[:train_size]
-    test = files[train_size:]
-    print "Number of trains:", len(train), "Number of tests:", len(test)
-    train_dataset = make_train_dataset(train)
-    test_dataset = make_test_dataset(test)
+    dataset = load_captcha_dataset(base_dir)
+    ziped_dataset = zip(*dataset)
+    random.shuffle(ziped_dataset)
+    dataset = zip(*ziped_dataset)
+    train_size = int(0.4*len(ziped_dataset))
+    train_dataset = (dataset[0][:train_size], dataset[1][:train_size])
+    test_dataset = (dataset[0][train_size:], dataset[1][train_size:])
+    print "Number of trains:", len(train_dataset), "Number of tests:", len(test_dataset)
     return train_dataset, test_dataset
 
 def largest_label_size(dataset):
     return max(map(len, dataset.values()))
 
-def test(model, test_dataset):
-    matches = 0
-    for labels,digits in test_dataset:
-        pred_labels = model.predict(digits)
-        if labels == ''.join(map(lambda x: str(int(x)), pred_labels)):
-            matches += 1
-    print 'Matches:', float(matches)/len(test_dataset)
-
 def main():
     if len(sys.argv) > 2:
-        train_dataset = make_train_dataset(get_files(sys.argv[1]))
+        train_dataset = load_captcha_dataset(sys.argv[1])
     else:
         train_dataset, test_dataset = generate_datasets(sys.argv[1])
     t0 = time.time()
-    model = SVM(train_dataset)
+    model = CaptchaDecoder(train_dataset[0], train_dataset[1])
     print 'Train time:', time.time() - t0
     if len(sys.argv) > 2:
         joblib.dump(model, sys.argv[2])
     else:
         t0 = time.time()
-        test(model, test_dataset)
+        print 'Matches:', model.score(test_dataset[0], test_dataset[1])
         print 'Test time:', time.time() - t0
 
 if __name__ == '__main__':