import t7 model

Author: mei jieru

README.md

@@ -3,3 +3,21 @@ Convolutional Recurrent Neural Network
 
 This software implements the Convolutional Recurrent Neural Network (CRNN) in pytorch.
 Origin software could be found in [crnn](https://github.com/bgshih/crnn)
+
+Run demo
+--------
+A demo program can be found in ``src/demo.py``. Before running the demo, download a pretrained model
+from [Baidu Netdisk](https://pan.baidu.com/s/1pLbeCND) or [Dropbox](https://www.dropbox.com/s/dboqjk20qjkpta3/crnn.pth?dl=0). 
+This pretrained model is converted from auther offered one by ``tool``.
+Put the downloaded model file ``crnn.pth`` into directory ``data/``. Then launch the demo by:
+
+    python demo.py
+
+The demo reads an example image and recognizes its text content.
+
+Example image:
+![Example Image](./data/demo.png)
+
+Expected output:
+    loading pretrained model from ./data/crnn.pth
+    a-----v--a-i-l-a-bb-l-ee-- => available

crnn_main.py

@@ -27,7 +27,7 @@
 parser.add_argument('--cuda', action='store_true', help='enables cuda')
 parser.add_argument('--ngpu', type=int, default=1, help='number of GPUs to use')
 parser.add_argument('--crnn', default='', help="path to crnn (to continue training)")
-parser.add_argument('--alphabet', type=str, default='abcdefghijklmnopqrstuvwxyz0123456789')
+parser.add_argument('--alphabet', type=str, default='0123456789abcdefghijklmnopqrstuvwxyz')
 parser.add_argument('--Diters', type=int, default=5, help='number of D iters per each G iter')
 parser.add_argument('--experiment', default=None, help='Where to store samples and models')
 parser.add_argument('--displayInterval', type=int, default=500, help='Interval to be displayed')
@@ -66,9 +66,9 @@
     train_dataset, batch_size=opt.batchSize,
     shuffle=True, sampler=sampler,
     num_workers=int(opt.workers),
-    collate_fn=dataset.alignCollate(imgH=opt.imgH,
-     keep_ratio=opt.keep_ratio))
-test_dataset = dataset.lmdbDataset(root=opt.valroot, transform=dataset.resizeNormalize((128, 32)))
+    collate_fn=dataset.alignCollate(imgH=opt.imgH, keep_ratio=opt.keep_ratio))
+test_dataset = dataset.lmdbDataset(
+    root=opt.valroot, transform=dataset.resizeNormalize((128, 32)))
 
 ngpu = int(opt.ngpu)
 nh = int(opt.nh)
@@ -114,11 +114,12 @@ def weights_init(m):
 
 # setup optimizer
 if opt.adam:
-    optimizer = optim.Adam(crnn.parameters(), lr=opt.lrD, betas=(opt.beta1, 0.999))
+    optimizer = optim.Adam(crnn.parameters(), lr=opt.lr,
+                           betas=(opt.beta1, 0.999))
 elif opt.adadelta:
-    optimizer = optim.Adadelta(crnn.parameters(), lr=opt.lrD)
+    optimizer = optim.Adadelta(crnn.parameters(), lr=opt.lr)
 else:
-    optimizer = optim.RMSprop(crnn.parameters(), lr=opt.lrD)
+    optimizer = optim.RMSprop(crnn.parameters(), lr=opt.lr)
 
 
 def val(net, dataset, criterion, max_iter=100):
@@ -129,7 +130,7 @@ def val(net, dataset, criterion, max_iter=100):
 
     net.eval()
     data_loader = torch.utils.data.DataLoader(
- dataset, shuffle=True, batch_size=opt.batchSize, num_workers=int(opt.workers))
+        dataset, shuffle=True, batch_size=opt.batchSize, num_workers=int(opt.workers))
     val_iter = iter(data_loader)
 
     i = 0
@@ -167,6 +168,9 @@ def val(net, dataset, criterion, max_iter=100):
     print('Test loss: %f, accuray: %f' % (loss_avg.val(), accuracy))
 
 
+#  val(crnn, test_dataset, criterion)
+#  exit(0)
+
 def trainBatch(net, criterion, optimizer):
     data = train_iter.next()
     cpu_images, cpu_texts = data
@@ -198,12 +202,14 @@ def trainBatch(net, criterion, optimizer):
         i += 1
 
         if i % opt.displayInterval == 0:
