yahoo / open_nsfw

Not Suitable for Work (NSFW) classification using deep neural network Caffe models.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

there is a RuntimeError: Pickling of "caffe._caffe.Net" instances is not enabled (http://www.boost.org/libs/python/doc/v2/pickle.html) when I try run open_nsfw in multi-process

kylezhang1983 opened this issue · comments

when I try run open_nsfw in multi-process , a runtimeError raised:
Traceback (most recent call last):
File "./verify_concurrency.py", line 152, in
main(sys.argv)
File "./verify_concurrency.py", line 149, in main
print res.get()
File "/usr/lib/python2.7/multiprocessing/pool.py", line 567, in get
raise self._value
RuntimeError: Pickling of "caffe._caffe.Net" instances is not enabled (http://www.boost.org/libs/python/doc/v2/pickle.html)

could some expert help check where is wrong as i'm rookie
my command is :
./verify_concurrency.py --model_def ./nsfw_model/deploy.prototxt --pretrained_model ./nsfw_model/resnet_50_1by2_nsfw.caffemodel /home/harold/nsfw/pic/1

below is my script verify_concurrency.py :
#!/usr/bin/env python
"""
Copyright 2016 Yahoo Inc.
Licensed under the terms of the 2 clause BSD license.
Please see LICENSE file in the project root for terms.
"""

import numpy as np
import os
import sys
import argparse
import glob
import time
from PIL import Image
from StringIO import StringIO
import caffe
from multiprocessing import Pool

def resize_image(data, sz=(256, 256)):
"""
Resize image. Please use this resize logic for best results instead of the
caffe, since it was used to generate training dataset
:param str data:
The image data
:param sz tuple:
The resized image dimensions
:returns bytearray:
A byte array with the resized image
"""
img_data = str(data)
im = Image.open(StringIO(img_data))
if im.mode != "RGB":
im = im.convert('RGB')
imr = im.resize(sz, resample=Image.BILINEAR)
fh_im = StringIO()
imr.save(fh_im, format='JPEG')
fh_im.seek(0)
return bytearray(fh_im.read())

def caffe_preprocess_and_compute(file_name, caffe_transformer=None, caffe_net=None,
output_layers=None):
"""
Run a Caffe network on an input image after preprocessing it to prepare
it for Caffe.
:param file name file_name:
PIL image to be input into Caffe.
:param caffe.Net caffe_net:
A Caffe network with which to process pimg afrer preprocessing.
:param list output_layers:
A list of the names of the layers from caffe_net whose outputs are to
to be returned. If this is None, the default outputs for the network
are returned.
:return:
Returns the requested outputs from the Caffe net.
"""
if caffe_net is not None:

    # Grab the default output names if none were requested specifically.
    if output_layers is None:
        output_layers = caffe_net.outputs
    pimg = open(file_name).read()
    start_time = time.time()
    img_data_rs = resize_image(pimg, sz=(256, 256))
    image = caffe.io.load_image(StringIO(img_data_rs))

    H, W, _ = image.shape
    _, _, h, w = caffe_net.blobs['data'].data.shape
    h_off = max((H - h) / 2, 0)
    w_off = max((W - w) / 2, 0)
    crop = image[h_off:h_off + h, w_off:w_off + w, :]
    transformed_image = caffe_transformer.preprocess('data', crop)
    transformed_image.shape = (1,) + transformed_image.shape

    input_name = caffe_net.inputs[0]
    all_outputs = caffe_net.forward_all(blobs=output_layers,
                **{input_name: transformed_image})

    outputs = all_outputs[output_layers[0]][0].astype(float)
    
    return outputs
else:
    return []

end_time = time.time()
print("deal image:%s;start time:%.3f;end time:%.3f;score:%.3f;cost time:%.3f" % (file_name,start_time, end_time , outputs[1] , (end_time - start_time)) )

def main(argv):
pycaffe_dir = os.path.dirname(file)

parser = argparse.ArgumentParser()
# Required arguments: input file.
parser.add_argument(
    "input_dir",
    help="Directory of the input image file"
)

# Optional arguments.
parser.add_argument(
    "--model_def",
    help="Model definition file."
)
parser.add_argument(
    "--pretrained_model",
    help="Trained model weights file."
)

args = parser.parse_args()
# Pre-load caffe model.
nsfw_net = caffe.Net(args.model_def,  # pylint: disable=invalid-name
    args.pretrained_model, caffe.TEST)

# Load transformer
# Note that the parameters are hard-coded for best results
caffe_transformer = caffe.io.Transformer({'data': nsfw_net.blobs['data'].data.shape})
caffe_transformer.set_transpose('data', (2, 0, 1))  # move image channels to outermost
caffe_transformer.set_mean('data', np.array([104, 117, 123]))  # subtract the dataset-mean value in each channel
caffe_transformer.set_raw_scale('data', 255)  # rescale from [0, 1] to [0, 255]
caffe_transformer.set_channel_swap('data', (2, 1, 0))  # swap channels from RGB to BGR

start_time  = time.time()
print("The process starttime: %.3f" % start_time) 
file_list = os.listdir(args.input_dir)
multiprocess_pool = Pool(10)
result = []
for file in file_list:

    # Classify.

    p = args.input_dir + '/' + file
    #print("The image %s deal startat %.3f" % (p, time.time()) )
    result.append(multiprocess_pool.apply_async(caffe_preprocess_and_compute, args=(p, caffe_transformer, nsfw_net, ['prob'] ,)))
    #scores = caffe_preprocess_and_compute(p, caffe_transformer=caffe_transformer, caffe_net=nsfw_net, output_layers=['prob'])

    # Scores is the array containing SFW / NSFW image probabilities
    # scores[1] indicates the NSFW probability
    #print "NSFW score:  " , scores[1]
    
    #print("The image %s got score: %.3f at %.3f" % (p, scores[1] ,time.time()) )

print 'Waiting for all subprocesses done...'
multiprocess_pool.close()
multiprocess_pool.join()
print 'All subprocesses done.'
end_time  = time.time()  
print("The process endtime: %.3f" % end_time) 

print("The process cost time is %.3f for %d images" % ( (end_time - start_time), len(file_list)) )
for res in result:
    print res.get()    

if name == 'main':
main(sys.argv)

And If I don't check the subprocess result:

change line 131:
result.append(multiprocess_pool.apply_async(caffe_preprocess_and_compute, args=(p, caffe_transformer, nsfw_net, ['prob'] ,))) ---->multiprocess_pool.apply_async(caffe_preprocess_and_compute, args=(p, caffe_transformer, nsfw_net, ['prob'] ,))

it's ok

this issue caused by bad usage for multi-processing poll

@kylezhang1983 Hi, how did you solve this issue? I have the same problem as you did.

@taosean , look Caffee will solve it in Caffee2. So I change the framwork caffee to tensorflow.