mxnet-batch-predict.py

mxnet-batch-predict.py

# -*- coding: utf-8 -*-
import os, sys, codecs
import glob

import numpy as np

import cv2
import imagehash
from PIL import Image

Oxford_IMAGE_Path = '/export/home/liuyuzhong/dataset/Oxford-5k/oxbuild_images/'
Oxford_GT_Path    = '/export/home/liuyuzhong/dataset/Oxford-5k/gt_files_170407/'

MXNet_Model_Path = '/export/home/liuyuzhong/work/mxnet_models/'

import mxnet as mx
ctx = mx.gpu(0)
sym, arg_params, aux_params = mx.model.load_checkpoint(MXNet_Model_Path + '/resnet-18', 0)

from collections import namedtuple
Batch = namedtuple('Batch', ['data'])

all_layers = sym.get_internals()
sym_inner = all_layers['stage4_unit2_conv2_output']

mod = mx.mod.Module(symbol=sym_inner, label_names=None, context=ctx)
mod.bind(for_training=False, data_shapes=[('data', (1,3,224,224))])
mod.set_params(arg_params, aux_params)

def predict_path(paths):
    img_input = np.zeros((len(paths), 3, 1024, 768))
    for idx, path in enumerate(paths):
        img = cv2.imread(path)
        img = cv2.resize(img, (768, 1024), interpolation=cv2.INTER_CUBIC)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = np.swapaxes(img, 0, 2)
        img = np.swapaxes(img, 1, 2)
        img_input[idx] = img
        
    mod.forward(Batch([mx.nd.array(img_input)]))
    feat = mod.get_outputs()[0]
    
    # max-pooling
    feat = feat.max(2).max(-1)
    return feat.asnumpy()

def chunks(l, n):
    """Yield successive n-sized chunks from l."""
    for i in range(0, len(l), n):
        yield l[i:i + n]
        
img_paths = glob.glob(Oxford_IMAGE_Path + '*')
imgs_feat = []
for path_batch in chunks(img_paths[:], 10):
    # print('aaa')
    imgs_feat.append(predict_path(path_batch))
    
imgs_feat = np.concatenate(imgs_feat)