pycaffe 的使用

文章目录

    • 一、导入 caffe、设置设备以及加载网络
    • 二、分类预处理及预测
      • 1、caffe 接口实现
      • 2、自己实现
    • 三、查看 blobs 形状、params 形状及参数
    • 四、参考资料


一、导入 caffe、设置设备以及加载网络

import numpy as np
import os
import cv2
import sys
reload(sys)
sys.setdefaultencoding('utf-8')

# caffe setting
sys.path.insert(0, '/home/manzp/caffe_ssd/caffe/python')
import caffe

# run decice setting
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
caffe.set_mode_gpu()  # caffe.set_mode_cpu()

# initialize model and use test mode(eg: don't perform dropout)
net_plate = caffe.Net('model/res18_deploy.prototxt', 'model/res18_iter_20600.caffemodel', caffe.TEST)

二、分类预处理及预测

1、caffe 接口实现

######### 1、caffe 接口实现,注意操作的前后顺序 #########
transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})  # create transformer for the input called 'data' 
transformer.set_transpose('data', (2,0,1))  # H,W,C-->C,H,W
transformer.set_channel_swap('data', (2,1,0))  # RGB-->BGR
transformer.set_raw_scale('data', 255)  # rescale from [0, 1] to [0, 255]
transformer.set_mean('data', np.load(mean_file).mean(1).mean(1))  # subtract the dataset-mean value in each channel
transformer.set_input_scale('data', 1/128.0)
img = caffe.io.load_image(img_path,color=True)  # 默认读入彩图,需要灰度图的话改成 color=False
net.blobs['data'].data[...] = transformer.preprocess('data', img) # 执行预处理操作,并将处理后的图片载入到 data blob 中

Note:
1、caffe io 加载的图片类型为 RGB、通道格式为 H,W,C、像素范围是 [0-1]
2、而模型需要的图片类型为 BGR、通道格式为 C,H,W、像素范围为 [0-255],所以需要上面的转换


######### 2、caffe 封装的接口实现,默认 H,W,C-->C,H,W,且会 resize 到 img_dims #########
net_plate = caffe.Classifier(model_file='model/res18_deploy.prototxt',
                             pretrained_file='model/res18_iter_20600.caffemodel',
                             image_dims=(256, 256),
                             mean=np.array([127.5, 127.5, 127.5]),
                             input_scale=1/128.0,
                             raw_scale=255,
                             channel_swap=(2, 1, 0))

# inputs: iterable of (H x W x C) input ndarrays, can inference multiple imgs one-time,eg: inputs=[img, img, img]
# oversampl: Set the flag of ten-crop test
prob = net.predict(inputs=[img], oversample=False)[0]  # first image in the batch
index = prob.argmax()

2、自己实现

# -*- coding:utf-8 -*-

import numpy as np
import os
import cv2
import sys
reload(sys)
sys.setdefaultencoding('utf-8')

# caffe setting
sys.path.insert(0, '/home/manzp/caffe_ssd/caffe/python')
import caffe

# run decice setting
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
caffe.set_mode_gpu()  # caffe.set_mode_cpu()

# initialize model and use test mode(eg: don't perform dropout)
net_plate = caffe.Net('model/res18_deploy.prototxt', 'model/res18_iter_20600.caffemodel', caffe.TEST)


def preprocess(img_bgr_f, crop_size):
    """resize the image(min_side) proportionally
    Args:
        img_bgr_f: 0~255, BGR, uint8
        crop_size: min side crop size, int

    Returns:
        min side resize and center cropped img, float32
    """
    #  resize the image(min_side) proportionally
    h, w = img_bgr_f.shape[:2]
    if h > w:
        resize_ratio = float(crop_size) / w
        h_resized = int(np.round(resize_ratio * h))
        w_resized = crop_size
    else:
        resize_ratio = float(crop_size) / h
        h_resized = crop_size
        w_resized = int(np.round(resize_ratio * w))

    img_bgr_f_resized = cv2.resize(img_bgr_f, (w_resized, h_resized), interpolation=cv2.INTER_AREA)  # better for resize

    # center crop
    if h_resized < w_resized:
        off = (w_resized - h_resized) // 2
        img_bgr_f_cropped = img_bgr_f_resized[:, off:off + h_resized]
    else:
        off = (h_resized - w_resized) // 2
        img_bgr_f_cropped = img_bgr_f_resized[off:off + w_resized, :]

    return img_bgr_f_cropped.astype(np.float32)


def plate_cls(pre_img, net=net_plate):
    """
    Args:
        pre_img: min side resize and center cropped img, float32
        net: pre_load model

    Returns:
        max prob index of the first image in the batch
    """
    pre_img = pre_img.transpose((2, 0, 1))  # H,W,C-->C,H,W
    pre_img = (pre_img - 127.5) / 128.0

    net.blobs['data'].data[...] = pre_img
    out = net.forward()
    prob = out['prob'][0]

    return prob.argmax()  # prob.argsort()[::-1][:5], sort top five predictions from softmax output


if __name__ == '__main__':
    img_dir = 'test_imgs'
    save_dir = 'error_imgs'
    cnt = 0
    right_cnt = 0

    for img_name in os.listdir(img_dir):
        img_path = os.path.join(img_dir, img_name)
        img_bgr = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), 1)

        # preprocess and predict max prob index
        img_bgr_resized = preprocess(img_bgr, 112)
        predict_label = plate_cls(img_bgr_resized)

        # cal acc
        gt_label = int(img_name.split('_')[0])
        if predict_label == gt_label:
            right_cnt += 1
        else:
            print('img path is {}, predict_label is {}, gt_label is {}'.format(img_path, predict_label, gt_label))
            cv2.imencode('.jpg', img_bgr)[1].tofile(os.path.join(save_dir, img_name))

        cnt += 1

    print('accuracy is: {} {}/{}'.format(float(right_cnt) / cnt, right_cnt, cnt))


三、查看 blobs 形状、params 形状及参数

# -*- coding:utf-8 -*-

import os
import sys
reload(sys)
sys.setdefaultencoding('utf-8')

# caffe setting
sys.path.insert(0, '/home/manzp/caffe_ssd/caffe/python')
import caffe

# run decice setting
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
caffe.set_mode_gpu()  # caffe.set_mode_cpu()

# initialize model and use test mode(eg: don't perform dropout)
net_plate = caffe.Net('model/res18_deploy.prototxt', 'model/res18_iter_20600.caffemodel', caffe.TEST)

if __name__ == '__main__':
    # for each bottom and top blob in proto, get output shape
    for layer_name, blob in net_plate.blobs.iteritems():
        print layer_name + '\t' + str(blob.data.shape)

    # for each layer name in proto, get parameter data and shape(0 for weights, 1 for bias)
    for layer_name, param in net_plate.params.iteritems():
        if len(param) is 1:
            print layer_name + '\t' + str(param[0].data.shape)
        else:
            print layer_name + '\t' + str(param[0].data.shape), str(param[1].data.shape)

    # find specific layer name param(eg:BN 层, 均值、方差、缩放和平移参数, channel 维)
    print net_plate.params['bn_conv1'][0].data, net_plate.params['bn_conv1'][1].data
    print net_plate.params['scale_conv1'][0].data, net_plate.params['scale_conv1'][1].data

四、参考资料

1、https://github.com/BVLC/caffe/tree/master/python
2、PyCaffe 学习

你可能感兴趣的:(Caffe)