Skip to main content
 首页 » 编程设计

python之如何使用 Spark 和 Caffe 对图像进行分类

2024年10月24日28lidabo

我正在使用Caffe做图像分类,我可以使用MAC OS X,Pyhton吗?

现在我知道如何使用 Caffe 和 Spark python 对图像列表进行分类,但如果我想让它更快,我想使用 Spark。

因此,我尝试将图像分类应用于 RDD 的每个元素,该 RDD 从 image_path 列表创建。但是,Spark 不允许我这样做。

这是我的代码:

这是图像分类的代码:

# display image name, class number, predicted label 
def classify_image(image_path, transformer, net): 
    image = caffe.io.load_image(image_path) 
    transformed_image = transformer.preprocess('data', image) 
    net.blobs['data'].data[...] = transformed_image 
    output = net.forward() 
    output_prob = output['prob'][0] 
    pred = output_prob.argmax() 
 
    labels_file = caffe_root + 'data/ilsvrc12/synset_words.txt' 
    labels = np.loadtxt(labels_file, str, delimiter='\t') 
    lb = labels[pred] 
 
    image_name = image_path.split(images_folder_path)[1] 
 
    result_str = 'image: '+image_name+'  prediction: '+str(pred)+'  label: '+lb 
    return result_str 

此代码生成 Caffe 参数并在 RDD 的每个元素上应用 classify_image 方法:

def main(): 
    sys.path.insert(0, caffe_root + 'python') 
    caffe.set_mode_cpu() 
    model_def = caffe_root + 'models/bvlc_reference_caffenet/deploy.prototxt' 
    model_weights = caffe_root + 'models/bvlc_reference_caffenet/bvlc_reference_caffenet.caffemodel' 
 
    net = caffe.Net(model_def, 
                model_weights, 
                caffe.TEST) 
 
    mu = np.load(caffe_root + 'python/caffe/imagenet/ilsvrc_2012_mean.npy') 
    mu = mu.mean(1).mean(1) 
 
    transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape}) 
 
    transformer.set_transpose('data', (2,0,1)) 
    transformer.set_mean('data', mu) 
    transformer.set_raw_scale('data', 255) 
    transformer.set_channel_swap('data', (2,1,0)) 
 
    net.blobs['data'].reshape(50, 
                          3, 
                          227, 227) 
 
    image_list= [] 
    for image_path in glob.glob(images_folder_path+'*.jpg'): 
        image_list.append(image_path) 
 
    images_rdd = sc.parallelize(image_list) 
    transformer_bc = sc.broadcast(transformer) 
    net_bc = sc.broadcast(net) 
    image_predictions = images_rdd.map(lambda image_path: classify_image(image_path, transformer_bc, net_bc)) 
    print image_predictions 
 
if __name__ == '__main__': 
    main() 

可以看到,这里我尝试广播caffe参数,transformer_bc = sc.broadcast(transformer), net_bc = sc.broadcast(net) 错误是:

RuntimeError: Pickling of "caffe._caffe.Net" instances is not enabled

在我进行广播之前,错误是:

Driver stacktrace.... Caused by: org.apache.spark.api.python.PythonException: Traceback (most recent call last):....

那么,您知道吗,有什么方法既可以使用 Caffe 和 Spark 对图像进行分类,又可以利用 Spark?

请您参考如下方法:

当您使用复杂的非本地对象时,初始化必须直接转移到工作人员,例如使用单例模块:

net_builder.py:

import cafe  
 
net = None 
 
def build_net(*args, **kwargs): 
     ...  # Initialize net here 
     return net        
 
def get_net(*args, **kwargs): 
    global net 
    if net is None: 
        net = build_net(*args, **kwargs) 
    return net 

main.py:

import net_builder 
 
sc.addPyFile("net_builder.py") 
 
def classify_image(image_path, transformer, *args, **kwargs): 
    net = net_builder.get_net(*args, **kwargs) 

这意味着您还必须分发所有必需的文件。它可以手动完成或 using SparkFiles机制。

附带说明一下,您应该查看 SparkNet包。