Diff of /darkflow/utils/process.py [000000] .. [d34869]

Switch to side-by-side view

--- a
+++ b/darkflow/utils/process.py
@@ -0,0 +1,334 @@
+"""
+WARNING: spaghetti code.
+"""
+
+import numpy as np
+import pickle
+import os
+
+
+def parser(model):
+    """
+    Read the .cfg file to extract layers into `layers`
+    as well as model-specific parameters into `meta`
+    """
+
+    def _parse(l, i=1):
+        return l.split('=')[i].strip()
+
+    with open(model, 'rb') as f:
+        lines = f.readlines()
+
+    lines = [line.decode() for line in lines]
+
+    meta = dict();
+    layers = list()  # will contains layers' info
+    h, w, c = [int()] * 3;
+    layer = dict()
+    for line in lines:
+        line = line.strip()
+        line = line.split('#')[0]
+        if '[' in line:
+            if layer != dict():
+                if layer['type'] == '[net]':
+                    h = layer['height']
+                    w = layer['width']
+                    c = layer['channels']
+                    meta['net'] = layer
+                else:
+                    if layer['type'] == '[crop]':
+                        h = layer['crop_height']
+                        w = layer['crop_width']
+                    layers += [layer]
+            layer = {'type': line}
+        else:
+            try:
+                i = float(_parse(line))
+                if i == int(i): i = int(i)
+                layer[line.split('=')[0].strip()] = i
+            except:
+                try:
+                    key = _parse(line, 0)
+                    val = _parse(line, 1)
+                    layer[key] = val
+                except:
+                    'banana ninja yadayada'
+
+    meta.update(layer)  # last layer contains meta info
+    if 'anchors' in meta:
+        splits = meta['anchors'].split(',')
+        anchors = [float(x.strip()) for x in splits]
+        meta['anchors'] = anchors
+    meta['model'] = model  # path to cfg, not model name
+    meta['inp_size'] = [h, w, c]
+    return layers, meta
+
+
+def cfg_yielder(model, binary):
+    """
+    yielding each layer information to initialize `layer`
+    """
+    layers, meta = parser(model);
+    yield meta;
+    h, w, c = meta['inp_size'];
+    l = w * h * c
+
+    # Start yielding
+    flat = False  # flag for 1st dense layer
+    conv = '.conv.' in model
+    for i, d in enumerate(layers):
+        # -----------------------------------------------------
+        if d['type'] == '[crop]':
+            yield ['crop', i]
+        # -----------------------------------------------------
+        elif d['type'] == '[local]':
+            n = d.get('filters', 1)
+            size = d.get('size', 1)
+            stride = d.get('stride', 1)
+            pad = d.get('pad', 0)
+            activation = d.get('activation', 'logistic')
+            w_ = (w - 1 - (1 - pad) * (size - 1)) // stride + 1
+            h_ = (h - 1 - (1 - pad) * (size - 1)) // stride + 1
+            yield ['local', i, size, c, n, stride,
+                   pad, w_, h_, activation]
+            if activation != 'linear': yield [activation, i]
+            w, h, c = w_, h_, n
+            l = w * h * c
+        # -----------------------------------------------------
+        elif d['type'] == '[convolutional]':
+            n = d.get('filters', 1)
+            size = d.get('size', 1)
+            stride = d.get('stride', 1)
+            pad = d.get('pad', 0)
+            padding = d.get('padding', 0)
+            if pad: padding = size // 2
+            activation = d.get('activation', 'logistic')
+            batch_norm = d.get('batch_normalize', 0) or conv
+            yield ['convolutional', i, size, c, n,
+                   stride, padding, batch_norm,
+                   activation]
+            if activation != 'linear': yield [activation, i]
+            w_ = (w + 2 * padding - size) // stride + 1
+            h_ = (h + 2 * padding - size) // stride + 1
+            w, h, c = w_, h_, n
+            l = w * h * c
+        # -----------------------------------------------------
+        elif d['type'] == '[maxpool]':
+            stride = d.get('stride', 1)
+            size = d.get('size', stride)
+            padding = d.get('padding', (size - 1) // 2)
+            yield ['maxpool', i, size, stride, padding]
+            w_ = (w + 2 * padding) // d['stride']
+            h_ = (h + 2 * padding) // d['stride']
+            w, h = w_, h_
+            l = w * h * c
+        # -----------------------------------------------------
+        elif d['type'] == '[avgpool]':
+            flat = True;
+            l = c
+            yield ['avgpool', i]
+        # -----------------------------------------------------
+        elif d['type'] == '[softmax]':
+            yield ['softmax', i, d['groups']]
+        # -----------------------------------------------------
+        elif d['type'] == '[connected]':
+            if not flat:
+                yield ['flatten', i]
+                flat = True
+            activation = d.get('activation', 'logistic')
+            yield ['connected', i, l, d['output'], activation]
+            if activation != 'linear': yield [activation, i]
+            l = d['output']
+        # -----------------------------------------------------
+        elif d['type'] == '[dropout]':
+            yield ['dropout', i, d['probability']]
+        # -----------------------------------------------------
+        elif d['type'] == '[select]':
+            if not flat:
+                yield ['flatten', i]
+                flat = True
+            inp = d.get('input', None)
+            if type(inp) is str:
+                file = inp.split(',')[0]
+                layer_num = int(inp.split(',')[1])
+                with open(file, 'rb') as f:
+                    profiles = pickle.load(f, encoding='latin1')[0]
+                layer = profiles[layer_num]
+            else:
+                layer = inp
+            activation = d.get('activation', 'logistic')
+            d['keep'] = d['keep'].split('/')
+            classes = int(d['keep'][-1])
+            keep = [int(c) for c in d['keep'][0].split(',')]
+            keep_n = len(keep)
+            train_from = classes * d['bins']
+            for count in range(d['bins'] - 1):
+                for num in keep[-keep_n:]:
+                    keep += [num + classes]
+            k = 1
+            while layers[i - k]['type'] not in ['[connected]', '[extract]']:
+                k += 1
+                if i - k < 0:
+                    break
+            if i - k < 0:
+                l_ = l
+            elif layers[i - k]['type'] == 'connected':
+                l_ = layers[i - k]['output']
+            else:
+                l_ = layers[i - k].get('old', [l])[-1]
+            yield ['select', i, l_, d['old_output'],
+                   activation, layer, d['output'],
+                   keep, train_from]
+            if activation != 'linear': yield [activation, i]
+            l = d['output']
+        # -----------------------------------------------------
+        elif d['type'] == '[conv-select]':
+            n = d.get('filters', 1)
+            size = d.get('size', 1)
+            stride = d.get('stride', 1)
+            pad = d.get('pad', 0)
+            padding = d.get('padding', 0)
+            if pad: padding = size // 2
+            activation = d.get('activation', 'logistic')
+            batch_norm = d.get('batch_normalize', 0) or conv
+            d['keep'] = d['keep'].split('/')
+            classes = int(d['keep'][-1])
+            keep = [int(x) for x in d['keep'][0].split(',')]
+
+            segment = classes + 5
+            assert n % segment == 0, \
+                'conv-select: segment failed'
+            bins = n // segment
+            keep_idx = list()
+            for j in range(bins):
+                offset = j * segment
+                for k in range(5):
+                    keep_idx += [offset + k]
+                for k in keep:
+                    keep_idx += [offset + 5 + k]
+            w_ = (w + 2 * padding - size) // stride + 1
+            h_ = (h + 2 * padding - size) // stride + 1
+            c_ = len(keep_idx)
+            yield ['conv-select', i, size, c, n,
+                   stride, padding, batch_norm,
+                   activation, keep_idx, c_]
+            w, h, c = w_, h_, c_
+            l = w * h * c
+        # -----------------------------------------------------
+        elif d['type'] == '[conv-extract]':
+            file = d['profile']
+            with open(file, 'rb') as f:
+                profiles = pickle.load(f, encoding='latin1')[0]
+            inp_layer = None
+            inp = d['input']
+            out = d['output']
+            inp_layer = None
+            if inp >= 0:
+                inp_layer = profiles[inp]
+            if inp_layer is not None:
+                assert len(inp_layer) == c, \
+                    'Conv-extract does not match input dimension'
+            out_layer = profiles[out]
+
+            n = d.get('filters', 1)
+            size = d.get('size', 1)
+            stride = d.get('stride', 1)
+            pad = d.get('pad', 0)
+            padding = d.get('padding', 0)
+            if pad: padding = size // 2
+            activation = d.get('activation', 'logistic')
+            batch_norm = d.get('batch_normalize', 0) or conv
+
+            k = 1
+            find = ['[convolutional]', '[conv-extract]']
+            while layers[i - k]['type'] not in find:
+                k += 1
+                if i - k < 0: break
+            if i - k >= 0:
+                previous_layer = layers[i - k]
+                c_ = previous_layer['filters']
+            else:
+                c_ = c
+
+            yield ['conv-extract', i, size, c_, n,
+                   stride, padding, batch_norm,
+                   activation, inp_layer, out_layer]
+            if activation != 'linear': yield [activation, i]
+            w_ = (w + 2 * padding - size) // stride + 1
+            h_ = (h + 2 * padding - size) // stride + 1
+            w, h, c = w_, h_, len(out_layer)
+            l = w * h * c
+        # -----------------------------------------------------
+        elif d['type'] == '[extract]':
+            if not flat:
+                yield ['flatten', i]
+                flat = True
+            activation = d.get('activation', 'logistic')
+            file = d['profile']
+            with open(file, 'rb') as f:
+                profiles = pickle.load(f, encoding='latin1')[0]
+            inp_layer = None
+            inp = d['input']
+            out = d['output']
+            if inp >= 0:
+                inp_layer = profiles[inp]
+            out_layer = profiles[out]
+            old = d['old']
+            old = [int(x) for x in old.split(',')]
+            if inp_layer is not None:
+                if len(old) > 2:
+                    h_, w_, c_, n_ = old
+                    new_inp = list()
+                    for p in range(c_):
+                        for q in range(h_):
+                            for r in range(w_):
+                                if p not in inp_layer:
+                                    continue
+                                new_inp += [r + w * (q + h * p)]
+                    inp_layer = new_inp
+                    old = [h_ * w_ * c_, n_]
+                assert len(inp_layer) == l, \
+                    'Extract does not match input dimension'
+            d['old'] = old
+            yield ['extract', i] + old + [activation] + [inp_layer, out_layer]
+            if activation != 'linear': yield [activation, i]
+            l = len(out_layer)
+        # -----------------------------------------------------
+        elif d['type'] == '[route]':  # add new layer here
+            routes = d['layers']
+            if type(routes) is int:
+                routes = [routes]
+            else:
+                routes = [int(x.strip()) for x in routes.split(',')]
+            routes = [i + x if x < 0 else x for x in routes]
+            for j, x in enumerate(routes):
+                lx = layers[x];
+                xtype = lx['type']
+                _size = lx['_size'][:3]
+                if j == 0:
+                    h, w, c = _size
+                else:
+                    h_, w_, c_ = _size
+                    assert w_ == w and h_ == h, \
+                        'Routing incompatible conv sizes'
+                    c += c_
+            yield ['route', i, routes]
+            l = w * h * c
+        # -----------------------------------------------------
+        elif d['type'] == '[reorg]':
+            stride = d.get('stride', 1)
+            yield ['reorg', i, stride]
+            w = w // stride;
+            h = h // stride;
+            c = c * (stride ** 2)
+            l = w * h * c
+        # -----------------------------------------------------
+        else:
+            exit('Layer {} not implemented'.format(d['type']))
+
+        d['_size'] = list([h, w, c, l, flat])
+
+    if not flat:
+        meta['out_size'] = [h, w, c]
+    else:
+        meta['out_size'] = l