TensorFlow 2 Object Detection
Создание точных моделей машинного обучения, способных локализовать и идентифицировать несколько объектов на одном изображении, остается основной проблемой компьютерного зрения. API обнаружения объектов TensorFlow - это платформа с открытым исходным кодом, построенная на основе TensorFlow, которая упрощает создание, обучение и развертывание моделей обнаружения объектов. Детектирование объектов в TensorFlow 2 является мощнейшим средством, предназначенным для распознавания объектов на изображениях. Ссылка: https://github.com/tensorflow/models/tree/master/research/object_detection
Установка имеет свои особенности, в данной статье установка производится с учетом среды Kaggle
Установка Object Detection
Саму процедуру установки можно найти по следующим ссылкам:
- https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2.md
- https://www.kaggle.com/bhallaakshit/dicom-training-prediction-and-evaluation
Выглядит это примерно так:
!# 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 ..
Подготовка данных
Про подготовку данных для обучения написано тут:
Данные лучше засовывать в 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()