Заключительная часть цикла вводных заметок в машинное зрение посвящена использованию предобученных сетей в задаче детектирования. В частности использование сетей, обученных на датасете Coco на базе фреймворка TensorFlow Object Detection API.
В отличие от предыдущих статей, которые могли быть выполнены в стандартном окружении Python, эту часть рекомендуется выполнять в Jupyter Notebook, запущеном в контейнере (либо пользоваться облачными сервисами, например colab.research.google.com в силу того, что в процессе будут устанавливаться компоненты.
Установка TensorFlow Object Detection API
Первым делом необходимо установить необходимые компоненты TensorFlow Object Detection API. Для его установки необходимо клонировать репозиторий, скомпилировать модели и запустить тесты.
1 2 3 | pip install -U --pre tensorflow== "2.*" pip install tf_slim pip install pycocotools |
Получите tenorflow/models или cd в родительский каталог репозитория.
1 2 3 4 5 6 7 8 9 | import os import pathlib if "models" in pathlib.Path.cwd().parts: while "models" in pathlib.Path.cwd().parts: os.chdir( '..' ) elif not pathlib.Path( 'models' ).exists(): !git clone - - depth 1 https: / / github.com / tensorflow / models |
Скомпилируйте protobufs и установите пакет object_detection
1 2 3 | %% bash cd models /research/ protoc object_detection /protos/ *.proto --python_out=. |
1 2 3 | %% bash cd models /research/ protoc object_detection /protos/ *.proto --python_out=. |
Импортируем зависимости и добавим в системные пути расположение TensorFlow Object Detection API
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 | import sys, os sys.path.append( 'models/research' ) sys.path.append( 'models/research/object_detection' ) import numpy as np import six.moves.urllib as urllib import tarfile import zipfile import cv2 from PIL import Image from matplotlib import pyplot as plt import tensorflow as tf from object_detection.utils import ops as utils_ops from utils import label_map_util from utils import visualization_utils as vis_util |
Загрузка предобученной модели
Зададим URL, где хранятся предобученные модели и имя планируемой к использованию модели.
1 2 | model_name = 'mask_rcnn_inception_v2_coco_2018_01_28' |
Создадим переменные среды для работы с путями из shell
1 2 | os.environ[ 'MODEL_PATH' ] = model_path + model_name + '.tar.gz' os.environ[ 'MODEL_FILE_NAME' ] = model_name + '.tar.gz' |
Удалим загруженную ранее модель, если она есть, загрузим модель и распакуем ее.
1 2 3 | rm $MODEL_FILE_NAME wget $MODEL_PATH tar xfz $MODEL_FILE_NAME |
Подготовка модели к запуску
Загрузка модели в память:
1 2 3 4 5 6 7 8 | model_file_name = model_name + '/frozen_inference_graph.pb' detection_graph = tf.Graph() with detection_graph.as_default(): od_graph_def = tf.compat.v1.GraphDef() with tf.compat.v2.io.gfile.GFile(model_file_name, 'rb' ) as fid: serialized_graph = fid.read() od_graph_def.ParseFromString(serialized_graph) tf.import_graph_def(od_graph_def, name = '') |
Загрузка меток классов
1 2 3 | label_map = label_map_util.load_labelmap( 'models/research/object_detection/data/mscoco_label_map.pbtxt' ) categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes = 90 , use_display_name = True ) category_index = label_map_util.create_category_index(categories) |
Воспользуемся следующей функцией для обработки единичного изображения
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | def run_inference_for_single_image(image, graph): with graph.as_default(): with tf.compat.v1.Session() as sess: ops = tf.compat.v1.get_default_graph().get_operations() all_tensor_names = {output.name for op in ops for output in op.outputs} tensor_dict = {} for key in [ 'num_detections' , 'detection_boxes' , 'detection_scores' , 'detection_classes' , 'detection_masks' ]: tensor_name = key + ':0' if tensor_name in all_tensor_names: tensor_dict[key] = tf.compat.v1.get_default_graph().get_tensor_by_name( tensor_name) if 'detection_masks' in tensor_dict: detection_boxes = tf.squeeze(tensor_dict[ 'detection_boxes' ], [ 0 ]) detection_masks = tf.squeeze(tensor_dict[ 'detection_masks' ], [ 0 ]) real_num_detection = tf.cast(tensor_dict[ 'num_detections' ][ 0 ], tf.int32) detection_boxes = tf. slice (detection_boxes, [ 0 , 0 ], [real_num_detection, - 1 ]) detection_masks = tf. slice (detection_masks, [ 0 , 0 , 0 ], [real_num_detection, - 1 , - 1 ]) detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks( detection_masks, detection_boxes, image.shape[ 0 ], image.shape[ 1 ]) detection_masks_reframed = tf.cast( tf.greater(detection_masks_reframed, 0.5 ), tf.uint8) tensor_dict[ 'detection_masks' ] = tf.expand_dims( detection_masks_reframed, 0 ) image_tensor = tf.compat.v1.get_default_graph().get_tensor_by_name( 'image_tensor:0' ) # Запуск поиска объектов output_dict = sess.run(tensor_dict, feed_dict = {image_tensor: np.expand_dims(image, 0 )}) # Преобразование выходных данных из массивов float32 в нужный формат output_dict[ 'num_detections' ] = int (output_dict[ 'num_detections' ][ 0 ]) output_dict[ 'detection_classes' ] = output_dict[ 'detection_classes' ][ 0 ].astype(np.uint8) output_dict[ 'detection_boxes' ] = output_dict[ 'detection_boxes' ][ 0 ] output_dict[ 'detection_scores' ] = output_dict[ 'detection_scores' ][ 0 ] if 'detection_masks' in output_dict: output_dict[ 'detection_masks' ] = output_dict[ 'detection_masks' ][ 0 ] return output_dict |
Загрузим тестовое изображение
1 | wget https: //digiratory .ru /wp-content/uploads/172914606_546132cbe9_z1 .jpg |
Выведем загруженное изображение
1 2 3 4 | % matplotlib inline sample_image = cv2.cvtColor(cv2.imread( '172914606_546132cbe9_z1.jpg' ), cv2.COLOR_BGR2RGB) plt.imshow(sample_image) plt.show() |

Получим отклик нейронной сети на изображение
1 | output_dict = run_inference_for_single_image(sample_image, detection_graph) |
Визуализируем результат детектирования и сегментации.
1 2 3 4 5 6 7 8 9 10 11 12 13 | vis_util.visualize_boxes_and_labels_on_image_array( sample_image, output_dict[ 'detection_boxes' ], output_dict[ 'detection_classes' ], output_dict[ 'detection_scores' ], category_index, instance_masks = output_dict.get( 'detection_masks' ), use_normalized_coordinates = True , line_thickness = 0 , min_score_thresh = . 6 ) plt.figure(figsize = ( 12 , 8 )) plt.grid( False ) plt.imshow(sample_image) |

Вы можете попробовать запустить различные предобученные модели, предложенные разработчиками и посмотреть на различие результатов.