Three Visits to Shredded Paper Restoration: CNN-Based Shredded Paper Restoration

三顾碎纸复原：基于CNN的碎纸复原 苏剑林 · 2016-11-25 · Machine Learning Python modeling image deep learning

Revisiting the Problem: CNN-Based Reassembly of Shredded Paper

Problem Recap

I have to say, Problem B from the 2013 National Mathematical Modeling Contest is truly a once-in-a-century gem among modeling competition problems: the statement is concise, its implications rich, its possible approaches diverse, and its potential for extension enormous — so much so that I simply can't let it go. Because of this problem, I've already written two posts on kexue.fm: Solo Modeling: Reassembling Shredded Paper and A Year Late: Revisiting the Shredded Paper Reassembly Problem. Back when I first tackled this problem, I only had a bit of mathematical modeling knowledge. Ever since learning data mining — and especially deep learning — I've wanted to redo this problem, but kept putting it off out of laziness. These past few days I finally got around to implementing it.

For readers unfamiliar with the problem statement, please refer to the two earlier posts. The shredded-paper reassembly problem comes with five attachments, representing five kinds of "shredded fragments," i.e., fragments of five different granularities. Attachments 1 and 2 are not difficult, while the real challenge concentrates in attachments 3, 4, and 5 — and the difficulty of implementing a solution for 3, 4, and 5 is essentially the same. The most natural approach to this problem is a greedy algorithm: pick an arbitrary image, find the fragment that best matches it, then continue matching the next one. For the greedy algorithm to work well, the key is finding a good distance function to judge whether two fragments are adjacent (horizontally adjacent — we don't consider vertical adjacency here). more

Both of my earlier posts used the Euclidean distance between edge vectors, and also mentioned metrics like the correlation coefficient. But when applied to attachments 3, 4, and 5, these don't work well, because the fragments in attachments 3, 4, and 5 are quite small, leaving little usable edge information. So relying on the edges alone isn't enough — we also need to take multiple factors into account, such as line spacing, average line position, and so on. But how exactly should these be combined? It's very hard to hand-craft a good function for this. So why not just hand it over to a model? Just throw a convolutional neural network (CNN) at it!

Constructing Training Samples

Specifically, after eyeballing the fragments a bit, we get a rough picture:

1. 44-point bold font;
2. If we're dealing with attachment 3, the content is Chinese; if attachment 4, it's English;
3. Fragments have a fixed size of 72×180

With these characteristics in hand, we can gather a bunch of text, then, following this specification, construct a large batch of adjacent and non-adjacent samples with the same properties ourselves, and train a convolutional neural network on them — automatically obtaining a "distance function" this way. This kind of thing is trivial for anyone familiar with deep learning. I simply gathered some Chinese text and generated a batch of Chinese samples; the code looks roughly like this:

from PIL import Image, ImageFont, ImageDraw
import numpy as np
from scipy import misc
import pymongo
from tqdm import tqdm

texts = list(pymongo.MongoClient().weixin.text_articles.find().limit(1000))
text = texts[0]['text']
line_words = 30
font_size = 44
nb_columns = line_words*font_size/72+1

def gen_img(text):
    n = len(text) / line_words + 1
    size = (nb_columns*72, (n*font_size/180+1)*180)
    im = Image.new('L', size, 255)
    dr = ImageDraw.Draw(im)
    font = ImageFont.truetype('simhei.ttf', font_size)
    for i in range(n):
        dr.text((0, 70*i), text[line_words*i: line_words*(i+1)], font=font)
    im = np.array(im.getdata()).reshape((size[1], size[0]))
    r = []
    for j in range(size[1]/180):
        for i in range(size[0]/72):
            r.append(1-im[j*180:(j+1)*180, i*72:(i+1)*72].T/255.0)
    return r

sample = []
for i in tqdm(iter(texts)):
    sample.extend(gen_img(i['text']))

np.save('sample.npy', sample)
nb_samples = len(sample) - len(sample)/nb_columns

def data(sample, batch_size):
    sample_shuffle_totally = sample[:]
    sample_shuffle_in_line = sample[:]
    while True:
        np.random.shuffle(sample_shuffle_totally)
        for i in range(0, len(sample_shuffle_in_line), nb_columns):
            subsample = sample_shuffle_in_line[i: i+nb_columns]
            np.random.shuffle(subsample)
            sample_shuffle_in_line[i: i+nb_columns] = subsample
        x = []
        y = []
        for i in range(0, len(sample), nb_columns):
            subsample_1 = sample[i: i+nb_columns]
            for j in range(0, nb_columns-1):
                x.append(np.vstack((subsample_1[j], subsample_1[j+1])))
                y.append([1])
            subsample_2 = sample_shuffle_totally[i: i+nb_columns]
            for j in range(0, nb_columns-1):
                x.append(np.vstack((subsample_2[j], subsample_2[j+1])))
                y.append([0])
            subsample_3 = sample_shuffle_in_line[i: i+nb_columns]
            for j in range(0, nb_columns-1):
                x.append(np.vstack((subsample_3[j], subsample_3[j+1])))
                y.append([0])
            if len(y) >= batch_size:
                yield np.array(x), np.array(y)
                x = []
                y = []
        if y:
            yield np.array(x), np.array(y)
            x = []
            y = []

The process works like this: I gathered 1,000 articles, each several thousand characters long, printed each article evenly onto an image, then cut it up to obtain a batch of samples. The data object below is an iterator used to generate positive and negative samples, since loading everything into memory at once is infeasible — hence the need for an iterator. Even so, I was still being a bit lazy about it, because even with this approach, it ate up 18 GB of memory on my server. sample_shuffle_totally fully shuffles the samples to produce arbitrary negative samples; sample_shuffle_in_line only shuffles within the same line, producing negative samples that have the same line spacing and line position but differ only in content. I reshuffle at every iteration, which improves data efficiency (data augmentation), greatly increasing the effective number of negative samples seen during training.

One thing to note: we consider horizontal adjacency, but when Python reads a matrix, it reads top-to-bottom rather than left-to-right, so we need to transpose the image matrix. In addition, I normalized the image — originally a grayscale image with values in the range 0–255 — into the range 0–1, which speeds up convergence. Finally, I subtracted the image matrix from 1, effectively performing a color inversion, turning "black text on white background" into "white text on black background." This is because during network training, we want the input to contain a large number of zeros, which speeds up convergence — but in terms of color, white is 255 and black is 0, so "white text on black background" converges faster than "black text on white background."

Training the Model

Next, we use these samples to train a CNN — a very standard process: stacking three convolution-plus-pooling layers, followed by a softmax classifier. That's all there is to it.

Model architecture:

______________________________________________________________
Layer (type) Output Shape Param # Connected to
==============================================================
input_2 (InputLayer) (None, 144, 180) 0
______________________________________________________________
convolution1d_4 (Convolution1D) (None, 143, 32) 11552 input_2[0][0]
______________________________________________________________
maxpooling1d_4 (MaxPooling1D) (None, 71, 32) 0 convolution1d_4[0][0]
______________________________________________________________
convolution1d_5 (Convolution1D) (None, 70, 32) 2080 maxpooling1d_4[0][0]
______________________________________________________________
maxpooling1d_5 (MaxPooling1D) (None, 35, 32) 0 convolution1d_5[0][0]
______________________________________________________________
convolution1d_6 (Convolution1D) (None, 34, 32) 2080 maxpooling1d_5[0][0]
______________________________________________________________
maxpooling1d_6 (MaxPooling1D) (None, 17, 32) 0 convolution1d_6[0][0]
______________________________________________________________
flatten_2 (Flatten) (None, 544) 0 maxpooling1d_6[0][0]
______________________________________________________________
dense_3 (Dense) (None, 32) 17440 flatten_2[0][0]
______________________________________________________________
dense_4 (Dense) (None, 1) 33 dense_3[0][0]
==============================================================
Total params: 33185
______________________________________________________________

Code:

from keras.layers import Input, Convolution1D, MaxPooling1D, Flatten, Dense
from keras.models import Model

input = Input((144, 180))
cnn = Convolution1D(32, 2)(input)
cnn = MaxPooling1D(2)(cnn)
cnn = Convolution1D(32, 2)(cnn)
cnn = MaxPooling1D(2)(cnn)
cnn = Convolution1D(32, 2)(cnn)
cnn = MaxPooling1D(2)(cnn)
cnn = Flatten()(cnn)
dense = Dense(32, activation='relu')(cnn)
dense = Dense(1, activation='sigmoid')(dense)

model = Model(input=input, output=dense)
model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])
model.summary()

