imageNet_Features.py 文件源码

def compute_features(imgname, args):
    """
    Instantiate a classifier class, pass the images through the network and save features.
    Features are saved in .mat format
    """
    image_dims = [int(s) for s in args.images_dim.split(',')]
    if args.force_grayscale:
      channel_swap = None
      mean_file = None
    else:
      channel_swap = [int(s) for s in args.channel_swap.split(',')]
      mean_file = args.mean_file

    # Make classifier.
    classifier = caffe.Classifier(args.model_def, args.pretrained_model,
            image_dims=image_dims, gpu=args.gpu, mean_file=mean_file,
            input_scale=args.input_scale, channel_swap=channel_swap)

    if args.gpu:
        print 'GPU mode'



    outfname = imgname.replace('imageNetForWeb', 'imageNetForWeb_Features') + ".mat"
    print outfname
    if not path.exists(path.dirname(outfname)):
        os.makedirs(path.dirname(outfname))

    inputs = [caffe.io.load_image(imgname)]

    if args.force_grayscale:
        inputs = [rgb2gray(input) for input in inputs];

    print "Classifying %d inputs." % len(inputs)

    scores = classifier.predict(inputs, not args.center_only)
        # Now save features
    feature_dict = {}
    feature_dict['IMG_NAME'] = path.join(path.dirname(imgname), path.basename(imgname))
    feature_dict['fc7'] = sp.asarray(classifier.blobs['fc7'].data.squeeze(axis=(2,3)))
    feature_dict['fc8'] = sp.asarray(classifier.blobs['fc8'].data.squeeze(axis=(2,3)))
    feature_dict['prob'] = sp.asarray(classifier.blobs['prob'].data.squeeze(axis=(2,3)))
    feature_dict['scores'] = sp.asarray(scores)
    savemat(outfname, feature_dict)