caffe

1
#!/usr/bin/env python
2
from mincepie import mapreducer, launcher
3
import gflags
4
import os
5
import cv2
6
from PIL import Image
7

8
# gflags
9
gflags.DEFINE_string('image_lib', 'opencv',
10
                     'OpenCV or PIL, case insensitive. The default value is the faster OpenCV.')
11
gflags.DEFINE_string('input_folder', '',
12
                     'The folder that contains all input images, organized in synsets.')
13
gflags.DEFINE_integer('output_side_length', 256,
14
                     'Expected side length of the output image.')
15
gflags.DEFINE_string('output_folder', '',
16
                     'The folder that we write output resized and cropped images to')
17
FLAGS = gflags.FLAGS
18

19
class OpenCVResizeCrop:
20
    def resize_and_crop_image(self, input_file, output_file, output_side_length = 256):
21
        '''Takes an image name, resize it and crop the center square
22
        '''
23
        img = cv2.imread(input_file)
24
        height, width, depth = img.shape
25
        new_height = output_side_length
26
        new_width = output_side_length
27
        if height > width:
28
            new_height = output_side_length * height / width
29
        else:
30
            new_width = output_side_length * width / height
31
        resized_img = cv2.resize(img, (new_width, new_height))
32
        height_offset = (new_height - output_side_length) / 2
33
        width_offset = (new_width - output_side_length) / 2
34
        cropped_img = resized_img[height_offset:height_offset + output_side_length,
35
                                  width_offset:width_offset + output_side_length]
36
        cv2.imwrite(output_file, cropped_img)
37

38
class PILResizeCrop:
39
## http://united-coders.com/christian-harms/image-resizing-tips-every-coder-should-know/
40
    def resize_and_crop_image(self, input_file, output_file, output_side_length = 256, fit = True):
41
        '''Downsample the image.
42
        '''
43
        img = Image.open(input_file)
44
        box = (output_side_length, output_side_length)
45
        #preresize image with factor 2, 4, 8 and fast algorithm
46
        factor = 1
47
        while img.size[0]/factor > 2*box[0] and img.size[1]*2/factor > 2*box[1]:
48
            factor *=2
49
        if factor > 1:
50
            img.thumbnail((img.size[0]/factor, img.size[1]/factor), Image.NEAREST)
51

52
        #calculate the cropping box and get the cropped part
53
        if fit:
54
            x1 = y1 = 0
55
            x2, y2 = img.size
56
            wRatio = 1.0 * x2/box[0]
57
            hRatio = 1.0 * y2/box[1]
58
            if hRatio > wRatio:
59
                y1 = int(y2/2-box[1]*wRatio/2)
60
                y2 = int(y2/2+box[1]*wRatio/2)
61
            else:
62
                x1 = int(x2/2-box[0]*hRatio/2)
63
                x2 = int(x2/2+box[0]*hRatio/2)
64
            img = img.crop((x1,y1,x2,y2))
65

66
        #Resize the image with best quality algorithm ANTI-ALIAS
67
        img.thumbnail(box, Image.ANTIALIAS)
68

69
        #save it into a file-like object
70
        with open(output_file, 'wb') as out:
71
            img.save(out, 'JPEG', quality=75)
72

73
class ResizeCropImagesMapper(mapreducer.BasicMapper):
74
    '''The ImageNet Compute mapper. 
75
    The input value would be the file listing images' paths relative to input_folder.
76
    '''
77
    def map(self, key, value):
78
        if type(value) is not str:
79
            value = str(value)
80
        files = [value]
81
        image_lib = FLAGS.image_lib.lower()
82
        if image_lib == 'pil':
83
            resize_crop = PILResizeCrop()
84
        else:
85
            resize_crop = OpenCVResizeCrop()
86
        for i, line in enumerate(files):
87
            try:
88
                line = line.replace(FLAGS.input_folder, '').strip()
89
                line = line.split()
90
                image_file_name = line[0]
91
                input_file = os.path.join(FLAGS.input_folder, image_file_name)
92
                output_file = os.path.join(FLAGS.output_folder, image_file_name)
93
                output_dir = output_file[:output_file.rfind('/')]
94
                if not os.path.exists(output_dir):
95
                    os.makedirs(output_dir)
96
                feat = resize_crop.resize_and_crop_image(input_file, output_file,
97
                                                              FLAGS.output_side_length)
98
            except Exception, e:
99
                # we ignore the exception (maybe the image is corrupted?)
100
                print line, Exception, e
101
        yield value, FLAGS.output_folder
102

103
mapreducer.REGISTER_DEFAULT_MAPPER(ResizeCropImagesMapper)
104
mapreducer.REGISTER_DEFAULT_REDUCER(mapreducer.NoPassReducer)
105
mapreducer.REGISTER_DEFAULT_READER(mapreducer.FileReader)
106
mapreducer.REGISTER_DEFAULT_WRITER(mapreducer.FileWriter)
107
 
108
if __name__ == '__main__':
109
    launcher.launch()
110