caffe

1
import numpy as np
2

3

4
class SimpleTransformer:
5

6
    """
7
    SimpleTransformer is a simple class for preprocessing and deprocessing
8
    images for caffe.
9
    """
10

11
    def __init__(self, mean=[128, 128, 128]):
12
        self.mean = np.array(mean, dtype=np.float32)
13
        self.scale = 1.0
14

15
    def set_mean(self, mean):
16
        """
17
        Set the mean to subtract for centering the data.
18
        """
19
        self.mean = mean
20

21
    def set_scale(self, scale):
22
        """
23
        Set the data scaling.
24
        """
25
        self.scale = scale
26

27
    def preprocess(self, im):
28
        """
29
        preprocess() emulate the pre-processing occurring in the vgg16 caffe
30
        prototxt.
31
        """
32

33
        im = np.float32(im)
34
        im = im[:, :, ::-1]  # change to BGR
35
        im -= self.mean
36
        im *= self.scale
37
        im = im.transpose((2, 0, 1))
38

39
        return im
40

41
    def deprocess(self, im):
42
        """
43
        inverse of preprocess()
44
        """
45
        im = im.transpose(1, 2, 0)
46
        im /= self.scale
47
        im += self.mean
48
        im = im[:, :, ::-1]  # change to RGB
49

50
        return np.uint8(im)
51

52

53
class CaffeSolver:
54

55
    """
56
    Caffesolver is a class for creating a solver.prototxt file. It sets default
57
    values and can export a solver parameter file.
58
    Note that all parameters are stored as strings. Strings variables are
59
    stored as strings in strings.
60
    """
61

62
    def __init__(self, testnet_prototxt_path="testnet.prototxt",
63
                 trainnet_prototxt_path="trainnet.prototxt", debug=False):
64

65
        self.sp = {}
66

67
        # critical:
68
        self.sp['base_lr'] = '0.001'
69
        self.sp['momentum'] = '0.9'
70

71
        # speed:
72
        self.sp['test_iter'] = '100'
73
        self.sp['test_interval'] = '250'
74

75
        # looks:
76
        self.sp['display'] = '25'
77
        self.sp['snapshot'] = '2500'
78
        self.sp['snapshot_prefix'] = '"snapshot"'  # string within a string!
79

80
        # learning rate policy
81
        self.sp['lr_policy'] = '"fixed"'
82

83
        # important, but rare:
84
        self.sp['gamma'] = '0.1'
85
        self.sp['weight_decay'] = '0.0005'
86
        self.sp['train_net'] = '"' + trainnet_prototxt_path + '"'
87
        self.sp['test_net'] = '"' + testnet_prototxt_path + '"'
88

89
        # pretty much never change these.
90
        self.sp['max_iter'] = '100000'
91
        self.sp['test_initialization'] = 'false'
92
        self.sp['average_loss'] = '25'  # this has to do with the display.
93
        self.sp['iter_size'] = '1'  # this is for accumulating gradients
94

95
        if (debug):
96
            self.sp['max_iter'] = '12'
97
            self.sp['test_iter'] = '1'
98
            self.sp['test_interval'] = '4'
99
            self.sp['display'] = '1'
100

101
    def add_from_file(self, filepath):
102
        """
103
        Reads a caffe solver prototxt file and updates the Caffesolver
104
        instance parameters.
105
        """
106
        with open(filepath, 'r') as f:
107
            for line in f:
108
                if line[0] == '#':
109
                    continue
110
                splitLine = line.split(':')
111
                self.sp[splitLine[0].strip()] = splitLine[1].strip()
112

113
    def write(self, filepath):
114
        """
115
        Export solver parameters to INPUT "filepath". Sorted alphabetically.
116
        """
117
        f = open(filepath, 'w')
118
        for key, value in sorted(self.sp.items()):
119
            if not(type(value) is str):
120
                raise TypeError('All solver parameters must be strings')
121
            f.write('%s: %s\n' % (key, value))
122