FilesExpand file tree

 master

/

classifier.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

98 lines (84 loc) · 3.45 KB

 master

/

classifier.py

Copy path

Top

File metadata and controls

• Code
• Blame

98 lines (84 loc) · 3.45 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

#!/usr/bin/env python

"""

Classifier is an image classifier specialization of Net.

"""

import numpy as np

import caffe

class Classifier(caffe.Net):

"""

Classifier extends Net for image class prediction

by scaling, center cropping, or oversampling.

Parameters

----------

image_dims : dimensions to scale input for cropping/sampling.

Default is to scale to net input size for whole-image crop.

mean, input_scale, raw_scale, channel_swap: params for

preprocessing options.

"""

def __init__(self, model_file, pretrained_file, image_dims=None,

mean=None, input_scale=None, raw_scale=None,

channel_swap=None):

caffe.Net.__init__(self, model_file, pretrained_file, caffe.TEST)

# configure pre-processing

in_ = self.inputs[0]

self.transformer = caffe.io.Transformer(

{in_: self.blobs[in_].data.shape})

self.transformer.set_transpose(in_, (2, 0, 1))

if mean is not None:

self.transformer.set_mean(in_, mean)

if input_scale is not None:

self.transformer.set_input_scale(in_, input_scale)

if raw_scale is not None:

self.transformer.set_raw_scale(in_, raw_scale)

if channel_swap is not None:

self.transformer.set_channel_swap(in_, channel_swap)

self.crop_dims = np.array(self.blobs[in_].data.shape[2:])

if not image_dims:

image_dims = self.crop_dims

self.image_dims = image_dims

def predict(self, inputs, oversample=True):

"""

Predict classification probabilities of inputs.

Parameters

----------

inputs : iterable of (H x W x K) input ndarrays.

oversample : boolean

average predictions across center, corners, and mirrors

when True (default). Center-only prediction when False.

Returns

-------

predictions: (N x C) ndarray of class probabilities for N images and C

classes.

"""

# Scale to standardize input dimensions.

input_ = np.zeros((len(inputs),

self.image_dims[0],

self.image_dims[1],

inputs[0].shape[2]),

dtype=np.float32)

for ix, in_ in enumerate(inputs):

input_[ix] = caffe.io.resize_image(in_, self.image_dims)

if oversample:

# Generate center, corner, and mirrored crops.

input_ = caffe.io.oversample(input_, self.crop_dims)

else:

# Take center crop.

center = np.array(self.image_dims) / 2.0

crop = np.tile(center, (1, 2))[0] + np.concatenate([

-self.crop_dims / 2.0,

self.crop_dims / 2.0

])

crop = crop.astype(int)

input_ = input_[:, crop[0]:crop[2], crop[1]:crop[3], :]

# Classify

caffe_in = np.zeros(np.array(input_.shape)[[0, 3, 1, 2]],

dtype=np.float32)

for ix, in_ in enumerate(input_):

caffe_in[ix] = self.transformer.preprocess(self.inputs[0], in_)

out = self.forward_all(**{self.inputs[0]: caffe_in})

predictions = out[self.outputs[0]]

# For oversampling, average predictions across crops.

if oversample:

predictions = predictions.reshape((len(predictions) / 10, 10, -1))

predictions = predictions.mean(1)

return predictions