model.fit_generator(data(sample, batch_size=1024), 
                    nb_epoch=100, 
                    samples_per_epoch=nb_samples*3
                   )
model.save_weights('2013_suizhifuyuan_cnn.model')

In fact, 3 epochs are enough to reach 95% accuracy; the nb_epoch=100 was just picked arbitrarily. If you have plenty of compute and aren't in a hurry, there's no harm in running more epochs. In the end I got up to 97.7% accuracy.

Reassembly Results

Now we can test the performance of our trained "distance function."

import glob

img_names = glob.glob(u'附件3/*')
images = {}
for i in img_names:
    images[i] = 1 - misc.imread(i, flatten=True).T/255.0

def find_most_similar(img, images):
    imgs_ = np.array([np.vstack((images[img], images[i])) for i in images if i != img])
    img_names_ = [i for i in images if i != img]
    sims = model.predict(imgs_).reshape(-1)
    return img_names_[sims.argmax()]

img = img_names[14]
result = [img]
images_ = images.copy()
while len(images_) > 1:
    print len(images_)
    img_ = find_most_similar(img, images_)
    result.append(img_)
    del images_[img]
    img = img_

images_ = [images[i].T for i in result]
compose = (1 - np.hstack(images_))*255
misc.imsave('result.png', compose)

For attachment 3, the result of a single-pass greedy reassembly is (click to view the full-size image):

Reconstruction quality for attachment 3 (CNN + greedy algorithm)

For comparison, here's the result of a single-pass greedy reassembly using the old Euclidean distance approach:

Reconstruction quality for attachment 3 (Euclidean distance + greedy algorithm)

The improvement is quite clear. Even though this model was trained on Chinese-language corpora, applying it directly to attachment 4 still gives decent results:

Reconstruction quality for attachment 4 (CNN + greedy algorithm)

Likewise, as a comparison, here's the reassembly result for attachment 4 using Euclidean distance:

Reconstruction quality for attachment 4 (Euclidean distance + greedy algorithm)

So we can see that the model we obtained really does perform well, and it generalizes remarkably well. That said, directly reassembling English text this way isn't great — it would be better to include English corpora in training as well, to get better results.

Closing Thoughts

A truly good problem is always rich in substance and never goes stale — this is already the third post I've written about the shredded-paper reassembly problem, and there may well be a fourth, a fifth... Each new investigation goes a little deeper...