Bài viết này hướng dẫn cách thực hiện chuyển giao kiến thức (transfer learning) cho các mô hình phát hiện đối tượng bằng TensorFlow. Quá trình bao gồm chuẩn bị dữ liệu, xử lý dữ liệu, cấu hình mô hình và bắt đầu huấn luyện.
Chuẩn bị Dữ liệu
Để huấn luyện mô hình phát hiện đối tượng, bạn cần tạo một bộ dữ liệu bằng cách gán nhãn cho các hình ảnh. Công cụ LabelImg (https://github.com/heartexlabs/labelImg) có thể được sử dụng cho mục đích này. Sau khi gán nhãn, mỗi hình ảnh sẽ có một tệp XML tương ứng chứa thông tin về các đối tượng. Bộ dữ liệu nên được chia thành hai phần: train (huấn luyện) và test (kiểm thử). Tạo một thư mục có tên images trong thư mục models/research/object_detection của dự án và đặt các thư mục con train và test cùng với các hình ảnh và tệp XML của chúng vào đó.
Xử lý Dữ liệu
Hai bước chính trong xử lý dữ liệu là chuyển đổi tệp XML thành tệp CSV và sau đó chuyển đổi tệp CSV thành định dạng TFRecord mà TensorFlow có thể sử dụng để huấn luyện.
1. Chuyển đổi XML sang CSV
Tạo một tệp Python có tên xml_to_csv.py trong thư mục models/research/object_detection và thêm đoạn mã sau:
import os
import glob
import pandas as pd
import xml.etree.ElementTree as ET
def xml_to_csv(path):
xml_list = []
for xml_file in glob.glob(path + '/*.xml'):
tree = ET.parse(xml_file)
root = tree.getroot()
for member in root.findall('object'):
value = (root.find('filename').text,
int(root.find('size')[0].text),
int(root.find('size')[1].text),
member[0].text,
int(member[4][0].text),
int(member[4][1].text),
int(member[4][2].text),
int(member[4][3].text)
)
xml_list.append(value)
column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax']
xml_df = pd.DataFrame(xml_list, columns=column_name)
return xml_df
def main():
for folder in ['train','test']:
image_path = os.path.join(os.getcwd(), ('images/' + folder))
xml_df = xml_to_csv(image_path)
xml_df.to_csv(('images/' + folder + '_labels.csv'), index=None)
print('Successfully converted xml to csv.')
if __name__ == '__main__':
main()
Chạy tập lệnh này từ dòng lệnh:
python xml_to_csv.py
Thao tác này sẽ tạo ra các tệp train_labels.csv và test_labels.csv trong thư mục images.
2. Chuyển đổi CSV sang TFRecord
Tạo một tệp Python khác có tên generate_tfrecord.py trong thư mục models/research/object_detection. Lưu ý rằng bạn cần sửa đổi hàm class_text_to_int để phản ánh các lớp đối tượng cụ thể của bạn.
"""
Usage:
# Create train data:
python generate_tfrecord.py --csv_input=images/train_labels.csv --image_dir=images/train --output_path=train.record
# Create test data:
python generate_tfrecord.py --csv_input=images/test_labels.csv --image_dir=images/test --output_path=test.record
"""
from __future__ import division
from __future__ import print_function
from __future__ import absolute_import
import os
import io
import pandas as pd
# TensorFlow version check for compatibility
try:
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
except ImportError:
import tensorflow as tf
from PIL import Image
from object_detection.utils import dataset_util
from collections import namedtuple, OrderedDict
flags = tf.app.flags
flags.DEFINE_string('csv_input', '', 'Path to the CSV input')
flags.DEFINE_string('image_dir', '', 'Path to the image directory')
flags.DEFINE_string('output_path', '', 'Path to output TFRecord')
FLAGS = flags.FLAGS
# TODO: Replace this with your label map
def class_text_to_int(row_label):
if row_label == 'nine':
return 1
elif row_label == 'ten':
return 2
elif row_label == 'jack':
return 3
elif row_label == 'queen':
return 4
elif row_label == 'king':
return 5
elif row_label == 'ace':
return 6
else:
return None # Return None if the class is not recognized
def split_into_namedtuple(df, group):
data = namedtuple('data', ['filename', 'object'])
gb = df.groupby(group)
return [data(filename, gb.get_group(x)) for filename, x in zip(gb.groups.keys(), gb.groups)]
def create_tf_example(group, path):
with tf.gfile.GFile(os.path.join(path, '{}'.format(group.filename)), 'rb') as fid:
encoded_jpg = fid.read()
encoded_jpg_io = io.BytesIO(encoded_jpg)
image = Image.open(encoded_jpg_io)
width, height = image.size
filename = group.filename.encode('utf8')
image_format = b'jpg'
xmins, xmaxs, ymins, ymaxs = [], [], [], []
classes_text, classes = [], []
for index, row in group.object.iterrows():
xmins.append(row['xmin'] / width)
xmaxs.append(row['xmax'] / width)
ymins.append(row['ymin'] / height)
ymaxs.append(row['ymax'] / height)
classes_text.append(row['class'].encode('utf8'))
classes.append(class_text_to_int(row['class']))
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_jpg),
'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
def main(_):
writer = tf.python_io.TFRecordWriter(FLAGS.output_path)
image_directory = os.path.join(os.getcwd(), FLAGS.image_dir)
examples = pd.read_csv(FLAGS.csv_input)
grouped_examples = split_into_namedtuple(examples, 'filename')
for group in grouped_examples:
tf_record_example = create_tf_example(group, image_directory)
writer.write(tf_record_example.SerializeToString())
writer.close()
output_path_full = os.path.join(os.getcwd(), FLAGS.output_path)
print('Successfully created the TFRecords: {}'.format(output_path_full))
if __name__ == '__main__':
tf.app.run()
Thực thi tập lệnh để tạo tệp TFRecord:
python generate_tfrecord.py --csv_input=images/train_labels.csv --image_dir=images/train --output_path=train.record
python generate_tfrecord.py --csv_input=images/test_labels.csv --image_dir=images/test --output_path=test.record
Các tệp train.record và test.record sẽ được tạo trong thư mục models/research/object_detection.
Cấu hình Mô hình
Trước khi tiến hành cấu hình, hãy đảm bảo bạn đã tải xuống một mô hình phát hiện đối tượng được huấn luyện trước (ví dụ: faster_rcnn_inception_v2_coco_2018_01_28) và đặt nó vào thư mục models/research/object_detection.
1. Tệp labelmap.pbtxt
Tạo một tệp có phần mở rộng .pbtxt, ví dụ: labelmap.pbtxt, trong thư mục models/research/object_detection/training (tạo thư mục này nếu nó chưa tồn tại). Tệp này ánh xạ tên lớp đối tượng với các số nguyên.
item {
id: 1
name: 'your_class_1'
}
item {
id: 2
name: 'your_class_2'
}
# ... và cứ thế tiếp tục cho tất cả các lớp của bạn
2. Tệp Cấu hình Pipeline (Ví dụ: faster_rcnn_inception_v2_pets.config)
Sao chép tệp faster_rcnn_inception_v2_pets.config từ research/object_detection/samples/configs vào thư mục research/object_detection/training. Sau đó, chỉnh sửa tệp này:
- Dòng 9: Cập nhật
num_classesthành số lượng lớp đối tượng của bạn. - Dòng ~110: Thay đổi
fine_tune_checkpointđể trỏ đến tệpmodel.ckptcủa mô hình được huấn luyện trước. - Khối
train_input_reader(khoảng dòng 126 và 128): Cập nhậtinput_pathvàlabel_map_pathđể trỏ đến các tệptrain.recordvàlabelmap.pbtxtcủa bạn. - Dòng ~132: Đặt
num_examplesthành tổng số hình ảnh trong thư mục kiểm thử của bạn. - Khối
eval_input_reader(khoảng dòng 140 và 142): Cập nhậtinput_pathvàlabel_map_pathđể trỏ đến các tệptest.recordvàlabelmap.pbtxtcủa bạn.
Lưu ý: Các số dòng có thể thay đổi một chút tùy thuộc vào phiên bản. Hãy tìm các phần cấu hình tương ứng.
Cài đặt Môi trường
Đảm bảo bạn đã cài đặt tất cả các gói phụ thuộc cần thiết. Đặc biệt, bạn cần cài đặt thư viện slim và object_detection.
1. Cài đặt slim
Điều hướng đến thư mục models/research/slim và chạy:
python setup.py install
Nếu gặp lỗi "error: could not create 'build': File exists", hãy xóa thư mục build trong slim rồi chạy lại lệnh.
2. Cài đặt object-detection
Điều hướng đến thư mục models/research và chạy:
python setup.py install
Bắt đầu Huấn luyện
1. Tạo Tệp train.py
Sao chép tệp train.py từ thư mục object_detection/legacy vào thư mục object_detection. Sau đó, chạy lệnh sau từ thư mục object_detection:
python train.py --logtostderr --train_dir=training/ --pipeline_config_path=training/faster_rcnn_inception_v2_pets.config
2. Khởi chạy TensorBoard
Để theo dõi quá trình huấn luyện, hãy chạy TensorBoard từ thư mục object_detection:
tensorboard --logdir=training
Mở trình duyệt và truy cập địa chỉ được cung cấp (thường là http://localhost:6006).
Xuất Mô hình
Sau khi quá trình huấn luyện hoàn tất (bạn có thể dừng bằng Ctrl+C khi độ mất mát (loss) đã giảm xuống mức chấp nhận được), hãy xuất đồ thị suy luận (inference graph):
python export_inference_graph.py --input_type image_tensor --pipeline_config_path training/faster_rcnn_inception_v2_pets.config --trained_checkpoint_prefix training/model.ckpt-XXXX --output_directory inference_graph
Thay thế XXXX bằng số lần lặp huấn luyện cao nhất của tệp model.ckpt. Tệp frozen_inference_graph.pb sẽ được tạo trong thư mục inference_graph.
Kiểm thử Mô hình
Sử dụng các tập lệnh được cung cấp để kiểm thử mô hình đã huấn luyện trên hình ảnh hoặc luồng video. Bạn có thể cần cài đặt thư viện OpenCV (opencv-python) nếu chưa có.
# Đối với kiểm thử hình ảnh
python Object_detection_image.py
# Đối với kiểm thử video
python Object_detection_video.py
Lưu ý rằng việc cài đặt OpenCV đôi khi có thể gặp sự cố; hãy thử các phương pháp khác nhau như cài đặt qua conda hoặc tải xuống gói bánh xe phù hợp nếu cần.