image_recognition_client_test

#!/usr/bin/env python
import os
import argparse
import random
import cv2
from cv_bridge import CvBridge, CvBridgeError
from mas_perception_libs.utils import case_insensitive_glob
from mas_perception_libs.image_recognition_service import RecognizeImageServiceProxy


ALLOWED_FILE_TYPES = ('*.jpg', '*.jpeg', '*.png', '*.bmp')


def main(arguments):
    service_proxy = RecognizeImageServiceProxy(arguments.service_name, arguments.model_name,
                                                 arguments.preprocess_input_module)

    file_list = []
    for file_type in ALLOWED_FILE_TYPES:
        file_list.extend(case_insensitive_glob(os.path.join(arguments.test_dir, file_type)))

    test_file_list = []
    if arguments.num_samples > 0:
        num_samples = arguments.num_samples
        indices = random.sample(range(len(file_list)), num_samples)
        for index in indices:
            test_file_list.append(file_list[index])
    else:
        test_file_list = file_list

    cv_bridge = CvBridge()
    image_messages = []
    for filename in test_file_list:
        cv_image = cv2.imread(filename)
        try:
            image_messages.append(cv_bridge.cv2_to_imgmsg(cv_image, "bgr8"))
        except CvBridgeError as e:
            print('error converting "{0}" to ROS format: {1}'.format(filename, e.message))

    results = service_proxy.classify_image_messages(image_messages)
    print(results)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Tool to test model with test images using KerasImageClassifier class.'
                                                 'Assuming images to be of type jpg')
    parser.add_argument('--test-dir', '-t', required=True, help='directory with test images')
    parser.add_argument('--service-name', '-s', required=True, help='name of recognition service')
    parser.add_argument('--num-samples', '-n', type=int, default=-1, help='number of samples to test, if left blank,'
                                                                          ' take all samples.')
    parser.add_argument('--preprocess-input-module', '-p', default=None, help='module containing image preprocessing'
                                                                              ' function.')
    parser.add_argument('model_name', help='Keras model to be tested')
    args = parser.parse_args()
    main(args)