-            print('[%d/%d][%d/%d] Loss: %f' % (epoch, opt.niter, i, len(train_loader), loss_avg.val()))
+            print('[%d/%d][%d/%d] Loss: %f' %
+                  (epoch, opt.niter, i, len(train_loader), loss_avg.val()))
             loss_avg.reset()
 
         if i % opt.valInterval == 0:
             val(crnn, test_dataset, criterion)
 
         # do checkpointing
         if i % opt.saveInterval == 0:
-            torch.save(crnn.state_dict(), '{0}/netCRNN_{1}_{2}.pth'.format(opt.experiment, epoch, i))
+            torch.save(
+                crnn.state_dict(), '{0}/netCRNN_{1}_{2}.pth'.format(opt.experiment, epoch, i))

data/demo.png

@@ -0,0 +1,36 @@
+import torch
+from torch.autograd import Variable
+import utils
+import dataset
+from PIL import Image
+
+import models.crnn as crnn
+
+
+model_path = './data/crnn.pth'
+img_path = './data/demo.png'
+alphabet = '0123456789abcdefghijklmnopqrstuvwxyz'
+
+model = crnn.CRNN(32, 1, 37, 256, 1).cuda()
+print('loading pretrained model from %s' % model_path)
+model.load_state_dict(torch.load(model_path))
+
+converter = utils.strLabelConverter(alphabet)
+
+transformer = dataset.resizeNormalize((100, 32))
+image = Image.open(img_path).convert('L')
+image = transformer(image).cuda()
+image = image.view(1, *image.size())
+image = Variable(image)
+
+model.eval()
+preds = model(image)
+
+_, preds = preds.max(2)
+preds = preds.squeeze(2)
+preds = preds.transpose(1, 0).contiguous().view(-1)
+
+preds_size = Variable(torch.IntTensor([preds.size(0)]))
+raw_pred = converter.decode(preds.data, preds_size.data, raw=True)
+sim_pred = converter.decode(preds.data, preds_size.data, raw=False)
+print('%-20s => %-20s' % (raw_pred, sim_pred))

demo.py

@@ -2,6 +2,31 @@
 import torch.nn.parallel
 
 
+class BidirectionalLSTM(nn.Module):
+
+    def __init__(self, nIn, nHidden, nOut, ngpu):
+        super(BidirectionalLSTM, self).__init__()
+        self.ngpu = ngpu
+
+        self.rnn = nn.LSTM(nIn, nHidden, bidirectional=True)
+        self.embedding = nn.Linear(nHidden * 2, nOut)
+
+    def forward(self, input):
+        gpu_ids = None
+        if isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
+            gpu_ids = range(self.ngpu)
+        recurrent, _ = nn.parallel.data_parallel(
+            self.rnn, input, gpu_ids)  # [T, b, h * 2]
+
+        T, b, h = recurrent.size()
+        t_rec = recurrent.view(T * b, h)
+        output = nn.parallel.data_parallel(
+            self.embedding, t_rec, gpu_ids)  # [T * b, nOut]
+        output = output.view(T, b, -1)
+
+        return output
+
+
 class CRNN(nn.Module):
 
     def __init__(self, imgH, nc, nclass, nh, ngpu, n_rnn=2, leakyRelu=False):
@@ -30,45 +55,40 @@ def convRelu(i, batchNormalization=False):
                 cnn.add_module('relu{0}'.format(i), nn.ReLU(True))
 
         convRelu(0)
-        cnn.add_module('pooling{0}'.format(0), nn.MaxPool2d((2, 2),
-                                                            (2, 2)))  # 64x16x64
+        cnn.add_module('pooling{0}'.format(0), nn.MaxPool2d(2, 2))  # 64x16x64
         convRelu(1)
-        cnn.add_module('pooling{0}'.format(1), nn.MaxPool2d((2, 2),
-                                                            (2, 2)))  # 128x8x32
+        cnn.add_module('pooling{0}'.format(1), nn.MaxPool2d(2, 2))  # 128x8x32
         convRelu(2, True)
         convRelu(3)
         cnn.add_module('pooling{0}'.format(2), nn.MaxPool2d((2, 2),
-                                                            (2, 2)))  # 256x4x16
+                                                            (2, 1),
+                                                            (0, 1)))  # 256x4x16
         convRelu(4, True)
         convRelu(5)
         cnn.add_module('pooling{0}'.format(3), nn.MaxPool2d((2, 2),
-                                                            (2, 1)))  # 512x2x16
+                                                            (2, 1),
+                                                            (0, 1)))  # 512x2x16
         convRelu(6, True)  # 512x1x16
 
         self.cnn = cnn
