TensorFlow 2 Object Detection

Создание точных моделей машинного обучения, способных локализовать и идентифицировать несколько объектов на одном изображении, остается основной проблемой компьютерного зрения. API обнаружения объектов TensorFlow - это платформа с открытым исходным кодом, построенная на основе TensorFlow, которая упрощает создание, обучение и развертывание моделей обнаружения объектов. Детектирование объектов в TensorFlow 2 является мощнейшим средством, предназначенным для распознавания объектов на изображениях. Ссылка: https://github.com/tensorflow/models/tree/master/research/object_detection

Установка имеет свои особенности, в данной статье установка производится с учетом среды Kaggle

Установка Object Detection

Саму процедуру установки можно найти по следующим ссылкам:

Выглядит это примерно так:

!# Download models
!git clone --depth 1 https://github.com/tensorflow/models

!# Compile proto files 
! # sudo apt install -y protobuf-compiler # Already present
%cd models/research
!protoc object_detection/protos/*.proto --python_out=.
%cd ..
%cd ..

!# Install cocoapi
!pip install cython 
!git clone https://github.com/cocodataset/cocoapi.git
%cd cocoapi/PythonAPI
!make
%cd ..
%cd ..
!cp -r cocoapi/PythonAPI/pycocotools models/research/

!# Install object detection api
%cd models/research
!cp object_detection/packages/tf2/setup.py .
!python -m pip install .
%cd ..
%cd ..
Подготовка данных

Про подготовку данных для обучения написано тут:

https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/using_your_own_dataset.md

Данные лучше засовывать в tfrecord, чтобы можно было обучать и на TPU. По указанной ссылке показано, как это сделать. Для каждого изображения вызывается функция, которая создает порцию данных (Example). В данном случае она простая:

def create_cat_tf_example(encoded_cat_image_data):

   """Creates a tf.Example proto from sample cat image.

 

  Args:

    encoded_cat_image_data: The jpg encoded data of the cat image.

 

  Returns:

    example: The created tf.Example.

  """

 

  height = 1032.0

  width = 1200.0

  filename = 'example_cat.jpg'

  image_format = b'jpg'

 

  xmins = [322.0 / 1200.0]

  xmaxs = [1062.0 / 1200.0]

  ymins = [174.0 / 1032.0]

  ymaxs = [761.0 / 1032.0]

  classes_text = ['Cat']

  classes = [1]

 

  tf_example = tf.train.Example(features=tf.train.Features(feature={

      'image/height': dataset_util.int64_feature(height),

      'image/width': dataset_util.int64_feature(width),

      'image/filename': dataset_util.bytes_feature(filename),

      'image/source_id': dataset_util.bytes_feature(filename),

      'image/encoded': dataset_util.bytes_feature(encoded_image_data),

      'image/format': dataset_util.bytes_feature(image_format),

      'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),

      'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),

      'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),

      'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),

      'image/object/class/text': dataset_util.bytes_list_feature(classes_text),

      'image/object/class/label': dataset_util.int64_list_feature(classes),

  }))

  return tf_example

 

 Обязательные параметры при составлении порции - это списки 'image/object/bbox/' xmins,xmaxs,ymins,ymaxs (описывают область объекта на изображении), 'image/object/class/text'- список наименование классов, 'image/object/class/label' - индексы классов (начиная с 1, т.к. 0 зарезервирован, как отсутствие объектов), 'image/encoded' - само изображение, ну и доп. параметры, которые вы видите выше.

Чтобы записать tfrecord файл, нужно:

  • открыть запись в файл writer = tf.python_io.TFRecordWriter(путь куда писать);
  • в цикле перебора изображений вызывать функцию записи (в данном случае create_cat_tf_example), которая возвращает Example, и записать результат с помощью writer.write(tf_example.SerializeToString());
  • в конце завершить работу с файлом tfrecord так writer.close().
 Конфигурирование модели

Сначала модель загружается и распаковывается.

url = "http://download.tensorflow.org/models/object_detection/tf2/20200711/efficientdet_d1_coco17_tpu-32.tar.gz"

path = "./workspace/pretrained_models"

r = requests.get(url)

 

# Extract model

thetarfile = tarfile.open(

    fileobj = BytesIO(r.content),

    mode = "r|gz"

)

 

# Save model

 

thetarfile.extractall(path = path)

Дальше нужно изменить конфигурацию под вашу задачу, о чем можно почитать тут:

https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/configuring_jobs.md

Обычно изменяется pipeline.config, который перемещается в нужную папку:

fname = "pipeline.config"

model_name = "efficientdet_d1_coco17_tpu-32"

 

src = os.path.join(path, model_name, fname)

dst = src.replace("pretrained_", "").replace(fname, "")

 

os.makedirs(dst, exist_ok = True)

 

copy2(src, dst)

 

Потом:

!wget "http://download.tensorflow.org/models/object_detection/tf2/20200711/efficientdet_d1_coco17_tpu-32.tar.gz"

!tar -xf efficientdet_d1_coco17_tpu-32.tar.gz

 

!mv efficientdet_d1_coco17_tpu-32/checkpoint models/research/object_detection/test_data/

Но если вы работаете через jupyter notebook, то это можно сделать из него, сначала загрузив файл.

fpath = os.path.join(dst, fname)

 

config_dic = config_util.get_configs_from_pipeline_file(fpath)

При этом настройки для различных сетей будут отличаться. Ниже представлены 2 варианта.

# Retina net

 

fpath = os.path.join(dst, fname)

config_dic = config_util.get_configs_from_pipeline_file(fpath)

 

config_dic["model"].ssd.num_classes = len(LabelMap)

 

 

config_dic["train_config"].batch_size = 4

config_dic["train_config"].fine_tune_checkpoint = os.path.join(path, model_name, "checkpoint/ckpt-0")

config_dic["train_config"].fine_tune_checkpoint_type = "detection"

config_dic["train_config"].use_bfloat16 = False # Set to True if training on a TPU

config_dic["train_config"].num_steps = 16_000

config_dic["train_config"].data_augmentation_options[1].random_crop_image.max_aspect_ratio=2

 

 

config_dic["train_input_config"].label_map_path = path_label

config_dic["train_input_config"].tf_record_input_reader.input_path[:] = train_data

 

config_dic["eval_input_configs"][0].label_map_path = path_label

 

config_dic["eval_input_configs"][0].tf_record_input_reader.input_path[:] = valid_data

#Effnet

fpath = os.path.join(dst, fname)

config_dic = config_util.get_configs_from_pipeline_file(fpath)

 

config_dic["model"].ssd.num_classes = len(LabelMap)

config_dic["model"].ssd.image_resizer.keep_aspect_ratio_resizer.min_dimension = 300

 

 

config_dic["train_config"].batch_size = 4

config_dic["train_config"].fine_tune_checkpoint = os.path.join(path, model_name, "checkpoint/ckpt-0")

config_dic["train_config"].fine_tune_checkpoint_type = "detection"

config_dic["train_config"].use_bfloat16 = True # Set to True if training on a TPU

config_dic["train_config"].num_steps = 6_000

config_dic["train_config"].data_augmentation_options[1].random_scale_crop_and_pad_to_square.scale_min = 0.7

config_dic["train_config"].data_augmentation_options[1].random_scale_crop_and_pad_to_square.scale_max = 1.3

 

 

config_dic["train_input_config"].label_map_path = path_label

config_dic["train_input_config"].tf_record_input_reader.input_path[:] = train_data

 

config_dic["eval_input_configs"][0].label_map_path = path_label

 

config_dic["eval_input_configs"][0].tf_record_input_reader.input_path[:] = valid_data

 

Здесь path_label - это путь до файла с метками pbtxt. train_data и valid_data - это списки с полными путями до записей tfrecord

После изменений записывайте изменения в файл:

config = config_util.create_pipeline_proto_from_configs(config_dic)

 

config_util.save_pipeline_config(config, dst)

 Обучение

Процесс обучения простой. Вы вызываете команду:

!python workspace/model_main_tf2.py --model_dir=$dst --pipeline_config_path=$fpath 

И следите за изменениями потерь. Можно использовать tensorboard, чтобы смотреть графики online.

Экспортируйте модель:

!python workspace/exporter_main_v2.py --input_type=image_tensor --pipeline_config_path=$fpath --trained_checkpoint_dir=$dst --output_directory=workspace/exported_models/$model_name

tar_model_name = model_name + ".tar.gz"!tar -zcvf workspace/exported_models/$tar_model_name workspace/exported_models/$model_name

 Распознавание

Загрузите 1 картинку:

for raw_record in dataset.take(1): # Select one shard from the TFRecords dataset

        example = tf.train.Example()

        example.ParseFromString(raw_record.numpy())

# декодирование

img_encoded = example.features.feature['image/encoded'].bytes_list.value[0]

img = tf.io.decode_jpeg(img_encoded)

img = cv2.cvtColor(img.numpy(), cv2.COLOR_GRAY2RGB)

img = img[tf.newaxis, ...]

#загрузка модели

detector = tf.saved_model.load(os.path.join("workspace/exported_models", model_name, "saved_model"))

#Распознавание

result = detector(img)

#Преобразование результата

 

result = {k:v.numpy() for k, v in result.items()}

Ну и рисование и вывод средствами TensorFlow

viz_util.visualize_boxes_and_labels_on_image_array(

    image = img[0],

    boxes = result['detection_boxes'][0],

    classes = (result['detection_classes'][0]).astype(int),

    scores = result['detection_scores'][0],

    category_index = LabelMap,

    use_normalized_coordinates = True,

    min_score_thresh = 0.3,

    line_thickness = 3,

    max_boxes_to_draw = 100,

)

 

%matplotlib inline

 

plt.figure(figsize = (8, 8))

plt.imshow(img[0])

plt.title("Prediction")

 

plt.show()