-        self.rnn = nn.Sequential(nn.LSTM(512, nh, n_rnn, bidirectional=True))
-        self.text = nn.Sequential(nn.Linear(nh * 2, nclass))  # [T, b, nclass]
+        self.rnn = nn.Sequential(
+            BidirectionalLSTM(512, nh, nh, ngpu),
+            BidirectionalLSTM(nh, nh, nclass, ngpu)
+        )
 
     def forward(self, input):
         gpu_ids = None
         if isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
             gpu_ids = range(self.ngpu)
+
         # conv features
         conv = nn.parallel.data_parallel(self.cnn, input, gpu_ids)
-
-        # rnn features
         b, c, h, w = conv.size()
         assert h == 1, "the height of conv must be 1"
         conv = conv.squeeze(2)
         conv = conv.permute(2, 0, 1)  # [w, b, c]
-        recurrent, _ = nn.parallel.data_parallel(self.rnn, conv,
-                                                 gpu_ids)  # [T, b, h * 2]
 
-        # text classifier
-        T, b, h = recurrent.size()
-        t_rec = recurrent.view(T * b, h)
-        text = nn.parallel.data_parallel(self.text, t_rec,
-                                         gpu_ids)  # [T * b, nclass]
-        text = text.view(T, b, -1)
+        # rnn features
+        output = nn.parallel.data_parallel(self.rnn, conv, gpu_ids)
 
-        return text
+        return output

models/crnn.py

@@ -0,0 +1,89 @@
+require('table')
+require('torch')
+require('os')
+
+function clone(t)
+    -- deep-copy a table
+    if type(t) ~= "table" then return t end
+    local meta = getmetatable(t)
+    local target = {}
+    for k, v in pairs(t) do
+        if type(v) == "table" then
+            target[k] = clone(v)
+        else
+            target[k] = v
+        end
+    end
+    setmetatable(target, meta)
+    return target
+end
+
+
+function tableMerge(lhs, rhs)
+    output = clone(lhs)
+    for _, v in pairs(rhs) do
+        table.insert(output, v)
+    end
+    return output
+end
+
+
+function isInTable(val, val_list)
+    for _, item in pairs(val_list) do
+        if val == item then
+            return true
+        end
+    end
+    return false
+end
+
+
+function modelToList(model)
+    local ignoreList = {
+        'nn.Copy',
+        'nn.AddConstant',
+        'nn.MulConstant',
+        'nn.View',
+        'nn.Transpose',
+        'nn.SplitTable',
+        'nn.SharedParallelTable',
+        'nn.JoinTable',
+    }
+    local state = {}
+    local param
+    for i, layer in pairs(model.modules) do
+        local typeName = torch.type(layer)
+        if not isInTable(typeName, ignoreList) then
+            if typeName == 'nn.Sequential' or typeName == 'nn.ConcatTable' then
+                param = modelToList(layer)
+            elseif typeName == 'cudnn.SpatialConvolution' or typeName == 'nn.SpatialConvolution' then
+                param = layer:parameters()
+            elseif typeName == 'cudnn.SpatialBatchNormalization' or typeName == 'nn.SpatialBatchNormalization' then
+                param = layer:parameters()
+                bn_vars = {layer.running_mean, layer.running_var}
+                param = tableMerge(param, bn_vars)
+            elseif typeName == 'nn.LstmLayer' then
+                param =  layer:parameters()
+            elseif typeName == 'nn.BiRnnJoin' then
+                param =  layer:parameters()
+            elseif typeName == 'cudnn.SpatialMaxPooling' or typeName == 'nn.SpatialMaxPooling' then
+                param = {}
+            elseif typeName == 'cudnn.ReLU' or typeName == 'nn.ReLU' then
+                param = {}
+            else
+                print(string.format('Unknown class %s', typeName))
+                os.exit(0)
+            end
+            table.insert(state, {typeName, param})
+        else
+            print(string.format('pass %s', typeName))
+        end
+    end
+    return state
+end
+
+
+function saveModel(model, output_path)
+    local state =  modelToList(model)
+    torch.save(output_path, state)
+end